diff --git a/.gitignore b/.gitignore index 37a2dcd..22ab01b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ data +.DS_Store # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/README.md b/README.md index b97c3fb..499778e 100644 --- a/README.md +++ b/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. diff --git a/lambda/virtualizarr-processor/pyproject.toml b/lambda/virtualizarr-processor/pyproject.toml index 2960f0d..c0365e5 100644 --- a/lambda/virtualizarr-processor/pyproject.toml +++ b/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", ] diff --git a/lambda/virtualizarr-processor/virtualizarr_processor/processor.py b/lambda/virtualizarr-processor/virtualizarr_processor/processor.py index b68a446..3769976 100644 --- a/lambda/virtualizarr-processor/virtualizarr_processor/processor.py +++ b/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) diff --git a/lambda/virtualizarr-processor/virtualizarr_processor/typing.py b/lambda/virtualizarr-processor/virtualizarr_processor/typing.py index f365f85..bb62379 100644 --- a/lambda/virtualizarr-processor/virtualizarr_processor/typing.py +++ b/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 diff --git a/tests/test_example.py b/tests/test_example.py index 7073a2c..7df30f7 100644 --- a/tests/test_example.py +++ b/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) diff --git a/uv.lock b/uv.lock index b848a41..089002c 100644 --- a/uv.lock +++ b/uv.lock @@ -802,6 +802,28 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/70/44/5191d2e4026f86a2a109053e194d3ba7a31a2d10a9c2348368c63ed4e85a/pandas-2.3.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:3869faf4bd07b3b66a9f462417d0ca3a9df29a9f6abd5d0d0dbab15dac7abe87", size = 13202175, upload-time = "2025-09-29T23:31:59.173Z" }, ] +[[package]] +name = "pandera" +version = "0.31.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "packaging" }, + { name = "pydantic" }, + { name = "typeguard" }, + { name = "typing-extensions" }, + { name = "typing-inspect" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b7/e1/aaa14c989ffd30c7acb293fb986715517ca2b5b435ca291432535bb2b111/pandera-0.31.1.tar.gz", hash = "sha256:c75aa3868af15d4f9aa613acf1a7f436a518f81f1eb658ad630c1dbe1dab0f13", size = 729785, upload-time = "2026-04-15T03:18:59.967Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/2d/6ea7cad2c2f0625c4120bef5353ab7cf749141bf1d070011cebb72f68189/pandera-0.31.1-py3-none-any.whl", hash = "sha256:f9f1ff4852804e1a181a4cb968e732a492f4b6dbefe051a8c5500da43d5c326d", size = 386913, upload-time = "2026-04-15T03:18:58.358Z" }, +] + +[package.optional-dependencies] +xarray = [ + { name = "numpy" }, + { name = "xarray" }, +] + [[package]] name = "pathlib-abc" version = "0.5.2" @@ -1170,6 +1192,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, ] +[[package]] +name = "typing-inspect" +version = "0.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mypy-extensions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/dc/74/1789779d91f1961fa9438e9a8710cdae6bd138c80d7303996933d117264a/typing_inspect-0.9.0.tar.gz", hash = "sha256:b23fc42ff6f6ef6954e4852c1fb512cdd18dbea03134f91f856a95ccc9461f78", size = 13825, upload-time = "2023-05-24T20:25:47.612Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/65/f3/107a22063bf27bdccf2024833d3445f4eea42b2e598abfbd46f6a63b6cb0/typing_inspect-0.9.0-py3-none-any.whl", hash = "sha256:9ee6fc59062311ef8547596ab6b955e1b8aa46242d854bfc78f4f6b0eff35f9f", size = 8827, upload-time = "2023-05-24T20:25:45.287Z" }, +] + [[package]] name = "typing-inspection" version = "0.4.2" @@ -1366,12 +1401,14 @@ version = "0.1.0" source = { editable = "lambda/virtualizarr-processor" } dependencies = [ { name = "icechunk" }, + { name = "pandera", extra = ["xarray"] }, { name = "virtualizarr" }, ] [package.metadata] requires-dist = [ { name = "icechunk" }, + { name = "pandera", extras = ["xarray"], specifier = ">=0.31.0" }, { name = "virtualizarr" }, ]