Merge pull request #65 from CDCgov/rr-61-polarslazyframe

Rr 61 polarslazyframe

Author: ryanraaschCDC

cfa/dataops/catalog.py

@@ -6,12 +6,14 @@
 from configparser import ConfigParser
 from importlib import import_module
 from io import BytesIO
+from pathlib import PurePosixPath
 from types import SimpleNamespace
-from typing import Any, List, Sequence
+from typing import Any, List, Literal, Sequence, overload
 
 import pandas as pd
 import polars as pl
 import tomli
+from azure.identity import ManagedIdentityCredential
 
 from cfa.cloudops.blob_helpers import (
     read_blob_stream,
@@ -386,37 +388,111 @@ def download_version_to_local(
             self.ledger_entry(action="read")
         return written
 
+    @overload
     def get_dataframe(
         self,
-        output="pandas",
-        version="latest",
-        pl_lazy: bool = False,
+        output: Literal["pandas", "pd"] = "pandas",
+        version: str = "latest",
+        newest: bool = True,
+    ) -> pd.DataFrame: ...
+
+    @overload
+    def get_dataframe(
+        self,
+        output: Literal["polars", "pl"],
+        version: str = "latest",
+        newest: bool = True,
+    ) -> pl.DataFrame: ...
+
+    @overload
+    def get_dataframe(
+        self,
+        output: Literal["pl_lazy", "lazy"],
+        version: str = "latest",
         newest: bool = True,
-    ) -> pd.DataFrame | pl.DataFrame:
+    ) -> pl.LazyFrame: ...
+
+    def get_dataframe(
+        self,
+        output: Literal[
+            "pandas", "pd", "polars", "pl", "pl_lazy", "lazy"
+        ] = "pandas",
+        version: str = "latest",
+        newest: bool = True,
+    ) -> pd.DataFrame | pl.DataFrame | pl.LazyFrame:
         """Get the data as a pandas or polars dataframe
 
         Args:
             output (str, optional): the type of dataframe to return,
-                either 'pandas' or 'polars'. Defaults to "pandas".
+                either 'pandas' or 'polars' or 'pl_lazy'. Defaults to "pandas".
             version (str, optional): the version of the data to get.
                 Defaults to "latest".
-            pl_lazy (bool, optional): whether to return a lazy polars dataframe.
-                Defaults to False.
             newest (bool, optional): whether to get the newest matching version. Defaults to True.
                 False returns the oldest matching version.
 
         Raises:
-            ValueError: if output is not 'pandas' or 'polars'
+            ValueError: if output is not one of
+                'pandas', 'pd', 'polars', 'pl', 'pl_lazy', or 'lazy'
 
         Returns:
-            pd.DataFrame | pl.DataFrame: the dataframe
+            pd.DataFrame | pl.DataFrame | pl.LazyFrame: the dataframe
         """
-        if output not in ["pandas", "polars", "pd", "pl"]:
+        if output not in ["pandas", "polars", "pd", "pl", "pl_lazy", "lazy"]:
+            raise ValueError(
+                f"Output {output} needs to be 'pandas', 'polars', 'pd', 'pl', 'pl_lazy', or 'lazy'."
+            )
+        # Fetch version blobs once and validate before deriving file extension.
+        version_blobs = self._get_version_blobs(version=version, newest=newest)
+        if not version_blobs:
             raise ValueError(
-                f"Output {output} needs to be 'pandas', 'polars', 'pd, or 'pl'."
+                f"No blobs found for version '{version}' in container '{self.container}'."
             )
+        name = version_blobs[0]["name"]
+        file_ext = PurePosixPath(name).suffix.lstrip(".").lower()
+        if output in ["pl_lazy", "lazy"]:
+            if file_ext in ["parquet", "parq"]:
+                path = str(PurePosixPath(name).parent / f"*.{file_ext}")
+                fullpath = f"az://{self.container}/{path}"
+                df = pl.scan_parquet(
+                    fullpath,
+                    storage_options={"account_name": self.account},
+                    credential_provider=pl.CredentialProviderAzure(
+                        credential=ManagedIdentityCredential()
+                    ),
+                )
+                self.ledger_entry(action="read")
+                return df
+            elif file_ext == "csv":
+                path = str(PurePosixPath(name).parent / f"*.{file_ext}")
+                fullpath = f"az://{self.container}/{path}"
+                df = pl.scan_csv(
+                    fullpath,
+                    infer_schema_length=None,
+                    storage_options={"account_name": self.account},
+                    credential_provider=pl.CredentialProviderAzure(
+                        credential=ManagedIdentityCredential()
+                    ),
+                )
+                self.ledger_entry(action="read")
+                return df
+            elif file_ext == "ndjson" or file_ext == "jsonl":
+                path = str(PurePosixPath(name).parent / f"*.{file_ext}")
+                fullpath = f"az://{self.container}/{path}"
+                df = pl.scan_ndjson(
+                    fullpath,
+                    infer_schema_length=None,
+                    storage_options={"account_name": self.account},
+                    credential_provider=pl.CredentialProviderAzure(
+                        credential=ManagedIdentityCredential()
+                    ),
+                )
+                self.ledger_entry(action="read")
+                return df
+            else:
+                raise ValueError(
+                    f"Lazy loading not supported for {file_ext} files."
+                )
         blobs = self.read_blobs(version, newest=newest)
-        file_ext = self.get_file_ext(version=version)
         blob_bytes = [blob.content_as_bytes() for blob in blobs]
         blob_files = [BytesIO(pq) for pq in blob_bytes]
         if file_ext == "csv":
@@ -431,8 +507,6 @@ def get_dataframe(
                     ],
                     how="diagonal",
                 )
