Merge pull request #131 from rstudio/feat-feather-parquet

feat: add support for parquet and arrow storage

Author: machow

.github/workflows/ci.yml

@@ -71,7 +71,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           python -m pip install -r requirements/dev.txt
-          python -m pip install -e .
+          python -m pip install -e .[test]
       - name: run RStudio Connect
         run: |
           docker-compose up --build -d

pins/constructors.py

@@ -255,6 +255,8 @@ def board_github(
     Examples
     --------
 
+    >>> import pytest; pytest.skip()
+
     >>> import os
     >>> board = board_github("machow", "pins-python", "pins/tests/pins-compat")
     >>> board.pin_list()

pins/drivers.py

@@ -14,6 +14,15 @@
 REQUIRES_SINGLE_FILE = frozenset(["csv", "joblib", "file"])
 
 
+def _assert_is_pandas_df(x):
+    import pandas as pd
+
+    if not isinstance(x, pd.DataFrame):
+        raise NotImplementedError(
+            "Currently only pandas.DataFrame can be saved to a CSV."
+        )
+
+
 def load_data(
     meta: Meta,
     fs,
@@ -67,6 +76,16 @@ def load_data(
 
         return pd.read_csv(fs.open(path_to_file))
 
+    elif meta.type == "feather":
+        import pandas as pd
+
+        return pd.read_feather(fs.open(path_to_file))
+
+    elif meta.type == "parquet":
+        import pandas as pd
+
+        return pd.read_parquet(fs.open(path_to_file))
+
     elif meta.type == "table":
         import pandas as pd
 
@@ -93,28 +112,35 @@ def save_data(
     # TODO: would be useful to have singledispatch func for a "default saver"
     #       as argument to board, and then type dispatchers for explicit cases
     #       of saving / loading objects different ways.
+
+    if apply_suffix:
+        final_name = f"{fname}.{type}"
+    else:
+        final_name = fname
+
     if type == "csv":
-        import pandas as pd
+        _assert_is_pandas_df(obj)
 
-        if apply_suffix:
-            fname = f"{fname}.{type}"
+        obj.to_csv(final_name, index=False)
+
+    elif type == "feather":
+        _assert_is_pandas_df(obj)
+
+        obj.to_feather(final_name)
+
+    elif type == "parquet":
+        _assert_is_pandas_df(obj)
+
+        obj.to_parquet(final_name)
 
-        if not isinstance(obj, pd.DataFrame):
-            raise NotImplementedError(
-                "Currently only pandas.DataFrame can be saved to a CSV."
-            )
-        obj.to_csv(fname, index=False)
     elif type == "joblib":
         import joblib
 
-        if apply_suffix:
-            fname = f"{fname}.{type}"
-
-        joblib.dump(obj, fname)
+        joblib.dump(obj, final_name)
     else:
         raise NotImplementedError(f"Cannot save type: {type}")
 
-    return fname
+    return final_name
 
 
 def default_title(obj, name):

pins/tests/test_drivers.py

@@ -45,7 +45,15 @@ def test_default_title(obj, dst_title):
     assert res == dst_title
 
 
-def test_driver_roundtrip_csv(tmp_dir2):
+@pytest.mark.parametrize(
+    "type_",
+    [
+        "csv",
+        "feather",
+        "parquet",
+    ],
+)
+def test_driver_roundtrip(tmp_dir2, type_):
     # TODO: I think this test highlights the challenge of getting the flow
     # between metadata, drivers, and the metafactory right.
     # There is the name of the data (relative to the pin directory), and the full
@@ -55,14 +63,14 @@ def test_driver_roundtrip_csv(tmp_dir2):
     df = pd.DataFrame({"x": [1, 2, 3]})
 
     fname = "some_df"
-    type_ = "csv"
+    full_file = f"{fname}.{type_}"
 
     p_obj = tmp_dir2 / fname
     res_fname = save_data(df, p_obj, type_)
 
-    assert Path(res_fname).name == f"{fname}.csv"
+    assert Path(res_fname).name == full_file
 
-    meta = MetaRaw(f"{fname}.csv", type_, "my_pin")
+    meta = MetaRaw(full_file, type_, "my_pin")
     obj = load_data(meta, fsspec.filesystem("file"), tmp_dir2)
 
     assert df.equals(obj)

setup.cfg

@@ -53,6 +53,8 @@ test =
     pytest-dotenv
     s3fs
     adlfs
+    fastparquet
+    pyarrow
 
 
 [bdist_wheel]