fix: Correct schema deserialization on failure (#128)

Author: florian-klein

dataframely/failure.py

@@ -19,6 +19,7 @@
     from .schema import Schema
 
 RULE_METADATA_KEY = "dataframely_rule_columns"
+UNKNOWN_SCHEMA_NAME = "__DATAFRAMELY_UNKNOWN__"
 
 S = TypeVar("S", bound=BaseSchema)
 
@@ -176,7 +177,7 @@ def scan_parquet(
     def _from_parquet(
         cls, source: str | Path | IO[bytes], scan: bool, **kwargs: Any
     ) -> FailureInfo[Schema]:
-        from .schema import deserialize_schema
+        from .schema import Schema, deserialize_schema
 
         metadata = pl.read_parquet_metadata(source)
         schema_metadata = metadata.get(SCHEMA_METADATA_KEY)
@@ -189,10 +190,13 @@ def _from_parquet(
             if scan
             else pl.read_parquet(source, **kwargs).lazy()
         )
+        failure_schema = deserialize_schema(schema_metadata, strict=False) or type(
+            UNKNOWN_SCHEMA_NAME, (Schema,), {}
+        )
         return FailureInfo(
             lf,
             json.loads(rule_metadata),
-            schema=deserialize_schema(schema_metadata),
+            schema=failure_schema,
         )
 
 

dataframely/schema.py

@@ -959,27 +959,33 @@ def read_parquet_metadata_schema(
     """
     metadata = pl.read_parquet_metadata(source)
     if (schema_metadata := metadata.get(SCHEMA_METADATA_KEY)) is not None:
-        try:
-            return deserialize_schema(schema_metadata)
-        except (JSONDecodeError, plexc.ComputeError):
-            return None
+        return deserialize_schema(schema_metadata, strict=False)
     return None
 
 
-def deserialize_schema(data: str) -> type[Schema]:
+@overload
+def deserialize_schema(data: str, strict: Literal[True] = True) -> type[Schema]: ...
+
+
+@overload
+def deserialize_schema(data: str, strict: Literal[False]) -> type[Schema] | None: ...
+
+
+def deserialize_schema(data: str, strict: bool = True) -> type[Schema] | None:
     """Deserialize a schema from a JSON string.
 
     This method allows to dynamically load a schema from its serialization, without
     having to know the schema to load in advance.
 
     Args:
         data: The JSON string created via :meth:`Schema.serialize`.
+        strict: Whether to raise an exception if the schema cannot be deserialized.
 
     Returns:
         The schema loaded from the JSON data.
 
     Raises:
-        ValueError: If the schema format version is not supported.
+        ValueError: If the schema format version is not supported and ``strict=True``.
 
     Attention:
         This functionality is considered unstable. It may be changed at any time
@@ -988,10 +994,15 @@ def deserialize_schema(data: str) -> type[Schema]:
     See also:
         :meth:`Schema.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}")
-    return _schema_from_dict(decoded)
+    try:
+        decoded = json.loads(data, cls=SchemaJSONDecoder)
+        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:
+        if strict:
+            raise e from e
+        return None
 
 
 def _schema_from_dict(data: dict[str, Any]) -> type[Schema]:

tests/test_failure_info.py

@@ -9,6 +9,8 @@
 from polars.testing import assert_frame_equal
 
 import dataframely as dy
+from dataframely._serialization import SCHEMA_METADATA_KEY
+from dataframely.failure import RULE_METADATA_KEY, UNKNOWN_SCHEMA_NAME, FailureInfo
 
 
 class MySchema(dy.Schema):
@@ -68,7 +70,9 @@ def test_write_parquet_custom_metadata(tmp_path: Path) -> None:
     "read_fn",
     [dy.FailureInfo.read_parquet, dy.FailureInfo.scan_parquet],
 )
-def test_missing_metadata(tmp_path: Path, read_fn: Callable[[Path], None]) -> None:
+def test_missing_metadata(
+    tmp_path: Path, read_fn: Callable[[Path], FailureInfo]
+) -> None:
     df = pl.DataFrame(
         {
             "a": [4, 5, 6, 6, 7, 8],
@@ -79,3 +83,27 @@ def test_missing_metadata(tmp_path: Path, read_fn: Callable[[Path], None]) -> No
 
     with pytest.raises(ValueError, match=r"does not contain the required metadata"):
         read_fn(tmp_path / "failure.parquet")
+
+
+@pytest.mark.parametrize(
+    "read_fn",
+    [dy.FailureInfo.read_parquet, dy.FailureInfo.scan_parquet],
+)
+def test_invalid_schema_deserialization(
+    tmp_path: Path, read_fn: Callable[[Path], FailureInfo]
+) -> None:
+    df = pl.DataFrame(
+        {
+            "a": [1, 2, 3],
+            "b": [False, True, False],
+        }
+    )
+    df.write_parquet(
+        tmp_path / "failure.parquet",
+        metadata={
+            SCHEMA_METADATA_KEY: "{WRONG",
+            RULE_METADATA_KEY: '["b"]',
+        },
+    )
+    failure = read_fn(tmp_path / "failure.parquet")
+    assert failure.schema.__name__ == UNKNOWN_SCHEMA_NAME