Local Data Source (skeleton)

Local Data Source (skeleton)#

local_datasource is the source skeleton that ships with CORR-Vars. It carries the variable classes and the extraction pipeline, but no connection to any store: the query that reaches a warehouse table, a folder of parquet files, or an HTTP service is deployment-specific and is not part of the published package.

Two things follow from that:

The aggregation and derivation classes (NativeStatic, DerivedStatic, DerivedDynamic, ComplexVariable) compute from data that is already in the cohort. They need no store and work on any cohort — see Aggregated Variables and Custom Variables Guide.
Only NativeDynamic, which fetches new rows, needs a store. That is what binding the skeleton is for.

The source is pinned as local under [local].sources in src/corr_vars/concepts/routing.toml: its definitions are never fetched from the Concepts API, so it works with no network and no API key. Definitions come from the bundled mapping/vars.json, and py functions from mapping/variables.py, looked up by variable name.

The two seams#

Binding the skeleton to a store means replacing exactly two things.

1. `cohort.load_data` — the observation frame#

load_data() raises NotImplementedError. Replace it to build the cohort’s observation frame: one row per observation, carrying at least the level’s primary key and its t_min / t_max anchor columns (for icu_stay: icu_stay_id, icu_admission, icu_discharge).

2. `NativeExtractor.extract` — the raw rows#

NativeExtractor.extract() also raises NotImplementedError. Subclass it, implement extract(cohort, id_column) to return rows carrying at least id_column, recordtime and value, and point NativeDynamic.extractor_class at the subclass. NativeDynamic wraps the extractor in the BaseDynamic pipeline: cache → extract_from_db → time filter → cleaning → optional py post-processing.

Worked example: a folder of parquet files#

The following binds the skeleton to a directory in which each table is one parquet file, keyed on icu_stay_id:

from pathlib import Path

import polars as pl

import corr_vars.sources.local_datasource.cohort as ld_cohort
from corr_vars.sources.local_datasource import NativeDynamic, NativeExtractor

DATA_DIR = Path("/path/to/my_store")


class MyExtractor(NativeExtractor):
    """Read one parquet file per table and restrict it to the cohort's ids."""

    def extract(self, cohort, id_column):
        frame = pl.scan_parquet(DATA_DIR / f"{self.table}.parquet")

        # `where` is optional; when set it is an SQL-style boolean on the table.
        if self.where:
            frame = frame.filter(pl.sql_expr(self.where))

        ids = cohort.obs.select(id_column).unique()
        frame = frame.join(ids.lazy(), on=id_column, how="inner")

        frame = frame.select(id_column, "recordtime", "value")
        if self.value_dtype:
            frame = frame.with_columns(
                pl.col("value").cast(getattr(pl, self.value_dtype))
            )
        return frame.collect()


NativeDynamic.extractor_class = MyExtractor


def load_data(obs_level, source_kwargs):
    """Build the observation frame for `obs_level`."""
    stays = pl.read_parquet(DATA_DIR / "icu_stays.parquet")
    return stays.select("icu_stay_id", "icu_admission", "icu_discharge")


ld_cohort.load_data = load_data

With both seams filled, a NativeDynamic variable extracts from the store like any other native variable:

from corr_vars import Cohort

cohort = Cohort(
    obs_level="icu_stay",
    sources={"local_datasource": {}},
    project="my_project",
)

cohort.add_variable(NativeDynamic(
    var_name="blood_lactate",
    dynamic=True,
    table="labs",
    where="analyte = 'lactate'",
    value_dtype="Float64",
    cleaning={"value": {"low": 0.1, "high": 30.0}},
))

Configuration#

corr_vars.sources.local_datasource.config defines the SourceDict and SourceDictPartial TypedDicts, plus DEFAULTS and MIGRATIONS. The published schema is deliberately small — a database name and conn_args (hostname, username, password, password_file):

DEFAULTS: SourceDict = {
    "database": None,
    "conn_args": {
        "hostname": None,
        "username": None,
        "password": None,
        "password_file": False,
    },
}

A deployment extends these with whatever its store needs. MIGRATIONS maps legacy top-level keys to their canonical nested position, so old configuration files keep loading.

Variable definitions#

Once the store is bound, the source needs definitions. Either fill the bundled mapping, or route the source to a Concepts API endpoint in concepts/routing.toml (which unpins it from [local].sources).

`mapping/vars.json`#

The bundled file declares no variables. Add them under variables, and list the ones that should load automatically under corr_defaults.default_vars per observation level. The main definition fields are:

Field	Meaning
`type`	`native_dynamic`, `native_static`, `derived_static`, `derived_dynamic`, `complex` or `free_days`. Selects the class the loader builds.
`table`	Table in the backing store (native variables only).
`where`	Row filter for the extraction, in the store’s dialect (native variables only).
`value_dtype`	Polars dtype name the `value` column is cast to, e.g. `"Float64"`.
`cleaning`	Per-column plausibility bounds, e.g. `{"value": {"low": 0, "high": 300}}`.
`requires` / `base_var`	Dependencies. Derived and complex variables take a `requires` list; `native_static` takes a single `base_var`.
`compatible_with`	Observation levels the definition is valid for.

