Feat/multiindex with subset options

docs/api/phenotypes/further_value_filter_phenotype.md

@@ -0,0 +1,3 @@
+# FurtherValueFilterPhenotype
+
+::: phenex.phenotypes.further_value_filter_phenotype

mkdocs.yml

@@ -20,6 +20,7 @@ nav:
           - EventPhenotype: api/phenotypes/event_phenotype.md
           - CodelistPhenotype: api/phenotypes/codelist_phenotype.md
           - MeasurementPhenotype: api/phenotypes/measurement_phenotype.md
+          - FurtherValueFilterPhenotype: api/phenotypes/further_value_filter_phenotype.md
           - AgePhenotype: api/phenotypes/age_phenotype.md
           - SexPhenotype: api/phenotypes/sex_phenotype.md
           - DeathPhenotype: api/phenotypes/death_phenotype.md

phenex/__init__.py

@@ -15,6 +15,7 @@
     CODETYPE_INFO,
     MeasurementPhenotype,
     MeasurementChangePhenotype,
+    FurtherValueFilterPhenotype,
     AgePhenotype,
     SexPhenotype,
     BinPhenotype,
@@ -134,6 +135,7 @@
     "CODETYPE_INFO",
     "MeasurementPhenotype",
     "MeasurementChangePhenotype",
+    "FurtherValueFilterPhenotype",
     "AgePhenotype",
     "SexPhenotype",
     "BinPhenotype",

phenex/aggregators/aggregator.py

