feat: Support where() after sort_by() in ehrQL

opensafely-core · Mar 14, 2023 · 96d8e81 · 96d8e81
1 parent 00eceea
commit 96d8e81
Show file tree

Hide file tree

Showing 4 changed files with 95 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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",