Write to parquet

pyproject.toml

@@ -41,7 +41,7 @@ test = [
     "scipy",
     "pooch",
     "ruff",
-
+    "fastparquet",
 ]
 
 

virtualizarr/tests/test_kerchunk.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pandas as pd
 import pytest
 import ujson  # type: ignore
 import xarray as xr
@@ -127,6 +128,47 @@ def test_accessor_to_kerchunk_json(self, tmp_path):
         }
         assert loaded_refs == expected_ds_refs
 
+    def test_accessor_to_kerchunk_parquet(self, tmp_path):
+        chunks_dict = {
+            "0.0": {"path": "foo.nc", "offset": 100, "length": 100},
+            "0.1": {"path": "foo.nc", "offset": 200, "length": 100},
+        }
+        manifest = ChunkManifest(entries=chunks_dict)
+        arr = ManifestArray(
+            chunkmanifest=manifest,
+            zarray=dict(
+                shape=(2, 4),
+                dtype=np.dtype("<i8"),
+                chunks=(2, 2),
+                compressor=None,
+                filters=None,
+                fill_value=None,
+                order="C",
+            ),
+        )
+        ds = xr.Dataset({"a": (["x", "y"], arr)})
+
+        filepath = tmp_path / "refs"
+
+        ds.virtualize.to_kerchunk(filepath, format="parquet", record_size=2)
+
+        with open(tmp_path / "refs" / ".zmetadata") as f:
+            meta = ujson.load(f)
+            assert list(meta) == ["metadata", "record_size"]
+            assert meta["record_size"] == 2
+
+        df0 = pd.read_parquet(filepath / "a" / "refs.0.parq")
+
+        assert df0.to_dict() == {
+            "offset": {0: 100, 1: 200},
+            "path": {
+                0: "foo.nc",
+                1: "foo.nc",
+            },
+            "size": {0: 100, 1: 100},
+            "raw": {0: None, 1: None},
+        }
+
 
 def test_kerchunk_roundtrip_in_memory_no_concat():
     # Set up example xarray dataset

virtualizarr/xarray.py

@@ -359,15 +359,23 @@ def to_kerchunk(
     ) -> KerchunkStoreRefs: ...
 
     @overload
-    def to_kerchunk(self, filepath: str, format: Literal["json"]) -> None: ...
+    def to_kerchunk(self, filepath: str | Path, format: Literal["json"]) -> None: ...
 
     @overload
-    def to_kerchunk(self, filepath: str, format: Literal["parquet"]) -> None: ...
+    def to_kerchunk(
+        self,
+        filepath: str | Path,
+        format: Literal["parquet"],
+        record_size: int = 100_000,
+        categorical_threshold: int = 10,
+    ) -> None: ...
 
     def to_kerchunk(
         self,
-        filepath: Optional[str] = None,
+        filepath: str | Path | None = None,
         format: Union[Literal["dict"], Literal["json"], Literal["parquet"]] = "dict",
+        record_size: int = 100_000,
+        categorical_threshold: int = 10,
     ) -> Union[KerchunkStoreRefs, None]:
         """
         Serialize all virtualized arrays in this xarray dataset into the kerchunk references format.
@@ -379,6 +387,13 @@ def to_kerchunk(
         format : 'dict', 'json', or 'parquet'
             Format to serialize the kerchunk references as.
             If 'json' or 'parquet' then the 'filepath' argument is required.
+        record_size (parquet only): int
+            Number of references to store in each reference file (default 100,000). Bigger values
+            mean fewer read requests but larger memory footprint.
+        categorical_threshold (parquet only) : int
+            Encode urls as pandas.Categorical to reduce memory footprint if the ratio
+            of the number of unique urls to total number of refs for each variable
+            is greater than or equal to this number. (default 10)
 
         References
         ----------
@@ -397,6 +412,18 @@ def to_kerchunk(
 
             return None
         elif format == "parquet":
-            raise NotImplementedError()
+            from kerchunk.df import refs_to_dataframe
+
+            if isinstance(filepath, Path):
+                url = str(filepath)
+            elif isinstance(filepath, str):
+                url = filepath
+
+            return refs_to_dataframe(
+                refs,
+                url=url,
+                record_size=record_size,
+                categorical_threshold=categorical_threshold,
+            )
         else:
             raise ValueError(f"Unrecognized output format: {format}")