feat: support polars dataframe in CachedFoundryClient.save_dataset (#103)

* feat: support polars dataframe in CachedFoundryClient.save_dataset

* Update docstring

Co-authored-by: Nicolas Renkamp <nicornk@users.noreply.github.com>

---------

Co-authored-by: Nicolas Renkamp <nicornk@users.noreply.github.com>

Author: Waschenbacher

libs/foundry-dev-tools/src/foundry_dev_tools/cached_foundry_client.py

@@ -27,6 +27,7 @@
 
 if TYPE_CHECKING:
     import pandas as pd
+    import polars as pl
     import pyspark.sql
 
     from foundry_dev_tools.utils import api_types
@@ -185,7 +186,7 @@ def _download_dataset_to_cache_dir(
 
     def save_dataset(
         self,
-        df: pd.DataFrame | pyspark.sql.DataFrame,
+        df: pd.DataFrame | pyspark.sql.DataFrame | pl.DataFrame,
         dataset_path_or_rid: str,
         branch: str = "master",
         exists_ok: bool = False,
@@ -198,8 +199,8 @@ def save_dataset(
         Creates SNAPSHOT transactions by default.
 
         Args:
-            df (:external+pandas:py:class:`pandas.DataFrame` | :external+spark:py:class:`pyspark.sql.DataFrame`): A
-                pyspark  or pandas DataFrame to upload
+            df (:external+pandas:py:class:`pandas.DataFrame` | :external+polars:py:class:`polars.DataFrame` | :external+spark:py:class:`pyspark.sql.DataFrame`):
+                A pyspark, pandas or polars DataFrame to upload
             dataset_path_or_rid (str): Path or Rid of the dataset in which the object should be stored.
             branch (str): Branch of the dataset in which the object should be stored
             exists_ok (bool): By default, this method creates a new dataset.
@@ -217,7 +218,7 @@ def save_dataset(
             ValueError: when dataframe is None
             ValueError: when branch is None
 
-        """
+        """  # noqa: E501
         if df is None:
             msg = "Please provide a spark or pandas dataframe object with parameter 'df'"
             raise ValueError(msg)
@@ -227,6 +228,7 @@ def save_dataset(
 
         with tempfile.TemporaryDirectory() as path:
             from foundry_dev_tools._optional.pandas import pd
+            from foundry_dev_tools._optional.polars import pl
 
             if not pd.__fake__ and isinstance(df, pd.DataFrame):
                 df.to_parquet(
@@ -235,6 +237,12 @@ def save_dataset(
                     compression="snappy",
                     flavor="spark",
                 )
+            elif not pl.__fake__ and isinstance(df, pl.DataFrame):
+                df.write_parquet(
+                    os.sep.join([path + "/dataset.parquet"]),  # noqa: PTH118
+                    use_pyarrow=True,
+                    compression="snappy",
+                )
             else:
                 df.write.format("parquet").option("compression", "snappy").save(path=path, mode="overwrite")
 

tests/unit/test_cached_foundry_client.py

@@ -5,6 +5,7 @@
 from unittest.mock import MagicMock, patch
 
 import pandas as pd
+import polars as pl
 import pytest
 from pandas.testing import assert_frame_equal
 from pyspark.sql import SparkSession
@@ -34,22 +35,34 @@
 )
 @patch(API + ".infer_dataset_schema")
 @patch(API + ".upload_dataset_schema")
-def test_save_pandas(upload_dataset_schema, infer_dataset_schema, save_objects, temporary_directory, test_context_mock):
+@pytest.mark.parametrize(
+    "df",
+    [
+        pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}),
+        pl.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}),
+    ],
+)
+def test_save_pandas(
+    upload_dataset_schema, infer_dataset_schema, save_objects, temporary_directory, test_context_mock, df
+):
     mtp = "mock_tmp_path"
     dsn = "dataset.parquet"
     dsp = mtp + "/" + dsn
     temporary_directory.return_value.__enter__.return_value = mtp
     fdt = CachedFoundryClient(ctx=test_context_mock)
-    df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]})
-    with patch.object(df, "to_parquet") as pd_to_parquet:
+
+    method_to_patch = "to_parquet" if isinstance(df, pd.DataFrame) else "write_parquet"
+
+    with patch.object(df, method_to_patch) as df_to_parquet:
         dataset_rid, transaction_id = fdt.save_dataset(
             df,
             dataset_path_or_rid=DATASET_PATH,
             branch="master",
             exists_ok=True,
             mode="SNAPSHOT",
         )
-    assert pd_to_parquet.call_args[0][0] == dsp
+
+    assert df_to_parquet.call_args[0][0] == dsp
 
     assert save_objects.call_args[0][0] == {"spark/" + dsn: Path(dsp)}
     assert save_objects.call_args[0][1] == DATASET_PATH