diff --git a/CHANGELOG.md b/CHANGELOG.md index 75ea0e58..07eac21b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,8 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - +## [0.3.0] Unreleased +- Support for OMOP extension tables with extension classes and builders ## [0.2.0] - 2026-02-25 diff --git a/circe/execution/builders/registry.py b/circe/execution/builders/registry.py index 68a0a633..4c889623 100644 --- a/circe/execution/builders/registry.py +++ b/circe/execution/builders/registry.py @@ -21,10 +21,25 @@ def decorator(func: Callable[[Criteria, BuildContext], ir.Table]): def get_builder(criteria: Criteria): name = criteria.__class__.__name__ + builder = _REGISTRY.get(name) + if builder is not None: + return builder + + # Fall back to the global extension registry so that dynamically + # registered ibis builders (via circe.extensions.ibis_builder or + # ExtensionRegistry.register_ibis_builder) are discovered without + # any hard-coded imports of extension modules. try: - return _REGISTRY[name] - except KeyError as exc: - raise ValueError(f"No builder registered for criteria {name}") from exc + from ...extensions import get_registry + + builder = get_registry().get_ibis_builder(name) + except ImportError: + builder = None + + if builder is not None: + return builder + + raise ValueError(f"No builder registered for criteria {name}") def build_events(criteria: Criteria, ctx: BuildContext) -> ir.Table: diff --git a/circe/extensions/__init__.py b/circe/extensions/__init__.py index 5ba69400..9a49b49f 100644 --- a/circe/extensions/__init__.py +++ b/circe/extensions/__init__.py @@ -3,14 +3,14 @@ This module provides the central registry for managing extensions to circe-py, allowing external projects to register custom criteria classes, SQL builders, -and markdown renderers. +ibis execution builders, and markdown renderers. Decorator Usage --------------- Extension authors can use the provided decorator functions to register their classes automatically, rather than calling the registry methods directly:: - from circe.extensions import criteria_class, sql_builder, markdown_template + from circe.extensions import criteria_class, sql_builder, markdown_template, ibis_builder @criteria_class("WaveformOccurrence") class WaveformOccurrence(Criteria): @@ -23,6 +23,10 @@ class WaveformOccurrenceSqlBuilder(CriteriaSqlBuilder): @markdown_template(WaveformOccurrence, "waveform_occurrence.j2") class WaveformOccurrenceMarkdownRenderer: ... + + @ibis_builder("WaveformOccurrence") + def build_waveform_occurrence(criteria, ctx): + ... """ from pathlib import Path @@ -49,6 +53,9 @@ def __init__(self): # Maps criteria types to markdown template names self._markdown_templates: dict[type[Criteria], str] = {} + # Maps criteria names to ibis execution builders + self._ibis_builders: dict[str, Callable] = {} + # List of paths to search for Jinja2 templates self._template_paths: list[Path] = [] @@ -83,6 +90,29 @@ def register_markdown_template(self, criteria_cls: type["Criteria"], template_na """ self._markdown_templates[criteria_cls] = template_name + def register_ibis_builder(self, criteria_name: str, func: Callable) -> None: + """Register an ibis execution builder for a criteria type. + + The callable must accept ``(criteria, build_context)`` and return an + ``ibis.expr.types.Table``. + + Args: + criteria_name: The criteria class name (e.g. ``"WaveformOccurrence"``). + func: A callable ``(Criteria, BuildContext) -> ibis.Table``. + """ + self._ibis_builders[criteria_name] = func + + def get_ibis_builder(self, criteria_name: str) -> Optional[Callable]: + """Look up a registered ibis execution builder by criteria class name. + + Args: + criteria_name: The criteria class name. + + Returns: + The registered callable, or ``None`` if not found. + """ + return self._ibis_builders.get(criteria_name) + def add_template_path(self, path: Path) -> None: """Add a path to search for Jinja2 templates. @@ -225,3 +255,41 @@ def template_path(path: Union[str, Path]) -> None: template_path(Path(__file__).parent / "templates") """ _registry.add_template_path(Path(path)) + + +def ibis_builder(criteria_name: str) -> Callable: + """Function decorator that registers an ibis execution builder for a criteria type. + + The decorated function must accept ``(criteria, build_context)`` and return + an ``ibis.expr.types.Table`` following the standard pipeline contract + (person_id, event_id, start_date, end_date, visit_occurrence_id). + + This also inserts the builder into the low-level execution registry + (``circe.execution.builders.registry``) so that ``build_events`` can + discover it without any hard-coded imports. + + Args: + criteria_name: The criteria class name (e.g. ``"WaveformOccurrence"``). + + Example:: + + from circe.extensions import ibis_builder + + @ibis_builder("WaveformOccurrence") + def build_waveform_occurrence(criteria, ctx): + ... + """ + + def decorator(func: Callable) -> Callable: + _registry.register_ibis_builder(criteria_name, func) + # Also push into the low-level execution registry so build_events + # can resolve the builder without the fallback path. + try: + from circe.execution.builders.registry import register as _register_exec + + _register_exec(criteria_name)(func) + except ImportError: + pass + return func + + return decorator diff --git a/circe/extensions/waveform/builders/__init__.py b/circe/extensions/waveform/builders/__init__.py index 622b5a40..1fa74f63 100644 --- a/circe/extensions/waveform/builders/__init__.py +++ b/circe/extensions/waveform/builders/__init__.py @@ -1 +1,19 @@ -"""builders sub-package for the waveform extension.""" +"""builders sub-package for the waveform extension. + +Importing this package registers both the SQL builders (via @sql_builder / +@markdown_template decorators) and the ibis execution builders (via @register) +for all four waveform criteria types. +""" + +# SQL / markdown builders — decorators fire on import +# Ibis execution builders — @register decorators fire on import +from . import ( + ibis_waveform_channel_metadata, # noqa: F401 + ibis_waveform_feature, # noqa: F401 + ibis_waveform_occurrence, # noqa: F401 + ibis_waveform_registry, # noqa: F401 + waveform_channel_metadata, # noqa: F401 + waveform_feature, # noqa: F401 + waveform_occurrence, # noqa: F401 + waveform_registry, # noqa: F401 +) diff --git a/circe/extensions/waveform/builders/ibis_waveform_channel_metadata.py b/circe/extensions/waveform/builders/ibis_waveform_channel_metadata.py new file mode 100644 index 00000000..e5dd8d96 --- /dev/null +++ b/circe/extensions/waveform/builders/ibis_waveform_channel_metadata.py @@ -0,0 +1,83 @@ +"""Ibis execution builder for WaveformChannelMetadata criteria.""" + +from __future__ import annotations + +from circe.execution.build_context import BuildContext +from circe.execution.builders.common import ( + apply_concept_filters, + apply_numeric_range, + apply_text_filter, + standardize_output, +) +from circe.execution.builders.groups import apply_criteria_group +from circe.extensions import ibis_builder + +from ..criteria import WaveformChannelMetadata + + +@ibis_builder("WaveformChannelMetadata") +def build_waveform_channel_metadata(criteria: WaveformChannelMetadata, ctx: BuildContext): + """Build an ibis event table from waveform_channel_metadata. + + Channel metadata rows have no timestamps of their own; start_date/end_date + are sourced from the parent waveform_registry file bounds via a join. + person_id is resolved through the registry → occurrence chain. + """ + table = ctx.table("waveform_channel_metadata") + + # Join waveform_registry to get file dates and occurrence link + registry = ctx.table("waveform_registry").select( + "waveform_registry_id", + "waveform_occurrence_id", + "file_start_datetime", + "file_end_datetime", + ) + table = table.join( + registry, + table.waveform_registry_id == registry.waveform_registry_id, + ) + + # Join waveform_occurrence to get person_id and visit context + occurrence = ctx.table("waveform_occurrence").select( + "waveform_occurrence_id", + "person_id", + "visit_occurrence_id", + ) + table = table.join( + occurrence, + table.waveform_occurrence_id == occurrence.waveform_occurrence_id, + ) + + # Registry link filter + table = apply_numeric_range(table, "waveform_registry_id", criteria.waveform_registry_id) + + # Channel identification + if criteria.channel_concept_id: + table = apply_concept_filters(table, "channel_concept_id", criteria.channel_concept_id) + table = apply_text_filter(table, "waveform_channel_source_value", criteria.waveform_channel_source_value) + + # Metadata type + if criteria.metadata_concept_id: + table = apply_concept_filters(table, "metadata_concept_id", criteria.metadata_concept_id) + table = apply_text_filter(table, "metadata_source_value", criteria.metadata_source_value) + + # Metadata values + table = apply_numeric_range(table, "value_as_number", criteria.value_as_number) + if criteria.value_as_concept_id: + table = apply_concept_filters(table, "value_as_concept_id", criteria.value_as_concept_id) + + # Units + if criteria.unit_concept_id: + table = apply_concept_filters(table, "unit_concept_id", criteria.unit_concept_id) + + # Device / procedure linkage + table = apply_numeric_range(table, "device_exposure_id", criteria.device_exposure_id) + table = apply_numeric_range(table, "procedure_occurrence_id", criteria.procedure_occurrence_id) + + events = standardize_output( + table, + primary_key=criteria.get_primary_key_column(), + start_column=criteria.get_start_date_column(), + end_column=criteria.get_end_date_column(), + ) + return apply_criteria_group(events, criteria.correlated_criteria, ctx) diff --git a/circe/extensions/waveform/builders/ibis_waveform_feature.py b/circe/extensions/waveform/builders/ibis_waveform_feature.py new file mode 100644 index 00000000..0692595d --- /dev/null +++ b/circe/extensions/waveform/builders/ibis_waveform_feature.py @@ -0,0 +1,77 @@ +"""Ibis execution builder for WaveformFeature criteria.""" + +from __future__ import annotations + +from circe.execution.build_context import BuildContext +from circe.execution.builders.common import ( + apply_concept_filters, + apply_date_range, + apply_numeric_range, + apply_text_filter, + standardize_output, +) +from circe.execution.builders.groups import apply_criteria_group +from circe.extensions import ibis_builder + +from ..criteria import WaveformFeature + + +@ibis_builder("WaveformFeature") +def build_waveform_feature(criteria: WaveformFeature, ctx: BuildContext): + """Build an ibis event table from waveform_feature. + + waveform_feature stores derived measurements (heart rate, SpO2, arrhythmia + detections, AI embeddings, etc.). person_id and visit_occurrence_id are + resolved by joining to waveform_occurrence. + """ + table = ctx.table("waveform_feature") + + # Join waveform_occurrence to obtain person_id and visit context + occurrence = ctx.table("waveform_occurrence").select( + "waveform_occurrence_id", + "person_id", + "visit_occurrence_id", + ) + table = table.join( + occurrence, + table.waveform_occurrence_id == occurrence.waveform_occurrence_id, + ) + + # Parent link filters + table = apply_numeric_range(table, "waveform_occurrence_id", criteria.waveform_occurrence_id) + table = apply_numeric_range(table, "waveform_registry_id", criteria.waveform_registry_id) + table = apply_numeric_range(table, "waveform_channel_metadata_id", criteria.waveform_channel_metadata_id) + + # Feature type (e.g., heart rate, SpO2, QRS) + if criteria.feature_concept_id: + table = apply_concept_filters(table, "feature_concept_id", criteria.feature_concept_id) + + # Algorithm used to derive feature + if criteria.algorithm_concept_id: + table = apply_concept_filters(table, "algorithm_concept_id", criteria.algorithm_concept_id) + table = apply_text_filter(table, "algorithm_source_value", criteria.algorithm_source_value) + + # Temporal window for feature + table = apply_date_range(table, "waveform_feature_start_timestamp", criteria.feature_start_timestamp) + table = apply_date_range(table, "waveform_feature_end_timestamp", criteria.feature_end_timestamp) + + # Feature values + table = apply_numeric_range(table, "value_as_number", criteria.value_as_number) + if criteria.value_as_concept_id: + table = apply_concept_filters(table, "value_as_concept_id", criteria.value_as_concept_id) + + # Units + if criteria.unit_concept_id: + table = apply_concept_filters(table, "unit_concept_id", criteria.unit_concept_id) + + # Links to standard OMOP tables + table = apply_numeric_range(table, "measurement_id", criteria.measurement_id) + table = apply_numeric_range(table, "observation_id", criteria.observation_id) + + events = standardize_output( + table, + primary_key=criteria.get_primary_key_column(), + start_column=criteria.get_start_date_column(), + end_column=criteria.get_end_date_column(), + ) + return apply_criteria_group(events, criteria.correlated_criteria, ctx) diff --git a/circe/extensions/waveform/builders/ibis_waveform_occurrence.py b/circe/extensions/waveform/builders/ibis_waveform_occurrence.py new file mode 100644 index 00000000..342f6046 --- /dev/null +++ b/circe/extensions/waveform/builders/ibis_waveform_occurrence.py @@ -0,0 +1,75 @@ +"""Ibis execution builder for WaveformOccurrence criteria.""" + +from __future__ import annotations + +from circe.execution.build_context import BuildContext +from circe.execution.builders.common import ( + apply_concept_filters, + apply_date_range, + apply_numeric_range, + apply_text_filter, + standardize_output, +) +from circe.execution.builders.groups import apply_criteria_group +from circe.extensions import ibis_builder + +from ..criteria import WaveformOccurrence + + +@ibis_builder("WaveformOccurrence") +def build_waveform_occurrence(criteria: WaveformOccurrence, ctx: BuildContext): + """Build an ibis event table from waveform_occurrence. + + Output columns match the standard pipeline contract: + person_id, event_id, start_date, end_date, visit_occurrence_id. + """ + table = ctx.table("waveform_occurrence") + + # Concept filter + if criteria.waveform_occurrence_concept_id: + table = apply_concept_filters( + table, + "waveform_occurrence_concept_id", + criteria.waveform_occurrence_concept_id, + ) + + # Temporal bounds + table = apply_date_range( + table, + "waveform_occurrence_start_datetime", + criteria.occurrence_start_datetime, + ) + table = apply_date_range( + table, + "waveform_occurrence_end_datetime", + criteria.occurrence_end_datetime, + ) + + # Visit context + table = apply_numeric_range(table, "visit_occurrence_id", criteria.visit_occurrence_id) + table = apply_numeric_range(table, "visit_detail_id", criteria.visit_detail_id) + + # File count + table = apply_numeric_range(table, "num_of_files", criteria.num_of_files) + + # Source value text filter + table = apply_text_filter( + table, + "waveform_occurrence_source_value", + criteria.waveform_occurrence_source_value, + ) + + # Sequence/chain filtering + table = apply_numeric_range( + table, + "preceding_waveform_occurrence_id", + criteria.preceding_waveform_occurrence_id, + ) + + events = standardize_output( + table, + primary_key=criteria.get_primary_key_column(), + start_column=criteria.get_start_date_column(), + end_column=criteria.get_end_date_column(), + ) + return apply_criteria_group(events, criteria.correlated_criteria, ctx) diff --git a/circe/extensions/waveform/builders/ibis_waveform_registry.py b/circe/extensions/waveform/builders/ibis_waveform_registry.py new file mode 100644 index 00000000..988328fb --- /dev/null +++ b/circe/extensions/waveform/builders/ibis_waveform_registry.py @@ -0,0 +1,77 @@ +"""Ibis execution builder for WaveformRegistry criteria.""" + +from __future__ import annotations + +from circe.execution.build_context import BuildContext +from circe.execution.builders.common import ( + apply_concept_filters, + apply_date_range, + apply_numeric_range, + apply_text_filter, + standardize_output, +) +from circe.execution.builders.groups import apply_criteria_group +from circe.extensions import ibis_builder + +from ..criteria import WaveformRegistry + + +@ibis_builder("WaveformRegistry") +def build_waveform_registry(criteria: WaveformRegistry, ctx: BuildContext): + """Build an ibis event table from waveform_registry. + + waveform_registry rows have file-level temporal bounds but no person_id + directly — person_id is carried via the waveform_occurrence join. + The output follows the standard pipeline contract via standardize_output. + """ + table = ctx.table("waveform_registry") + + # Join to waveform_occurrence to obtain person_id and visit context. + occurrence = ctx.table("waveform_occurrence").select( + "waveform_occurrence_id", + "person_id", + "visit_occurrence_id", + ) + base_columns = list(table.columns) + ["person_id", "visit_occurrence_id"] + table = table.join( + occurrence, + table.waveform_occurrence_id == occurrence.waveform_occurrence_id, + ) + # Keep only required columns to avoid ambiguity + keep = [c for c in base_columns if c in table.columns] + # deduplicate while preserving order + seen: set[str] = set() + unique_keep: list[str] = [] + for c in keep: + if c not in seen: + seen.add(c) + unique_keep.append(c) + table = table.select(unique_keep) + + # Parent occurrence link + table = apply_numeric_range(table, "waveform_occurrence_id", criteria.waveform_occurrence_id) + + # File temporal bounds + table = apply_date_range(table, "file_start_datetime", criteria.file_start_datetime) + table = apply_date_range(table, "file_end_datetime", criteria.file_end_datetime) + + # File format + if criteria.file_extension_concept_id: + table = apply_concept_filters( + table, + "file_extension_concept_id", + criteria.file_extension_concept_id, + ) + table = apply_text_filter(table, "file_extension_source_value", criteria.file_extension_source_value) + + # Visit context (denormalized) + table = apply_numeric_range(table, "visit_occurrence_id", criteria.visit_occurrence_id) + table = apply_numeric_range(table, "visit_detail_id", criteria.visit_detail_id) + + events = standardize_output( + table, + primary_key=criteria.get_primary_key_column(), + start_column=criteria.get_start_date_column(), + end_column=criteria.get_end_date_column(), + ) + return apply_criteria_group(events, criteria.correlated_criteria, ctx) diff --git a/docs/_static/.gitkeep b/docs/_static/.gitkeep deleted file mode 100644 index 8b137891..00000000 --- a/docs/_static/.gitkeep +++ /dev/null @@ -1 +0,0 @@ - diff --git a/docs/developer/extensions.rst b/docs/developer/extensions.rst index 69904f52..16b09cbd 100644 --- a/docs/developer/extensions.rst +++ b/docs/developer/extensions.rst @@ -1,34 +1,61 @@ Extending circe_py =================== -This guide explains how to extend `circe_py` with custom criteria types. This is useful when you have data in your CDM that isn't part of the standard OMOP domains (e.g., weather data, genomic features, or specialized clinical registries). +This guide explains how to extend ``circe_py`` with custom criteria types. +This is useful when you have data in your CDM that isn't part of the standard +OMOP domains (e.g., weather data, genomic features, or specialized clinical +registries). Architecture Overview --------------------- -The extension system consists of three main components: +The extension system consists of four main components: -1. **Criteria Class**: A Pydantic model that defines the fields available in your new criteria. -2. **SQL Builder**: A class that translates your criteria into SQL. -3. **Markdown Template**: A Jinja2 template that generates a human-readable description. +1. **Criteria Class** – A Pydantic model that defines the fields available in your new criteria. +2. **SQL Builder** – A class that translates your criteria into template-based SQL strings. +3. **Ibis Execution Builder** – A function that builds an ``ibis`` table expression so the criteria can be executed against a live database through the ibis execution layer. +4. **Markdown Template** – A Jinja2 template that generates a human-readable description. -Registration is handled by the `ExtensionRegistry`. +All four components are registered through the ``ExtensionRegistry``. Registration +can be done either **programmatically** (calling methods on the registry) or +**declaratively** (using decorator helpers). Both approaches are shown below. + +.. important:: + + Every component is registered dynamically at runtime—there must be no + hard-coded imports of extension types inside the core library. This means + an extension only needs to be *imported* (or call its registration function) + for its criteria to become available everywhere: SQL generation, ibis + execution, JSON round-tripping, and markdown rendering. Example: Weather Conditions --------------------------- -Imagine you want to create a cohort based on weather conditions (e.g., "Patients diagnosed with asthma during extreme cold"). +Imagine you want to create a cohort based on weather conditions (e.g., +"Patients diagnosed with asthma during extreme cold"). The extension assumes a +custom ``weather_data`` table in your CDM with the following columns: + +* ``weather_id`` (PK) +* ``person_id`` +* ``weather_concept_id`` +* ``temp_c`` +* ``observation_date`` +* ``visit_occurrence_id`` (nullable) + Step 1: Define the Criteria Class ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Your class must inherit from `circe.cohortdefinition.criteria.Criteria`. Use Pydantic's `Field` and `AliasChoices` to maintain compatibility with both Pythonic (`snake_case`) and Java-style (`PascalCase`) field names. +Your class must inherit from ``circe.cohortdefinition.criteria.Criteria``. Use +Pydantic's ``Field`` and ``AliasChoices`` to maintain compatibility with both +Pythonic (``snake_case``) and Java-style (``PascalCase``) field names. .. code-block:: python from typing import Optional, List from pydantic import Field, AliasChoices from circe.cohortdefinition.criteria import Criteria, CriteriaGroup + from circe.cohortdefinition.core import NumericRange from circe.vocabulary.concept import Concept class WeatherCondition(Criteria): @@ -44,13 +71,24 @@ Your class must inherit from `circe.cohortdefinition.criteria.Criteria`. Use Pyd serialization_alias="TemperatureCelsius" ) + # --- helpers used by the ibis builder --- + def get_primary_key_column(self) -> str: + return "weather_id" + + def get_start_date_column(self) -> str: + return "observation_date" + + def get_end_date_column(self) -> str: + return "observation_date" + # Resolve forward references (required for complex criteria types) WeatherCondition.model_rebuild() Step 2: Implement the SQL Builder ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -The SQL builder must inherit from `circe.cohortdefinition.builders.base.CriteriaSqlBuilder`. +The SQL builder must inherit from +``circe.cohortdefinition.builders.base.CriteriaSqlBuilder``. .. code-block:: python @@ -93,10 +131,98 @@ The SQL builder must inherit from `circe.cohortdefinition.builders.base.Criteria query = query.replace("@whereClause", " AND ".join(where_clauses)) return query -Step 3: Register the Extension +Step 3: Implement the Ibis Execution Builder +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The ibis execution builder is a plain function that receives the criteria +instance and a ``BuildContext``, and returns an ``ibis.expr.types.Table`` that +conforms to the **standard pipeline output contract**: + +.. list-table:: + :header-rows: 1 + + * - Column + - Type + * - ``person_id`` + - ``int64`` + * - ``event_id`` + - ``int64`` + * - ``start_date`` + - ``timestamp`` + * - ``end_date`` + - ``timestamp`` + * - ``visit_occurrence_id`` + - ``int64`` (nullable) + +Use the shared helpers in ``circe.execution.builders.common`` to apply +filters—these handle ``NumericRange``, ``DateRange``, ``TextFilter``, and +concept list comparisons consistently. + +.. code-block:: python + + from circe.execution.build_context import BuildContext + from circe.execution.builders.common import ( + apply_concept_filters, + apply_numeric_range, + standardize_output, + ) + from circe.execution.builders.groups import apply_criteria_group + + def build_weather_condition(criteria: WeatherCondition, ctx: BuildContext): + """Build an ibis event table from weather_data.""" + table = ctx.table("weather_data") + + # Apply concept filter on weather type + if criteria.weather_concept_id: + table = apply_concept_filters( + table, + "weather_concept_id", + criteria.weather_concept_id, + ) + + # Apply temperature threshold + if criteria.temperature_celsius is not None: + table = table.filter(table.temp_c >= criteria.temperature_celsius) + + # Project to the standard output contract + events = standardize_output( + table, + primary_key=criteria.get_primary_key_column(), + start_column=criteria.get_start_date_column(), + end_column=criteria.get_end_date_column(), + ) + + # Handle any correlated criteria (inclusion/exclusion sub-groups) + return apply_criteria_group(events, criteria.correlated_criteria, ctx) + +``standardize_output`` renames and casts columns to match the pipeline +contract. ``apply_criteria_group`` recursively processes nested +``CriteriaGroup`` filters attached to your criteria, just as built-in criteria +do. + +Step 4: Create a Markdown Template +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Create a file named ``weather_condition.j2``: + +.. code-block:: jinja + + Weather condition: {{ criteria.weather_concept_id[0].concept_name if criteria.weather_concept_id else 'Any' }} + {% if criteria.temperature_celsius %} with temperature >= {{ criteria.temperature_celsius }}°C{% endif %}. + + +Step 5: Register the Extension ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Use the extension registry to link your classes and templates. +All four components must be registered before they can be used. There are two +equivalent approaches. + +Option A: Programmatic Registration +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Call the registry methods directly. This is useful when you want full control +over the registration order or need to register at a specific point in your +application startup. .. code-block:: python @@ -105,51 +231,161 @@ Use the extension registry to link your classes and templates. def register_weather_extension(): registry = get_registry() - - # 1. Register the Criteria Class + + # 1. Criteria class (JSON deserialization) registry.register_criteria_class("WeatherCondition", WeatherCondition) - - # 2. Register the SQL Builder + + # 2. SQL builder (template-based SQL generation) registry.register_sql_builder(WeatherCondition, WeatherConditionSqlBuilder) - - # 3. Register Markdown Template - # Ensure templates/weather_condition.j2 exists - template_path = Path(__file__).parent / "templates" - registry.add_template_path(template_path) + + # 3. Ibis builder (ibis execution layer) + registry.register_ibis_builder("WeatherCondition", build_weather_condition) + + # 4. Markdown template + tpl_path = Path(__file__).parent / "templates" + registry.add_template_path(tpl_path) registry.register_markdown_template(WeatherCondition, "weather_condition.j2") -Step 4: Create a Markdown Template -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +Option B: Decorator Registration +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -Create a file named `weather_condition.j2`: +Use the decorator helpers for a more declarative style. Each decorator fires +at *import time*, so merely importing the module is enough to register +everything. -.. code-block:: jinja +.. code-block:: python + + from circe.extensions import criteria_class, sql_builder, ibis_builder, markdown_template, template_path + from pathlib import Path + + @criteria_class("WeatherCondition") + class WeatherCondition(Criteria): + # ... fields as shown in Step 1 ... + pass + + @sql_builder(WeatherCondition) + class WeatherConditionSqlBuilder(CriteriaSqlBuilder[WeatherCondition]): + # ... methods as shown in Step 2 ... + pass + + @ibis_builder("WeatherCondition") + def build_weather_condition(criteria, ctx): + # ... implementation as shown in Step 3 ... + pass + + @markdown_template(WeatherCondition, "weather_condition.j2") + class WeatherConditionRenderer: + pass + + template_path(Path(__file__).parent / "templates") + +The ``@ibis_builder`` decorator simultaneously registers the function in +**both** the ``ExtensionRegistry`` (for discoverability via +``get_ibis_builder``) and the low-level execution registry (so +``build_events`` resolves it without any hard-coded imports). + +.. note:: + + The ``@ibis_builder`` decorator accepts the criteria class **name** (a + string), not the class itself. This matches the key used by ``get_builder`` + at execution time. - Weather condition: {{ criteria.weather_concept_id[0].concept_name if criteria.weather_concept_id else 'Any' }} - {% if criteria.temperature_celsius %} with temperature >= {{ criteria.temperature_celsius }}°C{% endif %}. Full End-to-End Usage --------------------- -Once registered, you can use your custom criteria just like any built-in type. +Once registered, you can use your custom criteria just like any built-in type +for SQL generation, markdown rendering, JSON round-tripping, *and* ibis +execution. .. code-block:: python from circe.cohortdefinition import CohortExpression, PrimaryCriteria from circe.vocabulary.concept import Concept - # Setup + # Ensure the extension is registered (import or call your registration function) register_weather_extension() # Define cohort weather_criteria = WeatherCondition( weather_concept_id=[Concept(concept_id=123, concept_name="Snowing")], - temperature_celsius=-5.0 + temperature_celsius=-5.0, ) - + cohort = CohortExpression( primary_criteria=PrimaryCriteria(criteria_list=[weather_criteria]) ) - # generate SQL or Markdown as usual - # ... + # --- SQL generation (template-based) --- + from circe.cohortdefinition.cohort_expression_query_builder import ( + CohortExpressionQueryBuilder, BuildExpressionQueryOptions, + ) + builder = CohortExpressionQueryBuilder() + sql_options = BuildExpressionQueryOptions() + sql_options.cdm_schema = "my_cdm" + sql = builder.build_expression_query(cohort, sql_options) + + # --- Ibis execution --- + import ibis + from circe.execution import IbisExecutor, ExecutionOptions + + conn = ibis.duckdb.connect("my_cdm.ddb") + options = ExecutionOptions(cdm_schema="main") + with IbisExecutor(conn, options) as executor: + result = executor.to_polars(cohort) + print(result) + + +How Ibis Builder Discovery Works +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +When ``build_events`` (the core entry point of the ibis pipeline) receives a +criteria instance, it resolves the builder in two steps: + +1. **Module-level registry** – checked first; this is where the built-in OMOP + domain builders (``ConditionOccurrence``, ``DrugExposure``, etc.) live. +2. **Extension registry fallback** – if the name is not found above, the global + ``ExtensionRegistry`` is consulted via ``get_ibis_builder(name)``. + +This two-tier lookup means that extensions never need to be hard-coded +anywhere in the core library. As long as the builder has been registered +(either via ``@ibis_builder`` or ``register_ibis_builder``), it will be +discovered at execution time. + +.. code-block:: text + + build_events(criteria, ctx) + └─ get_builder(criteria) + ├─ 1. _REGISTRY[criteria.__class__.__name__] ← built-in domains + └─ 2. ExtensionRegistry.get_ibis_builder(name) ← extensions + + +Available Common Helpers +~~~~~~~~~~~~~~~~~~~~~~~~ + +The ``circe.execution.builders.common`` module provides shared filter functions +that extension ibis builders should use for consistency: + +.. list-table:: + :header-rows: 1 + :widths: 40 60 + + * - Function + - Purpose + * - ``standardize_output(table, *, primary_key, start_column, end_column)`` + - Project and rename columns to the standard pipeline output contract. + * - ``apply_date_range(table, column, date_range)`` + - Filter rows by a ``DateRange`` (before/after/between). + * - ``apply_numeric_range(table, column, numeric_range)`` + - Filter rows by a ``NumericRange`` (eq, gt, lt, between, etc.). + * - ``apply_text_filter(table, column, text_filter)`` + - Filter rows by a ``TextFilter`` (contains, startswith, etc.). + * - ``apply_concept_filters(table, column, concepts)`` + - Filter rows where a concept column matches a list of ``Concept`` objects. + * - ``apply_codeset_filter(table, column, codeset_id, ctx)`` + - Filter rows by a pre-compiled concept set from the ``BuildContext``. + * - ``apply_criteria_group(events, group, ctx)`` + - Recursively apply correlated ``CriteriaGroup`` inclusion/exclusion logic. + +See the built-in builders (e.g., ``condition_occurrence.py``) or the waveform +extension builders for real-world usage examples. diff --git a/tests/test_ibis_builder_registration.py b/tests/test_ibis_builder_registration.py new file mode 100644 index 00000000..7400c701 --- /dev/null +++ b/tests/test_ibis_builder_registration.py @@ -0,0 +1,166 @@ +""" +Tests for dynamic ibis execution builder registration via the extension registry. + +Verifies that: +1. The @ibis_builder decorator registers builders in both the extension registry + and the low-level execution registry. +2. ExtensionRegistry.register_ibis_builder works for programmatic registration. +3. get_builder falls back to the extension registry when a criteria name is not + in the module-level _REGISTRY. +4. No hardcoding of waveform or other extension types is required. +""" + +import pytest + +from circe.cohortdefinition.criteria import Criteria +from circe.execution.builders.registry import _REGISTRY, get_builder +from circe.extensions import get_registry, ibis_builder + +# --------------------------------------------------------------------------- +# Decorator-based registration +# --------------------------------------------------------------------------- + + +class TestIbisBuilderDecorator: + """@ibis_builder decorator registers a function in both registries.""" + + def test_decorator_registers_in_extension_registry(self): + @ibis_builder("_TestDecoratorExt") + def _build_test_ext(criteria, ctx): + return "ext_table" + + reg = get_registry() + assert reg.get_ibis_builder("_TestDecoratorExt") is _build_test_ext + + def test_decorator_registers_in_execution_registry(self): + @ibis_builder("_TestDecoratorExec") + def _build_test_exec(criteria, ctx): + return "exec_table" + + assert "_TestDecoratorExec" in _REGISTRY + assert _REGISTRY["_TestDecoratorExec"] is _build_test_exec + + def test_decorator_preserves_function(self): + @ibis_builder("_TestDecoratorPreserve") + def _build_preserve(criteria, ctx): + return "preserved" + + assert _build_preserve("c", "x") == "preserved" + + +# --------------------------------------------------------------------------- +# Programmatic registration +# --------------------------------------------------------------------------- + + +class TestProgrammaticRegistration: + """ExtensionRegistry.register_ibis_builder works at runtime.""" + + def test_register_and_retrieve(self): + reg = get_registry() + + def my_builder(criteria, ctx): + return "programmatic" + + reg.register_ibis_builder("_TestProgrammatic", my_builder) + assert reg.get_ibis_builder("_TestProgrammatic") is my_builder + + def test_unregistered_returns_none(self): + reg = get_registry() + assert reg.get_ibis_builder("_NeverRegistered") is None + + +# --------------------------------------------------------------------------- +# Fallback in get_builder +# --------------------------------------------------------------------------- + + +class TestGetBuilderFallback: + """get_builder falls back to the extension registry for unknown criteria.""" + + def test_fallback_to_extension_registry(self): + from circe.cohortdefinition.criteria import CriteriaGroup # noqa: F401 + + reg = get_registry() + + class _FallbackCriteria(Criteria): + pass + + _FallbackCriteria.model_rebuild() + + def _fb_builder(criteria, ctx): + return "fallback" + + reg.register_ibis_builder("_FallbackCriteria", _fb_builder) + + # Should be found via the extension registry fallback + found = get_builder(_FallbackCriteria()) + assert found is _fb_builder + + def test_unknown_criteria_raises(self): + from circe.cohortdefinition.criteria import CriteriaGroup # noqa: F401 + + class _UnknownCriteria(Criteria): + pass + + _UnknownCriteria.model_rebuild() + + with pytest.raises(ValueError, match="No builder registered"): + get_builder(_UnknownCriteria()) + + +# --------------------------------------------------------------------------- +# Waveform builders registered via @ibis_builder +# --------------------------------------------------------------------------- + + +class TestWaveformIbisBuilders: + """Waveform extension ibis builders are registered dynamically.""" + + @pytest.fixture(autouse=True) + def _import_waveform(self): + import circe.extensions.waveform # noqa: F401 + + @pytest.mark.parametrize( + "name", + [ + "WaveformOccurrence", + "WaveformRegistry", + "WaveformChannelMetadata", + "WaveformFeature", + ], + ) + def test_registered_in_extension_registry(self, name): + reg = get_registry() + builder = reg.get_ibis_builder(name) + assert builder is not None + assert callable(builder) + + @pytest.mark.parametrize( + "name", + [ + "WaveformOccurrence", + "WaveformRegistry", + "WaveformChannelMetadata", + "WaveformFeature", + ], + ) + def test_registered_in_execution_registry(self, name): + assert name in _REGISTRY + assert callable(_REGISTRY[name]) + + @pytest.mark.parametrize( + "criteria_cls_name", + [ + "WaveformOccurrence", + "WaveformRegistry", + "WaveformChannelMetadata", + "WaveformFeature", + ], + ) + def test_get_builder_resolves(self, criteria_cls_name): + from circe.extensions.waveform import criteria as wf_criteria + + cls = getattr(wf_criteria, criteria_cls_name) + builder = get_builder(cls()) + assert callable(builder)