diff --git a/cfa/stf/data/get_data.py b/cfa/stf/data/get_data.py index ad8b56000..00464bc7c 100644 --- a/cfa/stf/data/get_data.py +++ b/cfa/stf/data/get_data.py @@ -1,5 +1,5 @@ import datetime as dt -from typing import Literal +from typing import Literal, overload import polars as pl from cfa.dataops import datacat @@ -13,6 +13,30 @@ NSSPDataset = Literal["gold", "comprehensive"] +@overload +def get_nhsn_hrd( + disease: str, + loc_abb: str, + prelim: bool = ..., + as_of: dt.date | None = ..., + start_date: dt.date | None = ..., + end_date: dt.date | None = ..., + lazy: Literal[True] = ..., +) -> pl.LazyFrame: ... + + +@overload +def get_nhsn_hrd( + disease: str, + loc_abb: str, + prelim: bool = ..., + as_of: dt.date | None = ..., + start_date: dt.date | None = ..., + end_date: dt.date | None = ..., + lazy: Literal[False] = ..., +) -> pl.DataFrame: ... + + def get_nhsn_hrd( disease: str, loc_abb: str, @@ -20,7 +44,8 @@ def get_nhsn_hrd( as_of: dt.date | None = None, start_date: dt.date | None = None, end_date: dt.date | None = None, -) -> pl.DataFrame: + lazy: bool = True, +) -> pl.DataFrame | pl.LazyFrame: """ Retrieve and filter NHSN hospital respiratory data based on specified criteria. @@ -42,10 +67,13 @@ def get_nhsn_hrd( The start date for the time period to include. If None, no lower bound is applied. end_date The end date for the time period to include. If None, no upper bound is applied. + lazy + Whether to return a lazy frame (defaults to True). If True, returns a + `pl.LazyFrame`; if False, returns a `pl.DataFrame`. Returns ------- - pl.DataFrame + pl.DataFrame | pl.LazyFrame Filtered data with columns: `weekendingdate`, `jurisdiction`, and `hospital_admissions`. """ @@ -67,8 +95,9 @@ def get_nhsn_hrd( datacat.public.stf.nhsn_hrd_prelim if prelim else datacat.public.stf.nhsn_hrd ) + output = "pl_lazy" if lazy else "pl" dat = datacat_dataset.load.get_dataframe( - output="pl", version=f"<={as_of.strftime('%Y-%m-%d')}" + output=output, version=f"<={as_of.strftime('%Y-%m-%d')}" ) filtered_dat = ( @@ -83,6 +112,30 @@ def get_nhsn_hrd( return filtered_dat +@overload +def get_nssp( + disease: str, + loc_abb: str, + dataset: NSSPDataset = ..., + as_of: dt.date | None = ..., + start_date: dt.date | None = ..., + end_date: dt.date | None = ..., + lazy: Literal[True] = ..., +) -> pl.LazyFrame: ... + + +@overload +def get_nssp( + disease: str, + loc_abb: str, + dataset: NSSPDataset = ..., + as_of: dt.date | None = ..., + start_date: dt.date | None = ..., + end_date: dt.date | None = ..., + lazy: Literal[False] = ..., +) -> pl.DataFrame: ... + + def get_nssp( disease: str, loc_abb: str, @@ -90,7 +143,8 @@ def get_nssp( as_of: dt.date | None = None, start_date: dt.date | None = None, end_date: dt.date | None = None, -) -> pl.DataFrame: + lazy: bool = True, +) -> pl.DataFrame | pl.LazyFrame: """ Retrieve and filter NSSP emergency department data. @@ -116,10 +170,13 @@ def get_nssp( Start date for filtering data (inclusive). If None, no lower bound is applied (defaults to None). end_date End date for filtering data (inclusive). If None, no upper bound is applied (defaults to None). + lazy + Whether to return a lazy frame (defaults to True). If True, returns a + `pl.LazyFrame`; if False, returns a `pl.DataFrame`. Returns ------- - pl.DataFrame + pl.DataFrame | pl.LazyFrame Aggregated ED counts with columns: `reference_date`, `disease`, and `value`. @@ -162,11 +219,16 @@ def get_nssp( # we only filter by geo_value if loc_abb is not "US". filters.append(pl.col("geo_value") == loc_abb) + output = "pl_lazy" if lazy else "pl" dat = datacat_dataset.load.get_dataframe( - output="pl", version=f"<={as_of.strftime('%Y-%m-%d')}" + output=output, version=f"<={as_of.strftime('%Y-%m-%d')}" ) - valid_locs = dat.unique("geo_value").get_column("geo_value").to_list() + ["US"] + valid_locs = ( + dat.select("geo_value").unique().collect() + if lazy + else dat.select("geo_value").unique() + ).get_column("geo_value").to_list() + ["US"] if loc_abb not in valid_locs: raise ValueError(f"Invalid location abbreviation: {loc_abb}") diff --git a/cfa/stf/data/get_nnh_pmfs.py b/cfa/stf/data/get_nnh_pmfs.py index 1846ddb6f..675c58496 100644 --- a/cfa/stf/data/get_nnh_pmfs.py +++ b/cfa/stf/data/get_nnh_pmfs.py @@ -1,32 +1,56 @@ import datetime as dt +from typing import Literal, overload import polars as pl from cfa.dataops import datacat def _extract_pmf( - df: pl.DataFrame, + df: pl.DataFrame | pl.LazyFrame, parameter_name: str, -) -> list: - df = df.filter(pl.col("parameter") == parameter_name) - if df.height != 1: +) -> list[float]: + pmf_df = df.filter(pl.col("parameter") == parameter_name) + + if isinstance(pmf_df, pl.LazyFrame): + pmf_df = pmf_df.collect() + + if pmf_df.height != 1: raise ValueError( f"Expected exactly one {parameter_name!r} parameter row, " - f"but found {df.height}. " - f"Rows={df.to_dicts()}" + f"but found {pmf_df.height}. " + f"Rows={pmf_df.to_dicts()}" ) - return df.item(0, "value").to_list() + + return pmf_df.item(0, "value").to_list() + + +@overload +def _filter_param_estimates( + disease: str, + as_of: dt.date | None = ..., + lazy: Literal[True] = ..., +) -> pl.LazyFrame: ... + + +@overload +def _filter_param_estimates( + disease: str, + as_of: dt.date | None = ..., + lazy: Literal[False] = ..., +) -> pl.DataFrame: ... def _filter_param_estimates( disease: str, as_of: dt.date | None = None, -) -> pl.DataFrame: + lazy: bool = True, +) -> pl.DataFrame | pl.LazyFrame: as_of = as_of or dt.date.max - dt.timedelta(days=1) - dat = datacat.public.stf.param_estimates.load.get_dataframe(output="pl") - return ( - dat.with_columns( + output = "pl_lazy" if lazy else "pl" + result = ( + datacat.public.stf.param_estimates.load.get_dataframe(output=output) + .with_columns( pl.col("start_date").fill_null(dt.date.min), pl.col("end_date").fill_null(dt.date.max), ) @@ -36,6 +60,7 @@ def _filter_param_estimates( as_of < pl.col("end_date"), ) ) + return result def get_nnh_generation_interval_pmf( @@ -71,10 +96,7 @@ def get_nnh_generation_interval_pmf( return _extract_pmf(dat_filtered, "generation_interval") -def get_nnh_delay_pmf( - disease: str, - as_of: dt.date | None = None, -) -> list[float]: +def get_nnh_delay_pmf(disease: str, as_of: dt.date | None = None) -> list[float]: """ Filter and extract the delay probability mass function (PMF) based on disease and date filters.