@@ -79,10 +79,14 @@ def __init__(
         self.preserve_nulls = preserve_nulls
 
     def aggregate(self, input_table: Table):
+        # Ensure INDEX_DATE is in the aggregation index if the table has it
+        agg_index = list(self.aggregation_index)
+        if "INDEX_DATE" in input_table.columns and "INDEX_DATE" not in agg_index:
+            agg_index.append("INDEX_DATE")
         # Define the window specification
         partition_cols = [
             getattr(input_table, col) if isinstance(col, str) else col
-            for col in self.aggregation_index
+            for col in agg_index
         ]
         window_spec = ibis.window(
             group_by=partition_cols, order_by=input_table[self.event_date_column]
@@ -124,7 +128,7 @@ def aggregate(self, input_table: Table):
 
         # Apply the distinct reduction if required
         if self.reduce:
-            selected_columns = self.aggregation_index + [self.event_date_column]
+            selected_columns = list(agg_index) + [self.event_date_column]
             input_table = input_table.select(selected_columns).distinct()
             input_table = input_table.mutate(VALUE=ibis.null().cast("int32"))
 
@@ -255,10 +259,12 @@ def __init__(
         # otherwise, row count is preserved
 
     def aggregate(self, input_table: Table):
+        # Ensure INDEX_DATE is in the aggregation index if the table has it
+        agg_index = list(self.aggregation_index)
+        if "INDEX_DATE" in input_table.columns and "INDEX_DATE" not in agg_index:
+            agg_index.append("INDEX_DATE")
         # Get the aggregation index columns
-        _aggregation_index_cols = [
-            getattr(input_table, col) for col in self.aggregation_index
-        ]
+        _aggregation_index_cols = [getattr(input_table, col) for col in agg_index]
 
         # Determine the aggregation column
         if self.aggregation_column is None:

phenex/core/cohort.py

@@ -43,6 +43,8 @@ class Cohort:
         description: A plain text description of the cohort.
         data_period: Restrict all input data to a specific date range. The input data will be modified to look as if data outside the data_period was never recorded before any phenotypes are computed. See DataPeriodFilterNode for details on how the input data are affected by this parameter.
         custom_reporters: Additional reporter instances to run on this cohort only, after the default Waterfall and Table1 reporters. Each reporter must implement ``execute(cohort)`` and ``to_json(path)``.
+        write_subset_tables_entry: If True (default), materialize the entry-subset tables to the destination database. If False, keep them as lazy expressions instead.
+        write_subset_tables_index: If True (default), materialize the index-subset tables to the destination database. If False, keep them as lazy expressions instead.
 
     Attributes:
         table (PhenotypeTable): The resulting index table after filtering (None until execute is called)
@@ -67,10 +69,41 @@ def __init__(
         description: Optional[str] = None,
         database: Optional[Database] = None,
         custom_reporters: Optional[List] = None,
+        return_index: str = "first",
+        max_index_dates: Optional[int] = None,
+        write_subset_tables_entry: bool = True,
+        write_subset_tables_index: bool = True,
     ):
         self.name = name
         self.description = description
         self.database = database
+        self.return_index = return_index
+        self.max_index_dates = max_index_dates
+
+        assert return_index in (
+            "first",
+            "last",
+            "all",
+        ), f"return_index must be 'first', 'last', or 'all', got '{return_index}'"
+        if max_index_dates is not None:
+            assert (
+                isinstance(max_index_dates, int) and max_index_dates > 0
+            ), f"max_index_dates must be a positive integer, got {max_index_dates}"
+
+        # When return_index requires multiple candidate dates, auto-set entry criterion
+        if return_index in ("last", "all"):
+            if (
+                hasattr(entry_criterion, "return_date")
+                and entry_criterion.return_date != "all"
+            ):
+                logger.info(
+                    f"Cohort '{name}': return_index='{return_index}' requires entry criterion "
+                    f"return_date='all'. Auto-setting from '{entry_criterion.return_date}'."
+                )
+                entry_criterion.return_date = "all"
+
+        self.write_subset_tables_entry = write_subset_tables_entry
+        self.write_subset_tables_index = write_subset_tables_index
         self.table = None  # Will be set during execution to index table
         self.subset_tables_entry = None  # Will be set during execution
         self.subset_tables_index = None  # Will be set during execution
@@ -123,6 +156,7 @@ def __init__(
         self.derived_tables_stage = None
         self.entry_stage = None
         self.index_stage = None
+        self.subset_index_stage = None
         self.derived_tables_post_entry_stage = None
         self.reporting_stage = None
 
@@ -340,6 +374,8 @@ def build_stages(self, tables: Dict[str, PhenexTable]):
             entry_phenotype=self.entry_criterion,
             inclusion_table_node=self.inclusions_table_node,
             exclusion_table_node=self.exclusions_table_node,
+            return_index=self.return_index,
+            max_index_dates=self.max_index_dates,
         )
         index_nodes.append(self.index_table_node)
 
@@ -363,10 +399,26 @@ def build_stages(self, tables: Dict[str, PhenexTable]):
             domains=entry_domains,
             index_phenotype=self.index_table_node,
         )
-        self.index_stage = NodeGroup(
-            name="index_stage",
-            nodes=self.subset_tables_index_nodes + index_nodes,
-        )
+        if self.write_subset_tables_index:
+            # Default: materialize the index-subset tables together with the
+            # index nodes in a single multithreaded stage.
+            self.subset_index_stage = None
+            self.index_stage = NodeGroup(
+                name="index_stage",
+                nodes=self.subset_tables_index_nodes + index_nodes,
+            )
+        else:
+            # Keep the index-subset tables in a separate stage so they can be
+            # executed without materializing to the destination database
+            # (see execute()).
+            self.index_stage = NodeGroup(
+                name="index_stage",
+                nodes=index_nodes,
+            )
+            self.subset_index_stage = NodeGroup(
+                name="subset_index_stage",
+                nodes=self.subset_tables_index_nodes,
+            )
 
         #
         # Post-index / reporting stage: OPTIONAL
@@ -618,14 +670,55 @@ def execute(
 
         logger.info(f"Cohort '{self.name}': executing entry stage ...")
 
-        self.entry_stage.execute(
-            tables=tables,
-            con=con,
-            overwrite=overwrite,
-            n_threads=n_threads,
-            lazy_execution=lazy_execution,
-            table_name_prefix=self.name,
-        )
+        if self.write_subset_tables_entry:
+            self.entry_stage.execute(
+                tables=tables,
+                con=con,
+                overwrite=overwrite,
+                n_threads=n_threads,
+                lazy_execution=lazy_execution,
+                table_name_prefix=self.name,
+            )
+        else:
+            # Execute entry criterion in-memory so .table stays on the source
+            # backend, avoiding cross-backend joins with subset tables.
+            self.entry_criterion.execute(
+                tables=tables,
+                con=con,
+                overwrite=overwrite,
+                n_threads=n_threads,
+                table_name_prefix=self.name,
+                lazy_execution=lazy_execution,
+            )
+            # # Write entry criterion and its dependencies to dest
+            # if con is not None:
+            #     prefix = re.sub(r"[^A-Za-z0-9_]", "_", self.name).upper()
+            #     node = self.entry_criterion
+            #     # if node.table is not None:
+            #     #     db_name = (
+            #     #         f"{prefix}__{node.name}"
+            #     #         if not node.name.startswith(prefix)
+            #     #         else node.name
+            #     #     )
+            #     #     con.create_table(node.table, db_name, overwrite=overwrite)
+            # # Remove entry_criterion from subset table children so it won't be
+            # # re-executed; its .table is already set and SubsetTable._execute
+            # # accesses it via self.index_phenotype.table.
+            for node in self.subset_tables_entry_nodes:
+                node._children = [
+                    c for c in node._children if c is not self.entry_criterion
+                ]
+            self.entry_stage.execute(
+                tables=tables,
+                con=None,
+                overwrite=overwrite,
+                n_threads=n_threads,
+                table_name_prefix=self.name,
+            )
+            # Restore children for correct dependency graphs in later stages
+            for node in self.subset_tables_entry_nodes:
+                node._children.insert(0, self.entry_criterion)
+
         self.subset_tables_entry = tables = self.get_subset_tables_entry(tables)
 
         logger.info(f"Cohort '{self.name}': completed entry stage.")
@@ -666,6 +759,27 @@ def execute(
         )
         self.table = self.index_table_node.table
 
+        if not self.write_subset_tables_index:
+            # Execute the index-subset tables in-memory so they are not
+            # materialized to the destination database. The index table is
+            # already computed, so detach it from the subset nodes' children to
+            # avoid re-executing it; SubsetTable accesses it via
+            # self.index_phenotype.table.
+            for node in self.subset_tables_index_nodes:
+                node._children = [
+                    c for c in node._children if c is not self.index_table_node
+                ]
+            self.subset_index_stage.execute(
+                tables=self.subset_tables_entry,
+                con=None,
+                overwrite=overwrite,
+                n_threads=n_threads,
+                table_name_prefix=self.name,
+            )
+            # Restore children for correct dependency graphs in later stages
+            for node in self.subset_tables_index_nodes:
+                node._children.insert(0, self.index_table_node)
+
         logger.info(f"Cohort '{self.name}': completed index stage.")
         logger.info(f"Cohort '{self.name}': executing reporting stage ...")
 

phenex/core/exclusions_table_node.py

@@ -20,17 +20,27 @@ def __init__(
         self.index_phenotype = index_phenotype
 
     def _execute(self, tables: Dict[str, Table]):
-        exclusions_table = self.index_phenotype.table.select(["PERSON_ID"])
+        # Build base from entry criterion; add INDEX_DATE if ALL phenotype tables have it
+        if self.phenotypes and all(
+            "INDEX_DATE" in pt.table.columns for pt in self.phenotypes
+        ):
+            join_keys = ["PERSON_ID", "INDEX_DATE"]
+            exclusions_table = self.index_phenotype.table.mutate(
+                INDEX_DATE=self.index_phenotype.table.EVENT_DATE
+            ).select(["PERSON_ID", "INDEX_DATE"])
+        else:
+            join_keys = ["PERSON_ID"]
+            exclusions_table = self.index_phenotype.table.select(["PERSON_ID"])
 
         for pt in self.phenotypes:
-            pt_table = pt.table.select(["PERSON_ID", "BOOLEAN"]).rename(
+            pt_table = pt.table.select([*join_keys, "BOOLEAN"]).rename(
                 **{
                     f"{pt.name}_BOOLEAN": "BOOLEAN",
                 }
             )
-            exclusions_table = exclusions_table.left_join(pt_table, ["PERSON_ID"])
-            columns = exclusions_table.columns
-            columns.remove("PERSON_ID_right")
+            exclusions_table = exclusions_table.left_join(pt_table, join_keys)
+            drop_cols = [f"{k}_right" for k in join_keys]
+            columns = [c for c in exclusions_table.columns if c not in drop_cols]
             exclusions_table = exclusions_table.select(columns)
 
         # fill all nones with False

phenex/core/inclusions_table_node.py

@@ -20,17 +20,28 @@ def __init__(
         self.index_phenotype = index_phenotype
 
     def _execute(self, tables: Dict[str, Table]):
-        inclusions_table = self.index_phenotype.table.select(["PERSON_ID"])
+        # Build base from entry criterion; add INDEX_DATE if ALL phenotype tables have it
+        if self.phenotypes and all(
+            "INDEX_DATE" in pt.table.columns for pt in self.phenotypes
+        ):
+            join_keys = ["PERSON_ID", "INDEX_DATE"]
+            # Derive (PERSON_ID, INDEX_DATE) pairs from entry criterion
+            inclusions_table = self.index_phenotype.table.mutate(
+                INDEX_DATE=self.index_phenotype.table.EVENT_DATE
+            ).select(["PERSON_ID", "INDEX_DATE"])
+        else:
+            join_keys = ["PERSON_ID"]
+            inclusions_table = self.index_phenotype.table.select(["PERSON_ID"])
 
         for pt in self.phenotypes:
-            pt_table = pt.table.select(["PERSON_ID", "BOOLEAN"]).rename(
+            pt_table = pt.table.select([*join_keys, "BOOLEAN"]).rename(
                 **{
                     f"{pt.name}_BOOLEAN": "BOOLEAN",
                 }
             )
-            inclusions_table = inclusions_table.left_join(pt_table, ["PERSON_ID"])
-            columns = inclusions_table.columns
-            columns.remove("PERSON_ID_right")
+            inclusions_table = inclusions_table.left_join(pt_table, join_keys)
+            drop_cols = [f"{k}_right" for k in join_keys]
+            columns = [c for c in inclusions_table.columns if c not in drop_cols]
             inclusions_table = inclusions_table.select(columns)
 
         # fill all nones with False