feat: Add Polars eager DataFrame support (#82)

Author: OlegWock

deepnote_toolkit/chart/deepnote_chart.py

@@ -28,7 +28,7 @@
 def _create_vf_runtime_for_dataframe(
     oc_df: oc.DataFrame, name: str
 ) -> vegafusion.VegaFusionRuntime:
-    if oc_df.native_type == "pandas":
+    if oc_df.native_type in ("pandas", "polars-eager"):
         return vegafusion.VegaFusionRuntime()
     elif oc_df.native_type == "pyspark":
         spark_df = oc_df.to_native()
@@ -55,7 +55,7 @@ def spark_executor(sql_query: str) -> pa.Table:
 
 
 def _create_vf_inline_dataset_from_dataframe(oc_df: oc.DataFrame) -> Any:
-    if oc_df.native_type == "pandas":
+    if oc_df.native_type in ("pandas", "polars-eager"):
         return oc_df.to_native()
     elif oc_df.native_type == "pyspark":
         from pyspark.sql.pandas.types import to_arrow_schema
@@ -141,7 +141,11 @@ def __init__(
             if filtered_df.native_type == "pandas":
                 sanitized_pandas = sanitize_dataframe_for_chart(filtered_df.to_native())
                 oc_sanitized_df = oc.DataFrame.from_native(sanitized_pandas)
-            elif filtered_df.native_type == "pyspark":
+            elif filtered_df.native_type in ("pyspark", "polars-eager"):
+                # We don't need to sanitize Spark DFs because they will processed by Spark itself and it can handle
+                # all data types by itself
+                # Polars is powered by Arrow, which is same format used internally by VegaFusion so there is no need
+                # to do any additional sanitization for it either
                 oc_sanitized_df = filtered_df
             else:
                 raise TypeError(

deepnote_toolkit/dataframe_utils.py

@@ -102,6 +102,9 @@ def dataframe_formatter(native_df: oc.NativeInputDF):
     mimebundle_formatter.for_type_by_name(
         "pyspark.pandas.frame", "DataFrame", dataframe_formatter
     )
+    mimebundle_formatter.for_type_by_name(
+        "polars.dataframe.frame", "DataFrame", dataframe_formatter
+    )
     logger.info("Attached mimebundle formatters")
 
 

deepnote_toolkit/ocelots/__init__.py

@@ -7,6 +7,7 @@
 from .dataframe import (
     DataFrame,
     is_wrapped_pandas_dataframe,
+    is_wrapped_polars_eager_dataframe,
     is_wrapped_pyspark_dataframe,
 )
 from .filters import Filter, FilterOperator
@@ -19,11 +20,13 @@
     NativeOutputType,
     PandasDF,
     PandasOnSparkDF,
+    PolarsEagerDF,
     PysparkDF,
     UnsupportedDataFrameException,
 )
 from .utils import (
     is_pandas_dataframe,
     is_pandas_on_spark_dataframe,
+    is_polars_eager_dataframe,
     is_pyspark_dataframe,
 )

deepnote_toolkit/ocelots/dataframe.py

@@ -23,6 +23,9 @@
 
 from deepnote_toolkit.ocelots.filters import Filter
 from deepnote_toolkit.ocelots.pandas.implementation import PandasImplementation
+from deepnote_toolkit.ocelots.polars.implementation_eager import (
+    PolarsEagerImplementation,
+)
 from deepnote_toolkit.ocelots.pyspark.implementation import PysparkImplementation
 from deepnote_toolkit.ocelots.types import (
     Column,
@@ -32,16 +35,20 @@
     NativeOutputType,
     PandasDF,
     PandasOnSparkDF,
+    PolarsEagerDF,
     PysparkDF,
     UnsupportedDataFrameException,
 )
 from deepnote_toolkit.ocelots.utils import (
     is_pandas_dataframe,
     is_pandas_on_spark_dataframe,
+    is_polars_eager_dataframe,
     is_pyspark_dataframe,
 )
 
-Implementation = Union[PandasImplementation, PysparkImplementation]
+Implementation = Union[
+    PandasImplementation, PysparkImplementation, PolarsEagerImplementation
+]
 
 T = TypeVar("T", bound=NativeOutputDF)
 FromNativeT = TypeVar("FromNativeT", bound=NativeOutputDF)
@@ -86,6 +93,7 @@ def is_supported(cls, df: Any) -> bool:
             is_pandas_dataframe(df)
             or is_pyspark_dataframe(df)
             or is_pandas_on_spark_dataframe(df)
+            or is_polars_eager_dataframe(df)
         )
 
     # Special case for Pandas-on-Spark DFs, as they aren't wrapped directly, but converted
@@ -117,6 +125,8 @@ def from_native(cls, df: NativeInputDF):
             return cls(PandasImplementation(df))
         if is_pyspark_dataframe(df):
             return cls(PysparkImplementation(df))
+        if is_polars_eager_dataframe(df):
+            return cls(PolarsEagerImplementation(df))
         if is_pandas_on_spark_dataframe(df):
             # NOTE: we accept Pandas-on-Spark dataframes, but we convert them into Spark and
             # work like with it same as with normal Spark DF from that.
@@ -144,18 +154,27 @@ def from_native(cls, df: NativeInputDF):
             return cls(PysparkImplementation(df.to_spark()))
 
         raise UnsupportedDataFrameException(
-            f"expected Pandas or PySpark dataframe, got {type(df)}"
+            f"expected Pandas, PySpark, or Polars dataframe, got {type(df)}"
         )
 
     @property
     def native_type(self) -> NativeOutputType:
         """Get the native type of the dataframe.
 
         Returns:
-            NativeType: Either 'pandas' or 'pyspark'
+            NativeOutputType: Either 'pandas', 'pyspark', or 'polars-eager'
         """
         return self._implementation.name
 
+    @property
+    def lazy(self) -> bool:
+        """Whether the underlying dataframe uses lazy evaluation.
+
+        Lazy dataframes (e.g. PySpark) defer computation until results are collected,
+        while eager ones (e.g. pandas, Polars) evaluate immediately.
+        """
+        return self._implementation.lazy
+
     @property
     def columns(self) -> Tuple[Column, ...]:
         """Get the list of columns in the dataframe.
@@ -425,3 +444,9 @@ def is_wrapped_pandas_dataframe(df: DataFrame) -> TypeGuard[DataFrame[PandasDF]]
 
 def is_wrapped_pyspark_dataframe(df: DataFrame) -> TypeGuard[DataFrame[PysparkDF]]:
     return df.native_type == "pyspark"
+
+
+def is_wrapped_polars_eager_dataframe(
+    df: DataFrame,
+) -> TypeGuard[DataFrame[PolarsEagerDF]]:
+    return df.native_type == "polars-eager"

deepnote_toolkit/ocelots/pandas/implementation.py

@@ -4,6 +4,7 @@
 import pandas as pd
 from typing_extensions import Self
 
+from deepnote_toolkit.logging import LoggerManager
 from deepnote_toolkit.ocelots.constants import (
     DEEPNOTE_INDEX_COLUMN,
     MAX_COLUMNS_TO_DISPLAY,
@@ -21,11 +22,14 @@
     flatten_column_name,
 )
 
+logger = LoggerManager().get_logger()
+
 
 class PandasImplementation:
     """Implementation of DataFrame methods for pandas dataframes."""
 
     name: Literal["pandas"] = "pandas"
+    lazy: bool = False
 
     def __init__(self, df: pd.DataFrame):
         self._df = df
@@ -284,7 +288,8 @@ def filter(self, *filters: Filter) -> Self:
 
                 masks.append(mask)
 
-            except (ValueError, TypeError):
+            except (ValueError, TypeError) as e:
+                logger.warning("Skipping filter on column %r: %s", filter_obj.column, e)
                 continue
 
         if masks: