Add serialisation

Author: gabriel

dataframely/collection/collection.py

@@ -32,7 +32,7 @@
 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._typing import DataFrame, LazyFrame, Validation
 from dataframely.exc import (
     DeserializationError,
     ValidationError,
@@ -841,6 +841,7 @@ def serialize(cls) -> str:
                 name: {
                     "schema": info.schema._as_dict(),
                     "is_optional": info.is_optional,
+                    "is_lazy": info.is_lazy,
                     "ignored_in_filters": info.ignored_in_filters,
                     "inline_for_sampling": info.inline_for_sampling,
                 }
@@ -1329,11 +1330,14 @@ def deserialize_collection(data: str, strict: bool = True) -> type[Collection] |
 
         annotations: dict[str, Any] = {}
         for name, info in decoded["members"].items():
-            lf_type = LazyFrame[_schema_from_dict(info["schema"])]  # type: ignore
+            schema = _schema_from_dict(info["schema"])
+            # Default to lazy for backwards compatibility with old serialized data
+            is_lazy = info.get("is_lazy", True)
+            frame_type = LazyFrame[schema] if is_lazy else DataFrame[schema]  # type: ignore
             if info["is_optional"]:
-                lf_type = lf_type | None  # type: ignore
+                frame_type = frame_type | None  # type: ignore
             annotations[name] = Annotated[
-                lf_type,
+                frame_type,
                 CollectionMember(
                     ignored_in_filters=info["ignored_in_filters"],
                     inline_for_sampling=info["inline_for_sampling"],

tests/collection/test_serialization.py

@@ -84,6 +84,55 @@ def test_roundtrip_matches(collection: type[dy.Collection]) -> None:
     assert collection.matches(decoded)
 
 
+# --------------------------------- DATAFRAME MEMBERS -------------------------------- #
+
+
+class EagerSchema(dy.Schema):
+    id = dy.Int64(primary_key=True)
+
+
+class MixedEagerCollection(dy.Collection):
+    """Collection with mixed DataFrame and LazyFrame members."""
+
+    eager: dy.DataFrame[EagerSchema]
+    lazy: dy.LazyFrame[EagerSchema]
+
+
+def test_serialize_includes_is_lazy() -> None:
+    """Serialization includes the is_lazy field for each member."""
+    serialized = MixedEagerCollection.serialize()
+    decoded = json.loads(serialized)
+
+    assert decoded["members"]["eager"]["is_lazy"] is False
+    assert decoded["members"]["lazy"]["is_lazy"] is True
+
+
+def test_roundtrip_dataframe_members() -> None:
+    """DataFrame members round-trip correctly through serialization."""
+    serialized = MixedEagerCollection.serialize()
+    decoded = dy.deserialize_collection(serialized)
+
+    assert MixedEagerCollection.matches(decoded)
+    assert not decoded.members()["eager"].is_lazy
+    assert decoded.members()["lazy"].is_lazy
+
+
+def test_deserialize_without_is_lazy_defaults_to_lazy() -> None:
+    """Old serialized data without is_lazy defaults to lazy for backwards compat."""
+    collection = create_collection(
+        "test", {"s1": create_schema("schema1", {"a": dy.Int64()})}
+    )
+    serialized = collection.serialize()
+
+    # Remove is_lazy from serialized data to simulate old format
+    decoded_dict = json.loads(serialized)
+    del decoded_dict["members"]["s1"]["is_lazy"]
+    modified = json.dumps(decoded_dict)
+
+    result = dy.deserialize_collection(modified)
+    assert result.members()["s1"].is_lazy is True
+
+
 # ----------------------------- DESERIALIZATION FAILURES ----------------------------- #