`mapping/variables.py`#

One module-level function per variable that needs a py transformation, named exactly like its key in vars.json. The loader looks the function up by that name and injects it into the config as "py". A variable with no transformation needs no entry.

Note

A variable published through the CORR Concepts API carries its py source with it, so variables.py is only for the definitions a source bundles itself.

Class Reference#

class corr_vars.sources.local_datasource.extract.BaseDynamic(var_name, dynamic, requires=[], tmin=None, tmax=None, py=None, py_ready_polars=False, cleaning=None, allow_caching=True, screened_obs_level=None)[source]#

Bases: _ScreenedRequirements, Variable, ABC

Abstract base for all dynamic variables of this source.

Provides the shared extraction pipeline: optimised obs-level selection, caching, attribute building, time-filtering, and cleaning. Subclasses implement _extract_from_db() to supply the raw data.

Parameters:

var_name (str) – Variable name.
dynamic (bool) – Whether the variable is dynamic (time-series).
requires (list[str] | dict[str, RequirementsDict]) – Dependency variables loaded before extraction.
tmin (str | tuple[str, str] | TimeAnchor | None) – Minimum time anchor.
tmax (str | tuple[str, str] | TimeAnchor | None) – Maximum time anchor.
py (VariableCallable | None) – Optional transformation function applied after raw extraction.
py_ready_polars (bool) – Whether py accepts and returns polars DataFrames.
cleaning (dict[str, dict[Literal['low', 'high'], Any]] | None) – Value-range cleaning rules.
allow_caching (bool) – Whether to cache raw data in the cohort’s tmpdir.
screened_obs_level (Optional[Literal['patient', 'hospital_stay']]) – Default observation level used to scope DB queries when the time window is not bounded by hospital-stay bounds.

screened_obs_level: Literal[ObsLevel.HOSPITAL_STAY, ObsLevel.PATIENT] | None#

abstractmethod extract_from_db(cohort, id_column)[source]#

Return raw variable data; called by extract() on a cache miss.

The returned DataFrame must not yet have time-window columns or cleaning applied — those are handled by the shared pipeline.

Parameters:

cohort (Cohort)
id_column (Literal['case_id', 'patient_id'])

Return type:

DataFrame

extract(cohort)[source]#

Extracts data from the datasource. Usually follows this pattern.

# Load & Extract required variables
self._get_required_vars(cohort)

# This should change self.data either by returning or side effect
self.data = self._custom_extraction(cohort)

# Calls variable function to transform extracted data
self._call_var_function(cohort)

if self.dynamic:
    self._add_time_window(cohort)
    # Expects case_tmin, case_tmax for each primary key
    self._timefilter()
    self._apply_cleaning()

Parameters:: cohort (Cohort)
Return type:: DataFrame

class corr_vars.sources.local_datasource.extract.NativeExtractor(var_name, *, time_window=None, table=None, where=None, columns=None, value_dtype=None)[source]#

Bases: object

Raw extractor satisfying VariableProtocol.

Fetches a variable’s rows from the source’s backing store. No caching, time-filtering, or cleaning — those are handled by the wrapping NativeDynamic pipeline, which is what makes this class worth separating: an extractor can be swapped without touching that pipeline.

extract() is deliberately not implemented here. A deployment binds this source to its own store — a warehouse table, a set of parquet files, an HTTP service — and the query that reaches it is deployment-specific, so it is not part of the published package. Subclass this and implement extract(), then point NativeDynamic at the subclass through NativeDynamic.extractor_class.

Parameters:

var_name (str) – Variable name.
table (str | None) – Table name in the backing store.
where (str | None) – Row filter, in whatever dialect the store speaks.
columns (dict[str, str | dict[str, str]] | None) – Column alias/cast mapping for the query.
value_dtype (str | None) – Polars dtype name the value column is cast to.
time_window (TimeWindow | None) – Time window (stored to satisfy the protocol).

value_dtype: str | None#

extract(cohort, id_column)[source]#

Fetch this variable’s raw rows, restricted to the cohort’s ids.

Parameters:

cohort (Cohort) – Cohort whose ids the extract is restricted to.
id_column (Literal["case_id", "patient_id"]) – Key the store is queried and joined on.

Returns:

One row per record, carrying at least id_column, recordtime and value.

Return type:

pl.DataFrame

Raises:

NotImplementedError – Always. See the class docstring — the query that reaches a deployment’s store is not part of this package.

class corr_vars.sources.local_datasource.extract.NativeDynamic(var_name, dynamic, requires=[], tmin=None, tmax=None, py=None, py_ready_polars=False, cleaning=None, allow_caching=True, screened_obs_level=None, table=None, where=None, value_dtype=None, columns=None)[source]#

Bases: BaseDynamic

Store-backed dynamic variable with the full BaseDynamic pipeline.

Wraps a NativeExtractor to provide caching, time-filtering, cleaning, and optional py post-processing.

The extractor is the only part that touches a backing store, and the published NativeExtractor raises NotImplementedError. Bind this source to a store by subclassing it and setting extractor_class:

