dataframely 2.1.0__tar.gz → 2.2.0__tar.gz

{dataframely-2.1.0 → dataframely-2.2.0}/.github/copilot-instructions.md

@@ -202,7 +202,7 @@ validated_df: dy.DataFrame[MySchema] = MySchema.validate(df, cast=True)
 
 1. **Python code**: Run `pixi run pre-commit run` before committing
 2. **Rust code**: Run `pixi run postinstall` to rebuild, then run tests
-3. **Tests**: Ensure `pixi run test` passes
+3. **Tests**: Ensure `pixi run test` passes. If changes might affect storage backends, use `pixi run test -m s3`.
 4. **Documentation**: Update docstrings
 5. **API changes**: Ensure backward compatibility or document migration path
 

{dataframely-2.1.0 → dataframely-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataframely
-Version: 2.1.0
+Version: 2.2.0
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11

{dataframely-2.1.0 → dataframely-2.2.0}/dataframely/__init__.py

@@ -51,6 +51,7 @@ from .columns import (
     UInt64,
 )
 from .config import Config
+from .exc import DeserializationError
 from .filter_result import FailureInfo
 from .functional import (
     concat_collection_members,
@@ -106,4 +107,5 @@ __all__ = [
     "Array",
     "Object",
     "Validation",
+    "DeserializationError",
 ]

{dataframely-2.1.0 → dataframely-2.2.0}/dataframely/collection/collection.py

@@ -12,7 +12,7 @@ from collections.abc import Iterable, Mapping, Sequence
 from dataclasses import asdict
 from json import JSONDecodeError
 from pathlib import Path
-from typing import IO, Annotated, Any, Literal, cast
+from typing import IO, Annotated, Any, Literal, cast, overload
 
 import polars as pl
 import polars.exceptions as plexc
@@ -33,7 +33,11 @@ from dataframely._storage.constants import COLLECTION_METADATA_KEY
 from dataframely._storage.delta import DeltaStorageBackend
 from dataframely._storage.parquet import ParquetStorageBackend
 from dataframely._typing import LazyFrame, Validation
-from dataframely.exc import ValidationError, ValidationRequiredError
+from dataframely.exc import (
+    DeserializationError,
+    ValidationError,
+    ValidationRequiredError,
+)
 from dataframely.filter_result import FailureInfo
 from dataframely.random import Generator
 from dataframely.schema import _schema_from_dict
@@ -891,13 +895,13 @@ class Collection(BaseCollection, ABC):
                 - `"allow"`: The method tries to read the schema data from the parquet
                   files. If the stored collection schema matches this collection
                   schema, the collection is read without validation. If the stored
-                  schema mismatches this schema no metadata can be found in
+                  schema mismatches this schema, no valid metadata can be found in
                   the parquets, or the files have conflicting metadata,
                   this method automatically runs :meth:`validate` with `cast=True`.
                 - `"warn"`: The method behaves similarly to `"allow"`. However,
                   it prints a warning if validation is necessary.
                 - `"forbid"`: The method never runs validation automatically and only
-                  returns if the metadata stores a collection schema that matches
+                  returns if the metadata stores a valid collection schema that matches
                   this collection.
                 - `"skip"`: The method never runs validation and simply reads the
                   data, entrusting the user that the schema is valid. *Use this option
@@ -1184,7 +1188,12 @@ class Collection(BaseCollection, ABC):
                 members=cls.member_schemas().keys(), **kwargs
             )
 
-        collection_types = _deserialize_types(serialized_collection_types)
+        # Use strict=False when validation is "allow", "warn" or "skip" to tolerate
+        # missing or broken collection metadata.
+        strict = validation == "forbid"
+        collection_types = _deserialize_types(
+            serialized_collection_types, strict=strict
+        )
         collection_type = _reconcile_collection_types(collection_types)
 
         if cls._requires_validation_for_reading_parquets(collection_type, validation):
@@ -1245,14 +1254,27 @@ def read_parquet_metadata_collection(
     """
     metadata = pl.read_parquet_metadata(source)
     if (schema_metadata := metadata.get(COLLECTION_METADATA_KEY)) is not None:
-        try:
-            return deserialize_collection(schema_metadata)
-        except (JSONDecodeError, plexc.ComputeError):
-            return None
+        return deserialize_collection(schema_metadata, strict=False)
     return None
 
 
-def deserialize_collection(data: str) -> type[Collection]:
+@overload
+def deserialize_collection(
+    data: str, strict: Literal[True] = True
+) -> type[Collection]: ...
+
+
+@overload
+def deserialize_collection(
+    data: str, strict: Literal[False]
+) -> type[Collection] | None: ...
+
+
+@overload
+def deserialize_collection(data: str, strict: bool) -> type[Collection] | None: ...
+
+
+def deserialize_collection(data: str, strict: bool = True) -> type[Collection] | None:
     """Deserialize a collection from a JSON string.
 
     This method allows to dynamically load a collection from its serialization, without
@@ -1260,12 +1282,14 @@ def deserialize_collection(data: str) -> type[Collection]:
 
     Args:
         data: The JSON string created via :meth:`Collection.serialize`.
+        strict: Whether to raise an exception if the collection cannot be deserialized.
 
     Returns:
         The collection loaded from the JSON data.
 
     Raises:
-        ValueError: If the schema format version is not supported.
+        DeserializationError: If the collection can not be deserialized
+            and `strict=True`.
 
     Attention:
         The returned collection **cannot** be used to create instances of the
@@ -1280,34 +1304,41 @@ def deserialize_collection(data: str) -> type[Collection]:
     See also:
         :meth:`Collection.serialize` for additional information on serialization.
     """
-    decoded = json.loads(data, cls=SchemaJSONDecoder)
-    if (format := decoded["versions"]["format"]) != SERIALIZATION_FORMAT_VERSION:
-        raise ValueError(f"Unsupported schema format version: {format}")
-
-    annotations: dict[str, Any] = {}
-    for name, info in decoded["members"].items():
-        lf_type = LazyFrame[_schema_from_dict(info["schema"])]  # type: ignore
-        if info["is_optional"]:
-            lf_type = lf_type | None  # type: ignore
-        annotations[name] = Annotated[
-            lf_type,
-            CollectionMember(
-                ignored_in_filters=info["ignored_in_filters"],
-                inline_for_sampling=info["inline_for_sampling"],
-            ),
-        ]
-
-    return type(
-        f"{decoded['name']}_dynamic",
-        (Collection,),
-        {
-            "__annotations__": annotations,
-            **{
-                name: Filter(logic=lambda _, logic=logic: logic)  # type: ignore
-                for name, logic in decoded["filters"].items()
+    try:
+        decoded = json.loads(data, cls=SchemaJSONDecoder)
+        if (format := decoded["versions"]["format"]) != SERIALIZATION_FORMAT_VERSION:
+            raise ValueError(f"Unsupported schema format version: {format}")
+
+        annotations: dict[str, Any] = {}
+        for name, info in decoded["members"].items():
+            lf_type = LazyFrame[_schema_from_dict(info["schema"])]  # type: ignore
+            if info["is_optional"]:
+                lf_type = lf_type | None  # type: ignore
+            annotations[name] = Annotated[
+                lf_type,
+                CollectionMember(
+                    ignored_in_filters=info["ignored_in_filters"],
+                    inline_for_sampling=info["inline_for_sampling"],
+                ),
+            ]
+
+        return type(
+            f"{decoded['name']}_dynamic",
+            (Collection,),
+            {
+                "__annotations__": annotations,
+                **{
+                    name: Filter(logic=lambda _, logic=logic: logic)  # type: ignore
+                    for name, logic in decoded["filters"].items()
+                },
             },
-        },
-    )
+        )
+    except (ValueError, TypeError, JSONDecodeError, plexc.ComputeError) as e:
+        if strict:
+            raise DeserializationError(
+                "The Collection metadata could not be deserialized"
+            ) from e
+        return None
 
 
 # --------------------------------------- UTILS -------------------------------------- #
@@ -1333,14 +1364,15 @@ def _extract_keys_if_exist(
 
 def _deserialize_types(
     serialized_collection_types: Iterable[str | None],
+    strict: bool = True,
 ) -> list[type[Collection]]:
     collection_types = []
-    collection_type: type[Collection] | None = None
     for t in serialized_collection_types:
         if t is None:
             continue
-        collection_type = deserialize_collection(t)
-        collection_types.append(collection_type)
+        collection_type = deserialize_collection(t, strict=strict)
+        if collection_type is not None:
+            collection_types.append(collection_type)
 
     return collection_types
 

{dataframely-2.1.0 → dataframely-2.2.0}/dataframely/exc.py

@@ -41,3 +41,10 @@ class AnnotationImplementationError(ImplementationError):
 
 class ValidationRequiredError(Exception):
     """Error raised when validation is required when reading a parquet file."""
+
+
+# ---------------------------------- DESERIALIZATION --------------------------------- #
+
+
+class DeserializationError(Exception):
+    """Error raised when deserialization of a schema or collection fails."""

{dataframely-2.1.0 → dataframely-2.2.0}/dataframely/schema.py

@@ -40,7 +40,12 @@ from ._storage.parquet import (
 from ._typing import DataFrame, LazyFrame, Validation
 from .columns import Column, column_from_dict
 from .config import Config
-from .exc import SchemaError, ValidationError, ValidationRequiredError
+from .exc import (
+    DeserializationError,
+    SchemaError,
+    ValidationError,
+    ValidationRequiredError,
+)
 from .filter_result import FailureInfo, FilterResult, LazyFilterResult
 from .random import Generator
 
@@ -1238,8 +1243,13 @@ class Schema(BaseSchema, ABC):
         validation: Validation,
         source: str,
     ) -> DataFrame[Self] | LazyFrame[Self]:
+        # Use strict=False when validation is "allow", "warn" or "skip" to tolerate
+        # deserialization failures from old serialized formats.
+        strict = validation == "forbid"
         deserialized_schema = (
-            deserialize_schema(serialized_schema) if serialized_schema else None
+            deserialize_schema(serialized_schema, strict=strict)
+            if serialized_schema
+            else None
         )
 
         # Smart validation
@@ -1347,6 +1357,10 @@ def deserialize_schema(data: str, strict: Literal[True] = True) -> type[Schema]:
 def deserialize_schema(data: str, strict: Literal[False]) -> type[Schema] | None: ...
 
 
+@overload
+def deserialize_schema(data: str, strict: bool) -> type[Schema] | None: ...
+
+
 def deserialize_schema(data: str, strict: bool = True) -> type[Schema] | None:
     """Deserialize a schema from a JSON string.
 
@@ -1375,9 +1389,11 @@ def deserialize_schema(data: str, strict: bool = True) -> type[Schema] | None:
         if (format := decoded["versions"]["format"]) != SERIALIZATION_FORMAT_VERSION:
             raise ValueError(f"Unsupported schema format version: {format}")
         return _schema_from_dict(decoded)
-    except (ValueError, JSONDecodeError, plexc.ComputeError) as e:
+    except (ValueError, JSONDecodeError, plexc.ComputeError, TypeError) as e:
         if strict:
-            raise e from e
+            raise DeserializationError(
+                "The Schema metadata could not be deserialized"
+            ) from e
         return None
 
 

{dataframely-2.1.0 → dataframely-2.2.0}/dataframely/testing/storage.py

@@ -29,11 +29,11 @@ class SchemaStorageTester(ABC):
     def write_typed(
         self, schema: type[S], df: dy.DataFrame[S], path: str, lazy: bool
     ) -> None:
-        """Write a schema to the backend without recording schema information."""
+        """Write a schema to the backend and record schema information."""
 
     @abstractmethod
     def write_untyped(self, df: pl.DataFrame, path: str, lazy: bool) -> None:
-        """Write a schema to the backend and record schema information."""
+        """Write a schema to the backend without recording schema information."""
 
     @overload
     def read(
@@ -45,12 +45,22 @@ class SchemaStorageTester(ABC):
         self, schema: type[S], path: str, lazy: Literal[False], validation: Validation
     ) -> dy.DataFrame[S]: ...
 
+    @overload
+    def read(
+        self, schema: type[S], path: str, lazy: bool, validation: Validation
+    ) -> dy.LazyFrame[S] | dy.DataFrame[S]: ...
+
     @abstractmethod
     def read(
         self, schema: type[S], path: str, lazy: bool, validation: Validation
     ) -> dy.LazyFrame[S] | dy.DataFrame[S]:
         """Read from the backend, using schema information if available."""
 
+    @abstractmethod
+    def set_metadata(self, path: str, metadata: dict[str, Any]) -> None:
+        """Overwrite the metadata stored at the given path with the provided
+        metadata."""
+
 
 class ParquetSchemaStorageTester(SchemaStorageTester):
     """Testing interface for the parquet storage functionality of Schema."""
@@ -83,6 +93,11 @@ class ParquetSchemaStorageTester(SchemaStorageTester):
         self, schema: type[S], path: str, lazy: Literal[False], validation: Validation
     ) -> dy.DataFrame[S]: ...
 
+    @overload
+    def read(
+        self, schema: type[S], path: str, lazy: bool, validation: Validation
+    ) -> dy.LazyFrame[S] | dy.DataFrame[S]: ...
+
     def read(
         self, schema: type[S], path: str, lazy: bool, validation: Validation
     ) -> dy.LazyFrame[S] | dy.DataFrame[S]:
@@ -93,6 +108,11 @@ class ParquetSchemaStorageTester(SchemaStorageTester):
         else:
             return schema.read_parquet(self._wrap_path(path), validation=validation)
 
+    def set_metadata(self, path: str, metadata: dict[str, Any]) -> None:
+        target = self._wrap_path(path)
+        data = pl.read_parquet(target)
+        data.write_parquet(target, metadata=metadata)
+
 
 class DeltaSchemaStorageTester(SchemaStorageTester):
     """Testing interface for the deltalake storage functionality of Schema."""
@@ -115,6 +135,11 @@ class DeltaSchemaStorageTester(SchemaStorageTester):
         self, schema: type[S], path: str, lazy: Literal[False], validation: Validation
     ) -> dy.DataFrame[S]: ...
 
