Polars Lazyframe Support

hamilton/function_modifiers/base.py

@@ -26,6 +26,7 @@
         "pandas",
         "plotly",
         "polars",
+        "polars_lazyframe",
         "pyspark_pandas",
         "spark",
         "dask",

hamilton/io/utils.py

@@ -62,14 +62,27 @@ def get_dataframe_metadata(df: pd.DataFrame) -> Dict[str, Any]:
     - the column names
     - the data types
     """
-    return {
-        DATAFRAME_METADATA: {
-            "rows": len(df),
-            "columns": len(df.columns),
-            "column_names": list(df.columns),
-            "datatypes": [str(t) for t in list(df.dtypes)],  # for serialization purposes
-        }
-    }
+    metadata = {}
+    try:
+        metadata["rows"] = len(df)
+    except TypeError:
+        metadata["rows"] = None
+
+    try:
+        metadata["columns"] = len(df.columns)
+    except (AttributeError, TypeError):
+        metadata["columns"] = None
+
+    try:
+        metadata["column_names"] = list(df.columns)
+    except (AttributeError, TypeError):
+        metadata["column_names"] = None
+
+    try:
+        metadata["datatypes"] = [str(t) for t in list(df.dtypes)]
+    except (AttributeError, TypeError):
+        metadata["datatypes"] = None
+    return {DATAFRAME_METADATA: metadata}
 
 
 def get_file_and_dataframe_metadata(path: str, df: pd.DataFrame) -> Dict[str, Any]:

hamilton/plugins/h_polars.py

@@ -40,6 +40,8 @@ def build_result(
             (value,) = outputs.values()  # this works because it's length 1.
             if isinstance(value, pl.DataFrame):  # it's a dataframe
                 return value
+            if isinstance(value, pl.LazyFrame):  # it's a lazyframe
+                return value.collect()
             elif not isinstance(value, pl.Series):  # it's a single scalar/object
                 key, value = outputs.popitem()
                 return pl.DataFrame({key: [value]})

hamilton/plugins/h_polars_lazyframe.py

@@ -0,0 +1,46 @@
+from typing import Any, Dict, Type, Union
+
+import polars as pl
+
+from hamilton import base
+
+
+class PolarsLazyFrameResult(base.ResultMixin):
+    """A ResultBuilder that produces a polars dataframe.
+
+    Use this when you want to create a polars dataframe from the outputs. Caveat: you need to ensure that the length
+    of the outputs is the same, otherwise you will get an error; mixed outputs aren't that well handled.
+
+    To use:
+
+    .. code-block:: python
+
+        from hamilton import base, driver
+        from hamilton.plugins import polars_extensions
+        polars_builder = polars_extensions.PolarsDataFrameResult()
+        adapter = base.SimplePythonGraphAdapter(polars_builder)
+        dr =  driver.Driver(config, *modules, adapter=adapter)
+        df = dr.execute([...], inputs=...)  # returns polars dataframe
+
+    Note: this is just a first attempt at something for Polars. Think it should handle more? Come chat/open a PR!
+    """
+
+    def build_result(
+        self, **outputs: Dict[str, Union[pl.Series, pl.LazyFrame, Any]]
+    ) -> pl.LazyFrame:
+        """This is the method that Hamilton will call to build the final result. It will pass in the results
+        of the requested outputs that you passed in to the execute() method.
+
+        Note: this function could do smarter things; looking for contributions here!
+
+        :param outputs: The results of the requested outputs.
+        :return: a polars DataFrame.
+        """
+        if len(outputs) == 1:
+            (value,) = outputs.values()  # this works because it's length 1.
+            if isinstance(value, pl.LazyFrame):  # it's a lazyframe
+                return value
+        return pl.LazyFrame(outputs)
+
+    def output_type(self) -> Type:
+        return pl.LazyFrame

hamilton/plugins/polars_extensions.py

@@ -1,4 +1,5 @@
 import dataclasses
+import logging
 from io import BytesIO, IOBase, TextIOWrapper
 from pathlib import Path
 from typing import (
@@ -69,6 +70,8 @@ def fill_with_scalar_polars(df: pl.DataFrame, column_name: str, scalar_value: An
 
 
 def register_types():
+    logger = logging.getLogger(__name__)
+    logger.info("registering lazyframe")
     """Function to register the types for this extension."""
     registry.register_types("polars", DATAFRAME_TYPE, COLUMN_TYPE)
 
@@ -166,8 +169,6 @@ def _get_loading_kwargs(self):
             kwargs["row_count_name"] = self.row_count_name
         if self.row_count_offset is not None:
             kwargs["row_count_offset"] = self.row_count_offset
-        if self.sample_size is not None:
-            kwargs["sample_size"] = self.sample_size
         if self.eol_char is not None:
             kwargs["eol_char"] = self.eol_char
         if self.raise_if_empty is not None:
@@ -176,6 +177,7 @@ def _get_loading_kwargs(self):
 
     def load_data(self, type_: Type) -> Tuple[DATAFRAME_TYPE, Dict[str, Any]]:
         df = pl.read_csv(self.file, **self._get_loading_kwargs())
+
         metadata = utils.get_file_and_dataframe_metadata(self.file, df)
         return df, metadata
 
@@ -206,7 +208,7 @@ class PolarsCSVWriter(DataSaver):
 
     @classmethod
     def applicable_types(cls) -> Collection[Type]:
-        return [DATAFRAME_TYPE]
+        return [DATAFRAME_TYPE, pl.LazyFrame]
 
     def _get_saving_kwargs(self):
         kwargs = {}
@@ -236,15 +238,12 @@ def _get_saving_kwargs(self):
             kwargs["quote_style"] = self.quote_style
         return kwargs
 
-    def save_data(self, data: DATAFRAME_TYPE) -> Dict[str, Any]:
+    def save_data(self, data: Union[DATAFRAME_TYPE, pl.LazyFrame]) -> Dict[str, Any]:
+        if isinstance(data, pl.LazyFrame):
+            data = data.collect()
         data.write_csv(self.file, **self._get_saving_kwargs())
         return utils.get_file_and_dataframe_metadata(self.file, data)
 
-    def load_data(self, type_: Type) -> Tuple[DATAFRAME_TYPE, Dict[str, Any]]:
-        df = pl.read_csv(self.file, **self._get_loading_kwargs())
-        metadata = utils.get_file_and_dataframe_metadata(self.file, df)
-        return df, metadata
-
     @classmethod
     def name(cls) -> str:
         return "csv"
@@ -330,7 +329,7 @@ class PolarsParquetWriter(DataSaver):
 
     @classmethod
     def applicable_types(cls) -> Collection[Type]:
-        return [DATAFRAME_TYPE]
+        return [DATAFRAME_TYPE, pl.LazyFrame]
 
     def _get_saving_kwargs(self):
         kwargs = {}
@@ -348,8 +347,12 @@ def _get_saving_kwargs(self):
             kwargs["pyarrow_options"] = self.pyarrow_options
         return kwargs
 
-    def save_data(self, data: DATAFRAME_TYPE) -> Dict[str, Any]:
+    def save_data(self, data: Union[DATAFRAME_TYPE, pl.LazyFrame]) -> Dict[str, Any]:
+        if isinstance(data, pl.LazyFrame):
+            data = data.collect()
+
         data.write_parquet(self.file, **self._get_saving_kwargs())
+
         return utils.get_file_and_dataframe_metadata(self.file, data)
 
     @classmethod
@@ -422,15 +425,17 @@ class PolarsFeatherWriter(DataSaver):
 
     @classmethod
     def applicable_types(cls) -> Collection[Type]:
-        return [DATAFRAME_TYPE]
+        return [DATAFRAME_TYPE, pl.LazyFrame]
 
     def _get_saving_kwargs(self):
         kwargs = {}
         if self.compression is not None:
             kwargs["compression"] = self.compression
         return kwargs
 
-    def save_data(self, data: DATAFRAME_TYPE) -> Dict[str, Any]:
+    def save_data(self, data: Union[DATAFRAME_TYPE, pl.LazyFrame]) -> Dict[str, Any]:
+        if isinstance(data, pl.LazyFrame):
+            data = data.collect()
         data.write_ipc(self.file, **self._get_saving_kwargs())
         return utils.get_file_and_dataframe_metadata(self.file, data)