Support writing polars.DataFrame to parquet.

Author: nathanjmcdougall

pins/drivers.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Sequence
+from typing import Literal, Sequence, assert_never
 
 from .config import PINS_ENV_INSECURE_READ, get_allow_pickle_read
 from .errors import PinsInsecureReadError
@@ -14,14 +14,38 @@
 
 
 def _assert_is_pandas_df(x, file_type: str) -> None:
-    import pandas as pd
+    df_family = _get_df_family(x)
 
-    if not isinstance(x, pd.DataFrame):
+    if df_family != "pandas":
         raise NotImplementedError(
             f"Currently only pandas.DataFrame can be saved as type {file_type!r}."
         )
 
 
+def _get_df_family(df) -> Literal["unknown", "pandas", "polars"]:
+    try:
+        import polars as pl
+    except ModuleNotFoundError:
+        is_polars_df = False
+    else:
+        is_polars_df = isinstance(df, pl.DataFrame)
+
+    import pandas as pd
+
+    is_pandas_df = isinstance(df, pd.DataFrame)
+
+    if not is_polars_df and not is_pandas_df:
+        return "unknown"
+    if is_polars_df and is_pandas_df:  # Hybrid DataFrame type!
+        return "unknown"
+    elif is_polars_df:
+        return "polars"
+    elif is_pandas_df:
+        return "pandas"
+    else:
+        assert_never(df)
+
+
 def load_path(meta, path_to_version):
     # Check that only a single file name was given
     fnames = [meta.file] if isinstance(meta.file, str) else meta.file
@@ -171,9 +195,17 @@ def save_data(obj, fname, type=None, apply_suffix: bool = True) -> "str | Sequen
         )
 
     elif type == "parquet":
-        _assert_is_pandas_df(obj, file_type=type)
-
-        obj.to_parquet(final_name)
+        df_family = _get_df_family(obj)
+        if df_family == "polars":
+            obj.write_parquet(final_name)
+        elif df_family == "pandas":
+            obj.to_parquet(final_name)
+        else:
+            msg = (
+                "Currently only pandas.DataFrame and polars.DataFrame can be saved to "
+                "a parquet file."
+            )
+            raise NotImplementedError(msg)
 
     elif type == "joblib":
         import joblib

pins/tests/test_drivers.py

@@ -76,6 +76,36 @@ def test_driver_roundtrip(tmp_path: Path, type_):
     assert df.equals(obj)
 
 
+@pytest.mark.parametrize(
+    "type_",
+    [
+        "parquet",
+    ],
+)
+def test_driver_polars_roundtrip(tmp_dir2, type_):
+    import polars as pl
+
+    df = pl.DataFrame({"x": [1, 2, 3]})
+
+    fname = "some_df"
+    full_file = f"{fname}.{type_}"
+
+    p_obj = tmp_dir2 / fname
+    res_fname = save_data(df, p_obj, type_)
+
+    assert Path(res_fname).name == full_file
+
+    meta = MetaRaw(full_file, type_, "my_pin")
+    pandas_df = load_data(
+        meta, fsspec.filesystem("file"), tmp_dir2, allow_pickle_read=True
+    )
+
+    # Convert from pandas to polars
+    obj = pl.DataFrame(pandas_df)
+
+    assert df.equals(obj)
+
+
 @pytest.mark.parametrize(
     "type_",
     [