-                if pl_lazy:
-                    df = df.lazy()
             return df
         elif file_ext == "json":
             if output in ["pandas", "pd"]:
@@ -446,10 +520,8 @@ def get_dataframe(
                     ],
                     how="diagonal",
                 )
-                if pl_lazy:
-                    df = df.lazy()
             return df
-        elif file_ext == "jsonl":
+        elif file_ext == "jsonl" or file_ext == "ndjson":
             if output in ["pandas", "pd"]:
                 df = pd.concat(
                     [pd.read_json(blob, lines=True) for blob in blob_files]
@@ -460,8 +532,6 @@ def get_dataframe(
                     [pl.read_ndjson(blob) for blob in blob_files],
                     how="diagonal",
                 )
-                if pl_lazy:
-                    df = df.lazy()
             return df
         elif file_ext == "parquet" or file_ext == "parq":
             if output in ["pandas", "pd"]:
@@ -474,8 +544,6 @@ def get_dataframe(
                     [pl.read_parquet(pq_file) for pq_file in blob_files],
                     how="diagonal",
                 )
-                if pl_lazy:
-                    df = df.lazy()
             return df
 
     def ledger_entry(self, action: str) -> None:

changelog.md

@@ -7,6 +7,10 @@ and this project adheres to [Calendar Versioning](https://calver.org/).
 The versioning pattern is `YYYY.MM.DD.micro(a/b/{none if release})
 
 ---
+## [2026.03.31.0]
+
+- updated polars lazyframe loading to correctly reference the Blob file without downloading
+
 ## [2026.03.02.0]
 
 ### Added

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cfa.dataops"
-version = "2026.03.02.0"
+version = "2026.03.31.0"
 description = "Data cataloging, ETL, modeling, verification, and validation for CFA"
 authors = [
     { name = "Phil Rogers", email = "ap66@cdc.gov" },

tests/test_datasets_catalog.py

@@ -205,7 +205,73 @@ def mock_read_blob_stream_json_df(
     assert isinstance(blobs_df, pd.DataFrame)
     blobs_df = datacat.tests.etl_test.load.get_dataframe(output="pl")
     assert isinstance(blobs_df, pl.DataFrame)
-    blobs_df = datacat.tests.etl_test.load.get_dataframe(
-        output="pl", pl_lazy=True
+    blobs_df = datacat.tests.etl_test.load.get_dataframe(output="pl")
+    assert isinstance(blobs_df, pl.DataFrame)
+
+
+def test_datasets_catalog_get_dataframe_pl_lazy(
+    mocker, mock_write_blob_stream, dataset_ns_map, dataset_defaults
+):
+    datacat = dict_to_sn(dataset_ns_map, dataset_defaults)
+    dataset_namespaces = get_dataset_dot_path(dataset_ns_map)
+    datacat.__setattr__("__namespace_list__", dataset_namespaces)
+
+    mocker.patch(
+        "cfa.dataops.catalog.write_blob_stream",
+        mock_write_blob_stream,
+    )
+    mocker.patch.object(
+        datacat.tests.etl_test.load,
+        "get_versions",
+        return_value=["2025-06-03T17-56-50"],
+    )
+    mocker.patch.object(
+        datacat.tests.etl_test.load,
+        "_get_version_blobs",
+        return_value=[
+            {
+                "name": "prefix_test/transformed/test_dataset/2025-06-03T17-56-50/data.parquet",
+                "container": "container_test",
+            }
+        ],
+    )
+    mocker.patch.object(
+        datacat.tests.etl_test.load,
+        "ledger_entry",
+        return_value=None,
+    )
+
+    # Avoid real auth wiring during test
+    mocker.patch(
+        "cfa.dataops.catalog.ManagedIdentityCredential", return_value=object()
+    )
+    mocker.patch(
+        "cfa.dataops.catalog.pl.CredentialProviderAzure",
+        side_effect=lambda credential: object(),
+    )
+
+    scan_calls = []
+
+    def mock_scan_parquet(path, **kwargs):
+        scan_calls.append((path, kwargs))
+        return pl.DataFrame([{"test": 1, "data": 2}]).lazy()
+
+    mocker.patch(
+        "cfa.dataops.catalog.pl.scan_parquet",
+        side_effect=mock_scan_parquet,
+    )
+
+    out_pl_lazy = datacat.tests.etl_test.load.get_dataframe(output="pl_lazy")
+    out_lazy = datacat.tests.etl_test.load.get_dataframe(output="lazy")
+
+    assert isinstance(out_pl_lazy, pl.LazyFrame)
+    assert isinstance(out_lazy, pl.LazyFrame)
+    assert len(scan_calls) == 2
+    assert (
+        scan_calls[0][0]
+        == "az://container_test/prefix_test/transformed/test_dataset/2025-06-03T17-56-50/*.parquet"
+    )
+    assert (
+        scan_calls[1][0]
+        == "az://container_test/prefix_test/transformed/test_dataset/2025-06-03T17-56-50/*.parquet"
     )
-    assert isinstance(blobs_df, pl.LazyFrame)