+    @overload
+    def read(
+        self, schema: type[S], path: str, lazy: bool, validation: Validation
+    ) -> dy.LazyFrame[S] | dy.DataFrame[S]: ...
+
     def read(
         self, schema: type[S], path: str, lazy: bool, validation: Validation
     ) -> dy.DataFrame[S] | dy.LazyFrame[S]:
@@ -122,6 +147,18 @@ class DeltaSchemaStorageTester(SchemaStorageTester):
             return schema.scan_delta(path, validation=validation)
         return schema.read_delta(path, validation=validation)
 
+    def set_metadata(self, path: str, metadata: dict[str, Any]) -> None:
+        df = pl.read_delta(path)
+        df.head(0).write_delta(
+            path,
+            delta_write_options={
+                "commit_properties": deltalake.CommitProperties(
+                    custom_metadata=metadata
+                ),
+            },
+            mode="overwrite",
+        )
+
 
 # ------------------------------- Collection -------------------------------------------
 
@@ -147,11 +184,36 @@ class CollectionStorageTester(ABC):
     def read(self, collection: type[C], path: str, lazy: bool, **kwargs: Any) -> C:
         """Read from the backend, using collection information if available."""
 
+    @abstractmethod
+    def set_metadata(self, path: str, metadata: dict[str, Any]) -> None:
+        """Overwrite the metadata stored at the given path with the provided
+        metadata."""
+
+    def _prefix_path(self, path: str, fs: AbstractFileSystem) -> str:
+        return f"{self._get_prefix(fs)}{path}"
+
+    @staticmethod
+    def _get_prefix(fs: AbstractFileSystem) -> str:
+        return (
+            ""
+            if fs.protocol == "file"
+            else (
+                f"{fs.protocol}://"
+                if isinstance(fs.protocol, str)
+                else f"{fs.protocol[0]}://"
+            )
+        )
+
 
 class ParquetCollectionStorageTester(CollectionStorageTester):
     def write_typed(
         self, collection: dy.Collection, path: str, lazy: bool, **kwargs: Any
     ) -> None:
+        if "metadata" in kwargs:  # pragma: no cover
+            raise KeyError(
+                "`metadata` kwarg will be ignored in `write_typed`. Use `set_metadata`."
+            )
+
         # Polars does not support partitioning via kwarg on sink_parquet
         if lazy:
             kwargs.pop("partition_by", None)
@@ -164,6 +226,11 @@ class ParquetCollectionStorageTester(CollectionStorageTester):
     def write_untyped(
         self, collection: dy.Collection, path: str, lazy: bool, **kwargs: Any
     ) -> None:
+        if "metadata" in kwargs:  # pragma: no cover
+            raise KeyError(
+                "Cannot set metadata through `write_untyped`. Use `set_metadata`."
+            )
+
         if lazy:
             collection.sink_parquet(path, **kwargs)
         else:
@@ -175,17 +242,8 @@ class ParquetCollectionStorageTester(CollectionStorageTester):
             df.write_parquet(file)
 
         fs: AbstractFileSystem = url_to_fs(path)[0]
-        prefix = (
-            ""
-            if fs.protocol == "file"
-            else (
-                f"{fs.protocol}://"
-                if isinstance(fs.protocol, str)
-                else f"{fs.protocol[0]}://"
-            )
-        )
         for file in fs.glob(fs.sep.join([path, "**", "*.parquet"])):
