Fix stale characteristics under lazy execution when the index change

phenex/core/cohort.py

@@ -853,6 +853,10 @@ def execute(
 
         logger.info(f"Cohort '{self.name}': executing index stage ...")
 
+        index_membership_changed = lazy_execution and Node._node_manager.node_changed(
+            self.index_table_node, con
+        )
+
         self.index_stage.execute(
             tables=self.subset_tables_entry,
             con=con,
@@ -902,6 +906,34 @@ def execute(
                     self.subset_tables_index[node.name] = type(entry_tbl)(filtered_ibis)
 
         if self.reporting_stage:
+            # If the index population changed, clear characteristics/outcomes
+            if index_membership_changed:
+                logger.info(
+                    f"Cohort '{self.name}': index population changed; invalidating cached "
+                    f"characteristics/outcomes so they recompute against the new index."
+                )
+                # Clear only reporting-only nodes. Entry/index-stage nodes
+                # don't depend on the index, so their caches are still valid
+                _protected = set()
+                for _stage in (self.entry_stage, self.index_stage):
+                    if _stage is not None:
+                        _protected.add(_stage.name)
+                        _protected.update(n.name for n in _stage.dependencies)
+
+                _seen = set()
+
+                def _clear_reporting_only(node):
+                    if node.name in _protected or node.name in _seen:
+                        return
+                    _seen.add(node.name)
+                    Node._node_manager.clear_cache(node, con=con, recursive=False)
+                    for _child in node.children:
+                        _clear_reporting_only(_child)
+
+                for _node in list(self.characteristics or []) + list(
+                    self.outcomes or []
+                ):
+                    _clear_reporting_only(_node)
             logger.info(f"Cohort '{self.name}': executing reporting stage ...")
             self.reporting_stage.execute(
                 tables=self.subset_tables_index,

phenex/node_manager.py

@@ -58,16 +58,7 @@ def should_rerun(self, node, con) -> bool:
         """
         reasons = []
 
-        # Get current execution context
-        current_execution_params = self._get_execution_params(con)
-
-        # Get current node hash
-        current_hash = self._get_node_hash(node)
-
-        # Look up previous run with same name and execution context
-        last_hash = self._getlasthash(
-            node.name, execution_params=current_execution_params
-        )
+        current_hash, last_hash = self._compare_to_last_run(node, con)
 
         # Determine if node should rerun
         node_changed = current_hash != last_hash
@@ -89,6 +80,36 @@ def should_rerun(self, node, con) -> bool:
 
         return should_rerun
 
+    def node_changed(self, node, con) -> bool:
+        """Return True if the node's definition differs from its last cached run."""
+        current_hash, last_hash = self._compare_to_last_run(node, con)
+        return current_hash != last_hash
+
+    def _compare_to_last_run(self, node, con):
+        """
+        Look up the node's current hash and the hash of its last cached run.
+
+        Single source of truth for how a node is compared against its cache:
+        both should_rerun() and node_changed() go through here, so the two can
+        never disagree about what counts as changed. Returns both hashes rather
+        than a bool because should_rerun() distinguishes "never executed"
+        (last_hash is None) from "definition changed" when it logs.
+
+        Parameters:
+            node: The Node object to look up
+            con: Database connector object (determines execution context)
+
+        Returns:
+            tuple: (current_hash, last_hash); last_hash is None if this node has
+            never run in this execution context.
+        """
+        return (
+            self._get_node_hash(node),
+            self._getlasthash(
+                node.name, execution_params=self._get_execution_params(con)
+            ),
+        )
+
     def update_run_params(self, node, con) -> bool:
         """
         Update the run parameters for a node after execution.

phenex/test/cohort/test_cohort_lazy_execution.py

@@ -28,9 +28,11 @@
 from phenex.ibis_connect import DuckDBConnector
 from phenex.phenotypes import (
     AgePhenotype,
+    BinPhenotype,
     CodelistPhenotype,
     SexPhenotype,
     MeasurementPhenotype,
+    TimeRangePhenotype,
 )
 from phenex.reporting import TimeToEvent
 from phenex.test.cohort.test_mappings import (
@@ -343,6 +345,13 @@ def setUpClass(cls):
         cls._db_path = os.path.join(cls._tmpdir, "test_lazy.duckdb")
         cls._meta_db = os.path.join(cls._tmpdir, "phenex.db")
 
+        # Isolate node state
+        from phenex.node import Node
+        from phenex.node_manager import NodeManager
+
+        cls._orig_node_manager = Node._node_manager
+        Node._node_manager = NodeManager(db_name=cls._meta_db)
+
         cls.con = DuckDBConnector(DUCKDB_DEST_DATABASE=cls._db_path)
         cls.tables = _build_test_tables(cls.con)
         cls.cohort, cls.right_censor = _build_cohort(cls.tables)
@@ -351,6 +360,12 @@ def setUpClass(cls):
         # Enable debug logging so we can track execution
         logging.getLogger("phenex").setLevel(logging.DEBUG)
 
+    @classmethod
+    def tearDownClass(cls):
+        from phenex.node import Node
+
+        Node._node_manager = cls._orig_node_manager
+
     def test_01_first_execution_computes_everything(self):
         """First execution with lazy_execution=True should compute all nodes."""
         with _ExecutionTracker() as tracker:
@@ -497,13 +512,26 @@ def setUpClass(cls):
         cls._tmpdir = tempfile.mkdtemp()
         cls._db_path = os.path.join(cls._tmpdir, "test_lazy_subcohort.duckdb")
 
+        # Isolate node state from other test modules in the same session
+        from phenex.node import Node
+        from phenex.node_manager import NodeManager
+
+        cls._orig_node_manager = Node._node_manager
+        Node._node_manager = NodeManager(db_name=os.path.join(cls._tmpdir, "phenex.db"))
+
         cls.con = DuckDBConnector(DUCKDB_DEST_DATABASE=cls._db_path)
         cls.tables = _build_test_tables(cls.con)
         cls.cohort, cls.right_censor = _build_cohort(cls.tables)
         cls.subcohort = _build_subcohort(cls.cohort)
 
         logging.getLogger("phenex").setLevel(logging.DEBUG)
 
+    @classmethod
+    def tearDownClass(cls):
+        from phenex.node import Node
+
+        Node._node_manager = cls._orig_node_manager
+
     def test_01_parent_then_subcohort_first_run(self):
         """First run: parent cohort and subcohort should both fully execute."""
         with _ExecutionTracker() as tracker:
@@ -566,5 +594,139 @@ def test_02_parent_cached_subcohort_cached(self):
         self.assertIsNotNone(self.subcohort.index_table)
 
 
+class TestCohortLazyReporterIndexChange(unittest.TestCase):
+    """When an inclusion is added between two lazy executions, the
+    index shrinks, and the baseline-characteristic tables must be
+    recomputed against the new index."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls._tmpdir = tempfile.mkdtemp()
+        cls.con = DuckDBConnector(
+            DUCKDB_DEST_DATABASE=os.path.join(cls._tmpdir, "test_reporter_idx.duckdb")
+        )
+        cls.tables = _build_test_tables(cls.con)
+
+        from phenex.node import Node
+        from phenex.node_manager import NodeManager
+
+        cls._orig_node_manager = Node._node_manager
+        Node._node_manager = NodeManager(db_name=os.path.join(cls._tmpdir, "phenex.db"))
+        logging.getLogger("phenex").setLevel(logging.WARNING)
+
+    @classmethod
+    def tearDownClass(cls):
+        from phenex.node import Node
+
+        Node._node_manager = cls._orig_node_manager
+
+    def _make_cohort(self, with_extra_inclusion):
+        """Same entry + characteristics every time; only the inclusion list differs."""
+        entry = CodelistPhenotype(
+            name="entry_drug",
+            return_date="first",
+            codelist=Codelist(["d1"]).copy(use_code_type=False),
+            domain="DRUG_EXPOSURE",
+        )
+        inclusions = [
+            TimeRangePhenotype(
+                name="continuous_coverage",
+                relative_time_range=RelativeTimeRangeFilter(
+                    min_days=GreaterThanOrEqualTo(365), anchor_phenotype=entry
+                ),
+            )
+        ]
+        if with_extra_inclusion:
+            # Only patients P0-P6 have cond1, so adding this shrinks the cohort.
+            inclusions.append(
+                CodelistPhenotype(
+                    name="has_cond1",
+                    codelist=Codelist(["cond1"]).copy(use_code_type=False),
+                    domain="CONDITION_OCCURRENCE",
+                    relative_time_range=RelativeTimeRangeFilter(
+                        when="before",
+                        min_days=GreaterThanOrEqualTo(0),
+                        anchor_phenotype=entry,
+                    ),
+                )
+            )
+        age = AgePhenotype(name="age_", anchor_phenotype=entry)
+        return Cohort(
+            name="reporter_idx_cohort",  # SAME name both runs -> lazy matches node identities
+            entry_criterion=entry,
+            inclusions=inclusions,
+            exclusions=[],
+            characteristics=[age, BinPhenotype(name="age_category", phenotype=age)],
+            write_subset_tables_entry=False,
+            write_subset_tables_index=False,
+        )
+
+    @staticmethod
+    def _row_N(df, name):
+        for _, row in df.iterrows():
+            if str(row["Name"]).strip().lower() == name:
+                return int(str(row["N"]).replace(",", ""))
+        return None
+
+    def test_characteristics_recompute_when_index_changes(self):
+        # RUN 1: entry + coverage (all patients qualify)
+        c1 = self._make_cohort(with_extra_inclusion=False)
+        c1.execute(
+            tables=self.tables, con=self.con, overwrite=True, lazy_execution=True
+        )
+        n1 = self._row_N(c1.table1, "cohort")
+        age1 = self._row_N(c1.table1, "age")
+        self.assertEqual(
+            age1, n1, "RUN 1: everyone has an age, so Age count should equal N."
+        )
+
+        # RUN 2: add an inclusion that shrinks the cohort, re-execute lazily
+        c2 = self._make_cohort(with_extra_inclusion=True)
+        c2.execute(
+            tables=self.tables, con=self.con, overwrite=True, lazy_execution=True
+        )
+        n2 = self._row_N(c2.table1, "cohort")
+        age2 = self._row_N(c2.table1, "age")
+
+        self.assertLess(n2, n1, "adding the inclusion should shrink the cohort.")
+        # The bug leaves age2 == n1 (stale), so Age% = 100 * n1 / n2 > 100.
+        self.assertEqual(
+            age2,
+            n2,
+            f"Stale characteristic: Age count {age2} != new cohort N {n2}. "
+            f"Table1 Age% would be {100 * age2 / n2:.1f}% instead of 100%.",
+        )
+
+        # The invalidation must not touch the entry criterion
+        entry_tbl = next(
+            t
+            for t in self.con.dest_connection.list_tables()
+            if t.endswith("ENTRY_DRUG")
+        )
+        entry_n = self.con.dest_connection.table(entry_tbl).count().execute()
+        self.assertEqual(
+            entry_n,
+            n1,
+            f"Entry table poisoned: {entry_tbl} holds {entry_n} rows, expected {n1}. "
+            f"The cache invalidation recomputed the entry criterion against the "
+            f"index-restricted tables.",
+        )
+
+        # RUN 3: remove the inclusion again the cohort must grow back to n1.
+        c3 = self._make_cohort(with_extra_inclusion=False)
+        c3.execute(
+            tables=self.tables, con=self.con, overwrite=True, lazy_execution=True
+        )
+        n3 = self._row_N(c3.table1, "cohort")
+        age3 = self._row_N(c3.table1, "age")
+        self.assertEqual(
+            n3,
+            n1,
+            f"Cohort cannot grow back: N {n3} != {n1} after removing the inclusion. "
+            f"The entry population was permanently narrowed by the invalidation.",
+        )
+        self.assertEqual(age3, n3, "RUN 3: Age count should match restored cohort N.")
+
+
 if __name__ == "__main__":
     unittest.main()

phenex/test/reporting/test_reporter.py

@@ -291,17 +291,21 @@ def test_returns_absolute_path(self):
         reporter = SimpleReporter(decimal_places=1)
         reporter.execute()
 
-        with tempfile.TemporaryDirectory() as tmpdir:
-            # Use relative path
-            relative_path = "test.csv"
-            full_path = os.path.join(tmpdir, relative_path)
-
-            os.chdir(tmpdir)
-            filepath = reporter.to_csv(relative_path)
-
-            # Should return absolute path
-            assert os.path.isabs(filepath)
-            assert os.path.exists(filepath)
+        original_cwd = os.getcwd()
+        try:
+            with tempfile.TemporaryDirectory() as tmpdir:
+                # Use relative path
+                relative_path = "test.csv"
+                full_path = os.path.join(tmpdir, relative_path)
+
+                os.chdir(tmpdir)
+                filepath = reporter.to_csv(relative_path)
+
+                # Should return absolute path
+                assert os.path.isabs(filepath)
+                assert os.path.exists(filepath)
+        finally:
+            os.chdir(original_cwd)
 
 
 class TestReporterSubclassMustImplementExecute: