feat: Add DataFrame support in Collection

Author: gabriel

dataframely/collection/_base.py

@@ -15,6 +15,7 @@
 
 from dataframely._filter import Filter
 from dataframely._polars import FrameType
+from dataframely._typing import DataFrame as TypedDataFrame
 from dataframely._typing import LazyFrame as TypedLazyFrame
 from dataframely.exc import AnnotationImplementationError, ImplementationError
 from dataframely.schema import Schema
@@ -92,6 +93,8 @@ class MemberInfo(CollectionMember):
     schema: type[Schema]
     #: Whether the member is optional.
     is_optional: bool
+    #: Whether the member is a lazy frame.
+    is_lazy: bool = True
 
 
 @dataclass
@@ -249,31 +252,38 @@ def _derive_member_info(
                 raise AnnotationImplementationError(attr, type_annotation)
 
             not_none_args = [arg for arg in union_args if get_origin(arg) is not None]
-            if len(not_none_args) == 0 or not issubclass(
-                get_origin(not_none_args[0]), TypedLazyFrame
-            ):
+            if len(not_none_args) == 0:
                 raise AnnotationImplementationError(attr, type_annotation)
 
-            return MemberInfo(
-                schema=get_args(not_none_args[0])[0],
-                is_optional=True,
-                ignored_in_filters=collection_member.ignored_in_filters,
-                inline_for_sampling=collection_member.inline_for_sampling,
-                propagate_row_failures=collection_member.propagate_row_failures,
-            )
-        elif issubclass(origin, TypedLazyFrame):
-            # Happy path: required member
-            return MemberInfo(
-                schema=get_args(type_annotation)[0],
-                is_optional=False,
-                ignored_in_filters=collection_member.ignored_in_filters,
-                inline_for_sampling=collection_member.inline_for_sampling,
-                propagate_row_failures=collection_member.propagate_row_failures,
-            )
+            frame_origin = get_origin(not_none_args[0])
+            if frame_origin is None:
+                raise AnnotationImplementationError(attr, type_annotation)
+
+            schema = get_args(not_none_args[0])[0]
+            is_optional = True
+        elif issubclass(origin, (TypedLazyFrame, TypedDataFrame)):
+            frame_origin = origin
+            schema = get_args(type_annotation)[0]
+            is_optional = False
         else:
-            # Some other unknown annotation
             raise AnnotationImplementationError(attr, type_annotation)
 
+        if issubclass(frame_origin, TypedLazyFrame):
+            is_lazy = True
+        elif issubclass(frame_origin, TypedDataFrame):
+            is_lazy = False
+        else:
+            raise AnnotationImplementationError(attr, type_annotation)
+
+        return MemberInfo(
+            schema=schema,
+            is_optional=is_optional,
+            is_lazy=is_lazy,
+            ignored_in_filters=collection_member.ignored_in_filters,
+            inline_for_sampling=collection_member.inline_for_sampling,
+            propagate_row_failures=collection_member.propagate_row_failures,
+        )
+
     def __repr__(cls) -> str:
         parts = [f'[Collection "{cls.__class__.__name__}"]']
         parts.append(textwrap.indent("Members:", prefix=" " * 2))
@@ -344,6 +354,16 @@ def non_ignored_members(cls) -> set[str]:
             if not member.ignored_in_filters
         }
 
+    @classmethod
+    def lazy_members(cls) -> set[str]:
+        """The names of all members annotated as lazy frames."""
+        return {name for name, member in cls.members().items() if member.is_lazy}
+
+    @classmethod
+    def eager_members(cls) -> set[str]:
+        """The names of all members annotated as data frames (eager)."""
+        return {name for name, member in cls.members().items() if not member.is_lazy}
+
     @classmethod
     def _failure_propagating_members(cls) -> set[str]:
         """The names of all members of the collection that propagate individual row
@@ -371,20 +391,40 @@ def common_primary_key(cls) -> list[str]:
     def _filters(cls) -> dict[str, Filter[Self]]:
         return getattr(cls, _FILTER_ATTR)
 
-    def to_dict(self) -> dict[str, pl.LazyFrame]:
-        """Return a dictionary representation of this collection."""
+    def to_dict(self) -> dict[str, FrameType]:
+        """Return a dictionary representation of this collection.
+
+        Returns:
+            A dictionary mapping member names to their frames.
+            Members annotated with :class:`~dataframely.DataFrame` return DataFrames,
+            while members annotated with :class:`~dataframely.LazyFrame` return LazyFrames.
+        """
         return {
             member: getattr(self, member)
             for member in self.member_schemas()
             if getattr(self, member) is not None
         }
 
+    def _to_lazy_dict(self) -> dict[str, pl.LazyFrame]:
+        """Return a dictionary with all members as lazy frames (internal use)."""
+        return {
+            member: getattr(self, member).lazy()
+            for member in self.member_schemas()
+            if getattr(self, member) is not None
+        }
+
     @classmethod
     def _init(cls, data: Mapping[str, FrameType], /) -> Self:
         out = cls()
         for member_name, member in cls.members().items():
             if member.is_optional and member_name not in data:
                 setattr(out, member_name, None)
-            else:
+            elif member.is_lazy:
                 setattr(out, member_name, data[member_name].lazy())
+            else:
+                frame = data[member_name]
+                if isinstance(frame, pl.LazyFrame):
+                    setattr(out, member_name, frame.collect())
+                else:
+                    setattr(out, member_name, frame)
         return out

dataframely/collection/collection.py

@@ -735,7 +735,7 @@ def join(
                     how=how,
                     maintain_order=maintain_order,
                 )
-                for key, lf in self.to_dict().items()
+                for key, lf in self._to_lazy_dict().items()
             }
         )
 
@@ -795,9 +795,10 @@ def collect_all(self) -> Self:
             collection's members are still "lazy". However, they are "shallow-lazy",
             meaning they are obtained by calling `.collect().lazy()`.
         """