-            _delete_meta(f"{prefix}{file}")
+            _delete_meta(self._prefix_path(file, fs))
 
     def read(self, collection: type[C], path: str, lazy: bool, **kwargs: Any) -> C:
         if lazy:
@@ -193,11 +251,23 @@ class ParquetCollectionStorageTester(CollectionStorageTester):
         else:
             return collection.read_parquet(path, **kwargs)
 
+    def set_metadata(self, path: str, metadata: dict[str, Any]) -> None:
+        fs: AbstractFileSystem = url_to_fs(path)[0]
+        for file in fs.glob(fs.sep.join([path, "*.parquet"])):
+            file_path = self._prefix_path(file, fs)
+            df = pl.read_parquet(file_path)
+            df.write_parquet(file_path, metadata=metadata)
+
 
 class DeltaCollectionStorageTester(CollectionStorageTester):
     def write_typed(
         self, collection: dy.Collection, path: str, lazy: bool, **kwargs: Any
     ) -> None:
+        if "metadata" in kwargs:  # pragma: no cover
+            raise KeyError(
+                "`metadata` kwarg will be ignored in `write_typed`. Use `set_metadata`."
+            )
+
         extra_kwargs = {}
         if partition_by := kwargs.pop("partition_by", None):
             extra_kwargs["delta_write_options"] = {"partition_by": partition_by}
