dataset validation

.gitignore

@@ -1,4 +1,5 @@
 data
+.DS_Store
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

README.md

@@ -25,6 +25,8 @@ file.  Your class should follow the [VirtualizarrProcessor protocol](./lambda/vi
 
 - **process_file** This method should take a file uri and a session and use a  Virtualizarr parser to parse it and add the resulting ManifestStore or virtual dataset to the Icechunk store.
 
+- **validate_dataset** This method should validate the parsed xarray Dataset before writing it to Icechunk and return `True` or `False`. The sample processor uses [Pandera xarray validation](https://pandera.readthedocs.io/en/latest/xarray_guide/) as an example.
+
 - **commit_processed_files** This method commits all the changes made during the
   session in a single commit.
 

lambda/virtualizarr-processor/pyproject.toml

@@ -5,6 +5,7 @@ description = "Module for ingesting virtualizarr data to Icechunk"
 requires-python = ">=3.11"
 dependencies = [
     "icechunk",
+    "pandera[xarray]>=0.31.0",
     "virtualizarr",
 ]
 

lambda/virtualizarr-processor/virtualizarr_processor/processor.py

@@ -1,6 +1,6 @@
+import logging
 import os
 import tempfile
-from copy import Error
 from datetime import datetime
 from itertools import islice
 
@@ -9,13 +9,27 @@
 import obstore
 import xarray as xr
 from icechunk import Repository, Session
+from pandera.xarray import Coordinate, DatasetSchema, DataVar
 from virtualizarr.manifests import ChunkManifest, ManifestArray
 from zarr.codecs import BytesCodec
 from zarr.core.dtype import parse_data_type
 from zarr.core.metadata import ArrayV3Metadata
+import pandera as pa
 
 CHUNK_DIR = os.path.realpath(tempfile.gettempdir())
 CHUNK_DIRECTORY_URL_PREFIX = f"file://{CHUNK_DIR}/"
+logger = logging.getLogger(__name__)
+
+
+EXAMPLE_DATASET_SCHEMA = DatasetSchema(
+    data_vars={
+        "foo": DataVar(dtype=np.int64, dims=("y", "x")),
+    },
+    coords={
+        "time": Coordinate(dtype=np.datetime64, dims=("time",)),
+    },
+    strict=True,
+)
 
 
 def synthetic_vds(date: str) -> xr.Dataset:
@@ -80,6 +94,8 @@ def initialize_repo(self) -> Repository:
         if len(history) == 1:
             session = repo.writable_session("main")
             vds = synthetic_vds("2024-01-01")
+            if not self.validate_dataset(vds):
+                raise ValueError("Dataset validation failed for initialization dataset")
             vds.vz.to_icechunk(session.store, validate_containers=False)
             session.commit(message="Initialization")
         return repo
@@ -92,14 +108,27 @@ def process_file(self, file_key: str, session: Session) -> bool:
         result = False
         try:
             vds = synthetic_vds(file_key)
+            if not self.validate_dataset(vds):
+                return False
             vds.vz.to_icechunk(
                 session.store, append_dim="time", validate_containers=False
             )
             result = True
-        except Error:
+        except Exception:
             result = False
         return result
 
+    @classmethod
+    def validate_dataset(cls, dataset: xr.Dataset) -> bool:
+        try:
+            EXAMPLE_DATASET_SCHEMA.validate(dataset, lazy=True)
+            return True
+        except pa.errors.SchemaErrors as e:
+            logger.exception(
+                "Dataset validation failed:",
+            )
+            return False
+
     def commit_processed_files(self, session: Session) -> str:
         snapshot = session.commit(message=f"Append to {session.snapshot_id}")
         return str(snapshot)

lambda/virtualizarr-processor/virtualizarr_processor/typing.py

@@ -4,6 +4,7 @@
 from typing import Protocol, runtime_checkable
 
 import icechunk
+import xarray as xr
 from icechunk import Repository, Session
 
 
@@ -56,6 +57,21 @@ def process_file(self, file_key: str, session: Session) -> bool:
         """
         ...
 
+    @classmethod
+    def validate_dataset(cls, dataset: xr.Dataset) -> bool:
+        """
+        Validate a parsed virtual dataset before writing it to Icechunk.
+
+        Parameters
+        ----------
+            dataset: The parsed xarray Dataset.
+        Returns
+        -------
+        bool
+            True if the dataset passes validation.
+        """
+        ...
+
     def commit_processed_files(self, session: Session) -> str:
         """
         Commits the updates made by one or multiple calls to process_file

tests/test_example.py

@@ -2,7 +2,7 @@
 
 import icechunk
 from icechunk import Repository, Session
-from virtualizarr_processor.processor import Processor
+from virtualizarr_processor.processor import Processor, synthetic_vds
 from virtualizarr_processor.typing import VirtualizarrProcessor
 
 
@@ -27,6 +27,12 @@ def test_process_file(icechunk_session: Session) -> None:
     assert result
 
 
+def test_validate_dataset() -> None:
+    processor = Processor()
+    validated = processor.validate_dataset(synthetic_vds("2024-01-03"))
+    assert validated
+
+
 def test_commit_processed_files(icechunk_session: Session) -> None:
     processor = Processor()
     snapshot = processor.commit_processed_files(session=icechunk_session)