fix: Be more lenient when reading metadata (#114)

Author: borchero

dataframely/_serialization.py

@@ -46,7 +46,7 @@ def default(self, obj: Any) -> Any:
             case pl.Expr():
                 return {
                     "__type__": "expression",
-                    "value": obj.meta.serialize(format="json"),
+                    "value": base64.b64encode(obj.meta.serialize()).decode("utf-8"),
                 }
             case pl.LazyFrame():
                 return {
@@ -87,8 +87,14 @@ def object_hook(self, dct: dict[str, Any]) -> Any:
             case "tuple":
                 return tuple(dct["value"])
             case "expression":
-                data = BytesIO(cast(str, dct["value"]).encode("utf-8"))
-                return pl.Expr.deserialize(data, format="json")
+                value_str = cast(str, dct["value"]).encode("utf-8")
+                if value_str.startswith(b"{"):
+                    # NOTE: This branch is for backwards-compatibility only
+                    data = BytesIO(value_str)
+                    return pl.Expr.deserialize(data, format="json")
+                else:
+                    data = BytesIO(base64.b64decode(value_str))
+                    return pl.Expr.deserialize(data)
             case "lazyframe":
                 data = BytesIO(
                     base64.b64decode(cast(str, dct["value"]).encode("utf-8"))

dataframely/collection.py

@@ -870,7 +870,7 @@ def _from_parquet(
         if (collection_type is None) and (schema_file := path / "schema.json").exists():
             try:
                 collection_type = deserialize_collection(schema_file.read_text())
-            except JSONDecodeError:
+            except (JSONDecodeError, plexc.ComputeError):
                 pass
 
         return data, collection_type
@@ -947,12 +947,15 @@ def read_parquet_metadata_collection(
         source: Path to a parquet file or a file-like object that contains the metadata.
 
     Returns:
-        The collection that was serialized to the metadata or ``None`` if no collection metadata
-        is found.
+        The collection that was serialized to the metadata. ``None`` if no collection
+        metadata is found or the deserialization fails.
     """
     metadata = pl.read_parquet_metadata(source)
     if (schema_metadata := metadata.get(COLLECTION_METADATA_KEY)) is not None:
-        return deserialize_collection(schema_metadata)
+        try:
+            return deserialize_collection(schema_metadata)
+        except (JSONDecodeError, plexc.ComputeError):
+            return None
     return None
 
 

dataframely/schema.py

@@ -8,6 +8,7 @@
 import warnings
 from abc import ABC
 from collections.abc import Iterable, Mapping, Sequence
+from json import JSONDecodeError
 from pathlib import Path
 from typing import IO, Any, Literal, overload
 
@@ -970,12 +971,15 @@ def read_parquet_metadata_schema(
         source: Path to a parquet file or a file-like object that contains the metadata.
 
     Returns:
-        The schema that was serialized to the metadata or ``None`` if no schema metadata
-        is found.
+        The schema that was serialized to the metadata. ``None`` if no schema metadata
+        is found or the deserialization fails.
     """
     metadata = pl.read_parquet_metadata(source)
     if (schema_metadata := metadata.get(SCHEMA_METADATA_KEY)) is not None:
-        return deserialize_schema(schema_metadata)
+        try:
+            return deserialize_schema(schema_metadata)
+        except (JSONDecodeError, plexc.ComputeError):
+            return None
     return None
 
 

tests/collection/test_read_write_parquet.py

@@ -12,6 +12,7 @@
 from polars.testing import assert_frame_equal
 
 import dataframely as dy
+from dataframely._serialization import COLLECTION_METADATA_KEY
 from dataframely.collection import _reconcile_collection_types
 from dataframely.exc import ValidationRequiredError
 from dataframely.testing import create_collection, create_schema
@@ -378,3 +379,20 @@ def test_reconcile_collection_types(
     inputs: list[type[dy.Collection] | None], output: type[dy.Collection] | None
 ) -> None:
     assert output == _reconcile_collection_types(inputs)
+
+
+# ---------------------------------- MANUAL METADATA --------------------------------- #
+
+
+def test_read_invalid_parquet_metadata_collection(tmp_path: Path) -> None:
+    # Arrange
+    df = pl.DataFrame({"a": [1, 2, 3]})
+    df.write_parquet(
+        tmp_path / "df.parquet", metadata={COLLECTION_METADATA_KEY: "invalid"}
+    )
+
+    # Act
+    collection = dy.read_parquet_metadata_collection(tmp_path / "df.parquet")
+
+    # Assert
+    assert collection is None

tests/schema/test_read_write_parquet.py

@@ -10,6 +10,7 @@
 from polars.testing import assert_frame_equal
 
 import dataframely as dy
+from dataframely._serialization import SCHEMA_METADATA_KEY
 from dataframely.exc import ValidationRequiredError
 from dataframely.testing import create_schema
 
@@ -216,3 +217,18 @@ def test_read_write_parquet_validation_skip_invalid_schema(
 
     # Assert
     spy.assert_not_called()
+
+
+# ---------------------------------- MANUAL METADATA --------------------------------- #
+
+
+def test_read_invalid_parquet_metadata_schema(tmp_path: Path) -> None:
+    # Arrange
+    df = pl.DataFrame({"a": [1, 2, 3]})
+    df.write_parquet(tmp_path / "df.parquet", metadata={SCHEMA_METADATA_KEY: "invalid"})
+
+    # Act
+    schema = dy.read_parquet_metadata_schema(tmp_path / "df.parquet")
+
+    # Assert
+    assert schema is None

tests/test_serialization.py

@@ -0,0 +1,22 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+import json
+
+import polars as pl
+
+from dataframely._serialization import SchemaJSONDecoder
+
+
+def test_decode_json_expression() -> None:
+    # Arrange
+    expr = pl.col("a") + 1
+    encoded = json.dumps(
+        {"__type__": "expression", "value": expr.meta.serialize(format="json")}
+    )
+
+    # Act
+    decoded = json.loads(encoded, cls=SchemaJSONDecoder)
+
+    # Assert
+    assert expr.meta.eq(decoded)