@@ -207,6 +277,10 @@ class DeltaCollectionStorageTester(CollectionStorageTester):
     def write_untyped(
         self, collection: dy.Collection, path: str, lazy: bool, **kwargs: Any
     ) -> None:
+        if "metadata" in kwargs:  # pragma: no cover
+            raise KeyError(
+                "Cannot set metadata through `write_untyped`. Use `set_metadata`."
+            )
         collection.write_delta(path, **kwargs)
 
         # For each member table, write an empty commit
@@ -222,6 +296,23 @@ class DeltaCollectionStorageTester(CollectionStorageTester):
             return collection.scan_delta(source=path, **kwargs)
         return collection.read_delta(source=path, **kwargs)
 
+    def set_metadata(self, path: str, metadata: dict[str, Any]) -> None:
+        fs: AbstractFileSystem = url_to_fs(path)[0]
+        # For delta, we need to update metadata on each member table
+        for entry in fs.ls(path):
+            member_path = self._prefix_path(entry, fs)
+            if fs.isdir(member_path):
+                df = pl.read_delta(member_path)
+                df.head(0).write_delta(
+                    member_path,
+                    delta_write_options={
+                        "commit_properties": deltalake.CommitProperties(
+                            custom_metadata=metadata
+                        ),
+                    },
+                    mode="overwrite",
+                )
+
 
 # ------------------------------------ Failure info ------------------------------------
 class FailureInfoStorageTester(ABC):