Merge pull request #36 from OHDSI/fix/era-collapse-deterministic

fix(execution): make ERA collapse ordering deterministic

Author: egillax

circe/execution/engine/collapse.py

@@ -27,7 +27,11 @@ def _apply_censor_window(events, censor_window):
 def _collapse_era(intervals, era_pad: int):
     padded = intervals.mutate(_padded_end_date=(intervals.end_date + ibis.interval(days=int(era_pad))))
 
-    ordering = [padded.start_date]
+    ordering = [
+        padded.start_date,
+        padded._padded_end_date.desc(),
+        padded.end_date.desc(),
+    ]
     ordered_window = ibis.window(group_by=padded.person_id, order_by=ordering)
     cumulative_window = ibis.cumulative_window(group_by=padded.person_id, order_by=ordering)
     with_cummax = padded.mutate(_cummax_padded_end=padded._padded_end_date.max().over(cumulative_window))
@@ -44,7 +48,12 @@ def _collapse_era(intervals, era_pad: int):
 
     grouping_window = ibis.cumulative_window(
         group_by=marked.person_id,
-        order_by=[marked.start_date, marked._is_new_group.desc()],
+        order_by=[
+            marked.start_date,
+            marked._padded_end_date.desc(),
+            marked.end_date.desc(),
+            marked._is_new_group.desc(),
+        ],
     )
     group_index = marked._is_new_group.sum().over(grouping_window)
     grouped = marked.mutate(_group_idx=group_index)

tests/execution/test_end_strategy_censoring.py

@@ -283,6 +283,40 @@ def test_collapse_settings_era_merges_tied_start_dates_into_one_group():
     assert str(result.iloc[0]["end_date"])[:10] == "2020-01-05"
 
 
+def test_collapse_settings_era_merges_contained_intervals_after_tied_start_dates():
+    ibis = pytest.importorskip("ibis")
+    _ = pytest.importorskip("duckdb")
+
+    conn = ibis.duckdb.connect()
+    _seed_common_tables(conn, ibis)
+    conn.create_table(
+        "condition_occurrence",
+        obj=ibis.memtable(
+            {
+                "person_id": [1, 1, 1],
+                "condition_occurrence_id": [100, 101, 102],
+                "condition_concept_id": [111, 111, 111],
+                "condition_start_date": ["2020-01-01", "2020-01-01", "2020-01-10"],
+                "condition_end_date": ["2020-01-02", "2020-02-01", "2020-01-15"],
+                "visit_occurrence_id": [10, 10, 10],
+            }
+        ),
+        overwrite=True,
+    )
+
+    expression = CohortExpression(
+        concept_sets=[_make_concept_set(1, 111)],
+        primary_criteria=PrimaryCriteria(criteria_list=[ConditionOccurrence(codeset_id=1)]),
+        end_strategy=DateOffsetStrategy(offset=0, date_field="end_date"),
+        collapse_settings=CollapseSettings(era_pad=0),
+    )
+
+    result = build_cohort(expression, backend=conn, cdm_schema="main").execute()
+    assert len(result) == 1
+    assert str(result.iloc[0]["start_date"])[:10] == "2020-01-01"
+    assert str(result.iloc[0]["end_date"])[:10] == "2020-02-01"
+
+
 def test_apply_end_strategy_rejects_invalid_date_field_and_preserves_fallback_semantics():
     ibis_mod = pytest.importorskip("ibis")
     _ = pytest.importorskip("duckdb")