Add lazy option to cfa.stf.data functions

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,14 +13,39 @@
 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,
     prelim: bool = True,
     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,14 +112,39 @@ 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,
     dataset: NSSPDataset = "gold",
     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}")

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.