diff --git a/databuilder/query_language.py b/databuilder/query_language.py index 9f2caf3ad..5e6f49488 100644 --- a/databuilder/query_language.py +++ b/databuilder/query_language.py @@ -632,7 +632,7 @@ def take(self, series): # pragma: no cover return self.where(series) def where(self, series): - return EventFrame( + return self.__class__( qm.Filter( source=self.qm_node, condition=_convert(series), @@ -646,7 +646,7 @@ def drop(self, series): # pragma: no cover return self.except_where(series) def except_where(self, series): - return EventFrame( + return self.__class__( qm.Filter( source=self.qm_node, condition=qm.Function.Or( @@ -669,7 +669,7 @@ def sort_by(self, *order_series): return SortedEventFrame(qm_node) -class SortedEventFrame(BaseFrame): +class SortedEventFrame(EventFrame): def first_for_patient(self): return PatientFrame( qm.PickOneRowPerPatient( diff --git a/docs/includes/generated_docs/specs.md b/docs/includes/generated_docs/specs.md index 7efd17a1c..1440af3c4 100644 --- a/docs/includes/generated_docs/specs.md +++ b/docs/includes/generated_docs/specs.md @@ -388,6 +388,59 @@ returns the following patient series: +### 2.4 Mixing the order of `sort_by` and `where` operations + + +#### 2.4.1 Sort by before where + +This example makes use of an event-level table named `e` containing the following data: + +| patient|i1|i2 | +| - | - | - | +| 1|101|1 | +| 1|102|2 | +| 1|103|2 | +| 2|203|1 | +| 2|202|2 | +| 2|201|2 | + +``` +e.sort_by(e.i1).where(e.i1 > 102).first_for_patient().i1 +``` +returns the following patient series: + +| patient | value | +| - | - | +| 1|103 | +| 2|201 | + + + +#### 2.4.2 Sort by interleaved with where + +This example makes use of an event-level table named `e` containing the following data: + +| patient|i1|i2 | +| - | - | - | +| 1|101|1 | +| 1|102|2 | +| 1|103|2 | +| 2|203|1 | +| 2|202|2 | +| 2|201|2 | + +``` +e.sort_by(e.i1).where(e.i2 > 1).sort_by(e.i2).first_for_patient().i1 +``` +returns the following patient series: + +| patient | value | +| - | - | +| 1|102 | +| 2|201 | + + + ## 3 Aggregating event and patient frames diff --git a/tests/spec/sort_and_pick/test_sort_by_interleaved_with_where.py b/tests/spec/sort_and_pick/test_sort_by_interleaved_with_where.py new file mode 100644 index 000000000..21bf41bfe --- /dev/null +++ b/tests/spec/sort_and_pick/test_sort_by_interleaved_with_where.py @@ -0,0 +1,38 @@ +from ..tables import e + +title = "Mixing the order of `sort_by` and `where` operations" + +table_data = { + e: """ + | i1 | i2 + --+-----+--- + 1 | 101 | 1 + 1 | 102 | 2 + 1 | 103 | 2 + 2 | 203 | 1 + 2 | 202 | 2 + 2 | 201 | 2 + """, +} + + +def test_sort_by_before_where(spec_test): + spec_test( + table_data, + e.sort_by(e.i1).where(e.i1 > 102).first_for_patient().i1, + { + 1: 103, + 2: 201, + }, + ) + + +def test_sort_by_interleaved_with_where(spec_test): + spec_test( + table_data, + e.sort_by(e.i1).where(e.i2 > 1).sort_by(e.i2).first_for_patient().i1, + { + 1: 102, + 2: 201, + }, + ) diff --git a/tests/spec/toc.py b/tests/spec/toc.py index d9fb5d2e1..df47da4e0 100644 --- a/tests/spec/toc.py +++ b/tests/spec/toc.py @@ -9,6 +9,7 @@ "test_sort_by_column_and_pick", "test_sort_by_multiple_columns_and_pick", "test_sort_by_column_with_nulls_and_pick", + "test_sort_by_interleaved_with_where", ], "aggregate_frame": [ "test_exists_for_patient",