class MyExtractor(NativeExtractor):
    def extract(self, cohort, id_column):
        ...  # query your store, return one row per record

NativeDynamic.extractor_class = MyExtractor

Parameters:

table (str | None) – Table name in the backing store.
where (str | None) – Row filter appended to the generated query.
value_dtype (str | None) – Optional polars dtype name applied to the value column after extraction (e.g. "Float64").
columns (dict[str, str | dict[str, str]] | None) – Column alias/cast mapping for the query.

All other args are forwarded to BaseDynamic.

Note

Only use this class for defining custom variables. For CORR-specified variables, use cohort.add_variable() directly.

Parameters:

var_name (str)
dynamic (bool)
requires (list[str] | dict[str, RequirementsDict])
tmin (str | tuple[str, str] | TimeAnchor | None)
tmax (str | tuple[str, str] | TimeAnchor | None)
py (VariableCallable | None)
py_ready_polars (bool)
cleaning (dict[str, dict[Literal['low', 'high'], Any]] | None)
allow_caching (bool)
screened_obs_level (Optional[Literal['patient', 'hospital_stay']])

extractor_class#

Extractor the variable builds to reach the backing store. Rebind it on this class (or on a subclass) to plug in a deployment’s own query.

Parameters:

var_name (str)
time_window (TimeWindow | None)
table (str | None)
where (str | None)
columns (dict[str, str | dict[str, str]] | None)
value_dtype (str | None)

alias of NativeExtractor

extract_from_db(cohort, id_column)[source]#

Return raw variable data; called by extract() on a cache miss.

The returned DataFrame must not yet have time-window columns or cleaning applied — those are handled by the shared pipeline.

Parameters:

cohort (Cohort)
id_column (Literal['case_id', 'patient_id'])

Return type:

DataFrame

class corr_vars.sources.local_datasource.extract.ComplexVariable(var_name, dynamic, requires=[], tmin=None, tmax=None, py=None, py_ready_polars=False, cleaning=None, allow_caching=True, screened_obs_level=None)[source]#

Bases: BaseDynamic

Variable extracted entirely by a custom py() function.

No database query is issued; the variable function receives required_vars and must return a pl.DataFrame. Time-filtering is applied only when the function adds case_tmin/case_tmax columns to the result.

Parameters:

var_name (str) – Name of the variable.
dynamic (bool) – Whether the variable is dynamic.
requires (list[str] | dict[str, RequirementsDict]) – Variables loaded and passed to the function.
tmin (str | tuple[str, str] | TimeAnchor | None) – Minimum time anchor.
tmax (str | tuple[str, str] | TimeAnchor | None) – Maximum time anchor.
py (VariableCallable | None) – Variable function that produces the data.
py_ready_polars (bool) – Whether py accepts and returns polars DataFrames.
cleaning (dict[str, dict[Literal['low', 'high'], Any]] | None) – Dictionary of cleaning rules.
allow_caching (bool) – Whether to allow caching.
screened_obs_level (Optional[Literal['patient', 'hospital_stay']]) – Observation level used for scoping the cache key.

extract_from_db(cohort, id_column)[source]#

Return raw variable data; called by extract() on a cache miss.

The returned DataFrame must not yet have time-window columns or cleaning applied — those are handled by the shared pipeline.

Parameters:

cohort (Cohort)
id_column (Literal['case_id', 'patient_id'])

Return type:

DataFrame

extract(cohort)[source]#

Extracts data from the datasource. Usually follows this pattern.

# Load & Extract required variables
self._get_required_vars(cohort)

# This should change self.data either by returning or side effect
self.data = self._custom_extraction(cohort)

# Calls variable function to transform extracted data
self._call_var_function(cohort)

if self.dynamic:
    self._add_time_window(cohort)
    # Expects case_tmin, case_tmax for each primary key
    self._timefilter()
    self._apply_cleaning()

Parameters:: cohort (Cohort)
Return type:: DataFrame

corr_vars.sources.local_datasource.extract.VariableLoader(*args, **kwargs)[source]#

Factory function to create the correct variable type based on the arguments.

Return type:: Variable

corr_vars.sources.local_datasource.cohort.load_data(obs_level, source_kwargs)[source]#

Build the cohort’s observation frame from the backing store.

Parameters:

obs_level (ObsLevel) – Level of observation the frame is keyed at.
source_kwargs (SourceDict) – This source’s resolved configuration.

Returns:

One row per observation, carrying at least the level’s primary key and its t_min / t_max anchor columns.

Return type:

pl.DataFrame

Raises:

NotImplementedError – Always. local_datasource is a source skeleton; replace this function to bind it to a store.

Local Data Source (skeleton)

Contents

Local Data Source (skeleton)#

The two seams#

1. cohort.load_data — the observation frame#

2. NativeExtractor.extract — the raw rows#

Worked example: a folder of parquet files#

Configuration#

Variable definitions#

mapping/vars.json#

mapping/variables.py#

Class Reference#

1. `cohort.load_data` — the observation frame#

2. `NativeExtractor.extract` — the raw rows#

`mapping/vars.json`#

`mapping/variables.py`#