BigQuery

Requires pip install "ehr2rl[bigquery]". See the BigQuery guide for prerequisites, costs, and limits.

Loading a dataset

ehr2rl.bigquery.load_mimiciv_bigquery_dataset(client, cohort, preset, itemid_map, action_config)

Load a bounded MIMIC-IV cohort from BigQuery into an EHRDataset.

Validates the itemid map against live labels, then queries admissions, the preset’s chartevents and labevents, and inputevents for medication actions. Every query goes through client’s budget guard.

Parameters:
  • client (GuardedBigQueryClient) – Guarded client that runs every query.

  • cohort (BigQueryCohort) – Cohort definition.

  • preset (FeaturePreset) – Concepts to use as state features.

  • itemid_map (ItemIdMap) – Itemid map for features and medications. All of its itemids are validated.

  • action_config (ActionConfig) – Bins and windows for medication actions.

Returns:

One trajectory per admission, with metadata died, feature_names, action_names, action_sizes, and provenance.

Return type:

EHRDataset

Raises:

Cohorts

class ehr2rl.bigquery.CohortCriteria(cohort_size=25, min_age=None, max_age=None, admission_types=(), min_icu_los_hours=None, max_icu_los_hours=None)

Filters that define a bounded ICU cohort.

Bounds are inclusive. None or an empty tuple means no filter.

Parameters:
  • cohort_size (int) – Maximum number of ICU stays. The first stays by subject_id, hadm_id, and stay_id are taken, not a random sample.

  • min_age (int | None) – Minimum MIMIC-IV anchor_age in years.

  • max_age (int | None) – Maximum MIMIC-IV anchor_age in years.

  • admission_types (tuple[str, ...]) – Allowed admission_type values, such as "EW EMER.".

  • min_icu_los_hours (float | None) – Minimum ICU length of stay in hours.

  • max_icu_los_hours (float | None) – Maximum ICU length of stay in hours.

Raises:

ValueError – If cohort_size is not positive or a bound is negative.

class ehr2rl.bigquery.BigQueryCohort(criteria, mimic_version='3_1')

Build bounded MIMIC-IV v3.1 SQL queries for a cohort.

Every query is restricted to the cohort’s ICU stays, and event queries keep only events during those stays.

Parameters:
  • criteria (CohortCriteria) – Cohort filters.

  • mimic_version (str) – Must be "3_1".

Raises:

ValueError – If mimic_version is not "3_1".

admissions_sql()

Return SQL for the cohort’s admissions rows.

Return type:

str

vitals_sql(itemids)

Return SQL for chartevents rows of the given itemids during cohort stays.

Parameters:

itemids (Iterable[int]) – Chartevents itemids to select. At least one is required.

Return type:

str

labs_sql(itemids)

Return SQL for labevents rows of the given itemids during cohort stays.

Parameters:

itemids (Iterable[int]) – Labevents itemids to select. At least one is required.

Return type:

str

inputevents_sql(itemids)

Return SQL for inputevents rows of the given itemids overlapping cohort stays.

Parameters:

itemids (Iterable[int]) – Inputevents itemids to select. At least one is required.

Return type:

str

Guarded execution

class ehr2rl.bigquery.GuardedBigQueryClient(client, maximum_bytes_billed, cache_dir=None, query_timeout=None, disable_retries=False)

Run BigQuery queries only after a dry-run budget check.

Parameters:
  • client (Any) – A google.cloud.bigquery.Client.

  • maximum_bytes_billed (int) – Per-query byte cap. A query whose dry run exceeds it is not run, and the cap is also set on each job so BigQuery enforces it server-side.

  • cache_dir (str | Path | None) – Directory for cached results, created if missing. Results are not cached when None.

  • query_timeout (float | None) – Seconds to wait for each API call and result download. Must be positive.

  • disable_retries (bool) – Turn off the Google client’s automatic retries.

Raises:

ValueError – If maximum_bytes_billed or query_timeout is not positive.

query_dataframe(sql, *, cache_key_parts=())

Run sql after cache and dry-run checks and return its results.

Parameters:
  • sql (str) – Standard SQL query.

  • cache_key_parts (tuple[str, ...]) – Extra strings mixed into the cache key, such as a map version.

Returns:

The results and execution metadata.

Return type:

QueryResult

Raises:
  • BudgetExceededError – If the dry-run estimate exceeds maximum_bytes_billed.

  • BigQueryAuthError – If the dry run or query fails for any reason, including credentials, permissions, invalid SQL, and timeouts.

class ehr2rl.bigquery.QueryResult(dataframe, job_id, bytes_processed, query_hash)

A query’s results plus BigQuery execution metadata.

Parameters:
  • dataframe (DataFrame) – Query results.

  • job_id (str | None) – BigQuery job ID, or None when the result came from the local cache.

  • bytes_processed (int) – Dry-run estimate of bytes processed; 0 for cached results.

  • query_hash (str) – SHA-256 cache key of the SQL and cache key parts.

Errors

exception ehr2rl.bigquery.BigQueryError

Base error for ehr2rl BigQuery operations.

exception ehr2rl.bigquery.BudgetExceededError

Raised before execution when a dry run exceeds the configured budget.

exception ehr2rl.bigquery.BigQueryAuthError

Raised when a BigQuery dry run or query fails.

Despite the name, this covers every failure from the BigQuery client, including invalid SQL, timeouts, and network errors, not only authentication and permissions.