-        dfs = pl.collect_all(self.to_dict().values())
+        lazy_dict = self._to_lazy_dict()
+        dfs = pl.collect_all(lazy_dict.values())
         return self._init(
-            {key: dfs[i].lazy() for i, key in enumerate(self.to_dict().keys())}
+            {key: dfs[i].lazy() for i, key in enumerate(lazy_dict.keys())}
         )
 
     # --------------------------------- SERIALIZATION -------------------------------- #
@@ -1172,7 +1173,7 @@ def _write(self, backend: StorageBackend, **kwargs: Any) -> None:
         # Utility method encapsulating the interaction with the StorageBackend
 
         backend.write_collection(
-            self.to_dict(),
+            self._to_lazy_dict(),
             serialized_collection=self.serialize(),
             serialized_schemas={
                 key: schema.serialize() for key, schema in self.member_schemas().items()
@@ -1184,7 +1185,7 @@ def _sink(self, backend: StorageBackend, **kwargs: Any) -> None:
         # Utility method encapsulating the interaction with the StorageBackend
 
         backend.sink_collection(
-            self.to_dict(),
+            self._to_lazy_dict(),
             serialized_collection=self.serialize(),
             serialized_schemas={
                 key: schema.serialize() for key, schema in self.member_schemas().items()

dataframely/collection/filter_result.py

@@ -47,7 +47,7 @@ def collect_all(self, **kwargs: Any) -> CollectionFilterResult[C]:
             Until https://github.com/pola-rs/polars/pull/24129 is released, the
             performance advantage of this method is limited.
         """
-        members = self.result.to_dict()
+        members = self.result._to_lazy_dict()
         collected = pl.collect_all(
             itertools.chain(
                 members.values(),

tests/collection/test_dataframe_members.py

@@ -0,0 +1,151 @@
+# Copyright (c) QuantCo 2025-2026
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for dy.DataFrame members in collections.
+
+Members annotated with dy.DataFrame are collected once during _init and stored
+as DataFrames, while dy.LazyFrame members remain lazy.
+"""
+
+import polars as pl
+import pytest
+
+import dataframely as dy
+
+# ------------------------------------------------------------------------------------ #
+#                                        SCHEMA                                        #
+# ------------------------------------------------------------------------------------ #
+
+
+class UserSchema(dy.Schema):
+    id = dy.Integer(primary_key=True)
+    name = dy.String()
+
+
+class OrderSchema(dy.Schema):
+    id = dy.Integer(primary_key=True)
+    user_id = dy.Integer()
+    amount = dy.Float(min=0)
+
+
+class EagerCollection(dy.Collection):
+    """Collection with only DataFrame (eager) members."""
+
+    users: dy.DataFrame[UserSchema]
+    orders: dy.DataFrame[OrderSchema]
+
+
+class MixedCollection(dy.Collection):
+    """Collection with mixed DataFrame and LazyFrame members."""
+
+    users: dy.DataFrame[UserSchema]
+    orders: dy.LazyFrame[OrderSchema]
+
+
+class LazyCollection(dy.Collection):
+    """Collection with only LazyFrame members (traditional)."""
+
+    users: dy.LazyFrame[UserSchema]
+    orders: dy.LazyFrame[OrderSchema]
+
+
+class OptionalEagerCollection(dy.Collection):
+    """Collection with optional DataFrame member."""
+
+    users: dy.DataFrame[UserSchema]
+    orders: dy.DataFrame[OrderSchema] | None
+
+
+# ------------------------------------------------------------------------------------ #
+#                                       FIXTURES                                       #
+# ------------------------------------------------------------------------------------ #
+
+
+@pytest.fixture()
+def valid_data() -> dict[str, pl.DataFrame]:
+    return {
+        "users": pl.DataFrame({"id": [1, 2], "name": ["Alice", "Bob"]}),
+        "orders": pl.DataFrame(
+            {"id": [1, 2], "user_id": [1, 2], "amount": [10.0, 20.0]}
+        ),
+    }
+
+
+# ------------------------------------------------------------------------------------ #
+#                                    MEMBER INFO TESTS                                 #
+# ------------------------------------------------------------------------------------ #
+
+
+def test_eager_member_detection() -> None:
+    members = EagerCollection.members()
+    assert not members["users"].is_lazy
+    assert not members["orders"].is_lazy
+
+
+def test_lazy_member_detection() -> None:
+    members = LazyCollection.members()
+    assert members["users"].is_lazy
+    assert members["orders"].is_lazy
+
+
+def test_mixed_member_detection() -> None:
+    members = MixedCollection.members()
+    assert not members["users"].is_lazy
+    assert members["orders"].is_lazy
+
+
+def test_optional_eager_member_detection() -> None:
+    members = OptionalEagerCollection.members()
+    assert not members["users"].is_lazy
+    assert not members["orders"].is_lazy
+    assert not members["users"].is_optional
+    assert members["orders"].is_optional
+
+
+def test_lazy_members_helper() -> None:
+    assert EagerCollection.lazy_members() == set()
+    assert LazyCollection.lazy_members() == {"users", "orders"}
+    assert MixedCollection.lazy_members() == {"orders"}
+
+
+def test_eager_members_helper() -> None:
+    assert EagerCollection.eager_members() == {"users", "orders"}
+    assert LazyCollection.eager_members() == set()
+    assert MixedCollection.eager_members() == {"users"}
+
+
+# ------------------------------------------------------------------------------------ #
+#                                  ACCESS PATTERN TESTS                                #
+# ------------------------------------------------------------------------------------ #
+
+
+def test_eager_member_returns_dataframe(valid_data: dict[str, pl.DataFrame]) -> None:
+    collection = EagerCollection.validate(valid_data)
+    assert isinstance(collection.users, pl.DataFrame)
+    assert isinstance(collection.orders, pl.DataFrame)
+
+
+def test_lazy_member_returns_lazyframe(valid_data: dict[str, pl.DataFrame]) -> None:
+    collection = LazyCollection.validate(valid_data)
+    assert isinstance(collection.users, pl.LazyFrame)
+    assert isinstance(collection.orders, pl.LazyFrame)
+
+
+def test_mixed_collection_returns_correct_types(
+    valid_data: dict[str, pl.DataFrame],
+) -> None:
+    collection = MixedCollection.validate(valid_data)
+    assert isinstance(collection.users, pl.DataFrame)
+    assert isinstance(collection.orders, pl.LazyFrame)
+
+
+def test_to_dict_returns_correct_types(valid_data: dict[str, pl.DataFrame]) -> None:
+    eager = EagerCollection.validate(valid_data)
+    result = eager.to_dict()
+    assert isinstance(result["users"], pl.DataFrame)
+    assert isinstance(result["orders"], pl.DataFrame)
+
+    mixed = MixedCollection.validate(valid_data)
+    result = mixed.to_dict()
+    assert isinstance(result["users"], pl.DataFrame)
+    assert isinstance(result["orders"], pl.LazyFrame)

tests/collection/test_implementation.py

@@ -16,17 +16,17 @@ class MyTestSchema(dy.Schema):
     a = dy.Integer(primary_key=True)
 
 
-def test_annotation_type_failure() -> None:
-    with pytest.raises(
-        AnnotationImplementationError,
-    ):
-        create_collection(
-            "test",
-            {
-                "first": create_schema("first", {"a": dy.Integer()}),
-            },
-            annotation_base_class=dy.DataFrame,
-        )
+def test_annotation_dataframe_success() -> None:
+    """DataFrame annotations are now supported."""
+    collection = create_collection(
+        "test",
+        {
+            "first": create_schema("first", {"a": dy.Integer()}),
+        },
+        annotation_base_class=dy.DataFrame,
+    )
+    members = collection.members()
+    assert not members["first"].is_lazy
 
 
 def test_annotation_union_success() -> None:
@@ -40,14 +40,16 @@ def test_annotation_union_success() -> None:
 
 
 def test_annotation_union_with_data_frame() -> None:
-    """When we use a union annotation, it must contain one typed LazyFrame and None."""
-    with pytest.raises(AnnotationImplementationError):
-        create_collection_raw(
-            "test",
-            {
-                "first": dy.DataFrame[MyTestSchema] | None,
-            },
-        )
+    """DataFrame union with None is now supported for optional eager members."""
+    collection = create_collection_raw(
+        "test",
+        {
+            "first": dy.DataFrame[MyTestSchema] | None,
+        },
+    )
+    members = collection.members()
+    assert not members["first"].is_lazy
+    assert members["first"].is_optional
 
 
 def test_annotation_union_too_many_arg_failure() -> None: