dataframely 2.11.0__tar.gz → 2.12.0__tar.gz

{dataframely-2.11.0 → dataframely-2.12.0}/PKG-INFO

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

{dataframely-2.11.0 → dataframely-2.12.0}/dataframely/_polars.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import datetime as dt
-from typing import TypeVar
+from typing import Any, TypeVar
 
 import polars as pl
 from polars.datatypes import DataTypeClass
@@ -41,9 +41,18 @@ def timedelta_matches_resolution(d: dt.timedelta, resolution: str) -> bool:
     return datetime_matches_resolution(EPOCH_DATETIME + d, resolution)
 
 
-def collect_if(lf: pl.LazyFrame, condition: bool) -> pl.DataFrame | pl.LazyFrame:
-    """Collect a lazy frame if the original was eager, otherwise return the lazy
-    frame."""
+def collect_if(lf: pl.LazyFrame, condition: bool, **kwargs: Any) -> pl.LazyFrame:
+    """Collect a lazy frame based on `condition`."""
     if condition:
-        return lf.collect()
+        return lf.collect(**kwargs).lazy()
     return lf
+
+
+def collect_all_if(
+    lfs: dict[str, pl.LazyFrame], condition: bool, **kwargs: Any
+) -> dict[str, pl.LazyFrame]:
+    """Collect the lazy frames in the dictionary based on `condition`."""
+    if condition:
+        dfs = pl.collect_all(lfs.values(), **kwargs)
+        return {k: v.lazy() for k, v in zip(lfs.keys(), dfs)}
+    return lfs

{dataframely-2.11.0 → dataframely-2.12.0}/dataframely/collection/collection.py

@@ -8,11 +8,21 @@ import sys
 import textwrap
 import warnings
 from abc import ABC
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from dataclasses import asdict
 from json import JSONDecodeError
 from pathlib import Path
-from typing import IO, Annotated, Any, Literal, cast, overload
+from typing import (
+    IO,
+    Annotated,
+    Any,
+    Concatenate,
+    Literal,
+    ParamSpec,
+    TypeVar,
+    cast,
+    overload,
+)
 
 import polars as pl
 import polars.exceptions as plexc
@@ -21,7 +31,7 @@ from dataframely._compat import deltalake
 from dataframely._filter import Filter
 from dataframely._native import format_rule_failures
 from dataframely._plugin import all_rules_required
-from dataframely._polars import FrameType, collect_if
+from dataframely._polars import FrameType, collect_all_if
 from dataframely._serialization import (
     SERIALIZATION_FORMAT_VERSION,
     SchemaJSONDecoder,
@@ -53,6 +63,9 @@ else:
 
 _FILTER_COLUMN_PREFIX = "__DATAFRAMELY_FILTER_COLUMN__"
 
+P = ParamSpec("P")
+T = TypeVar("T")
+
 
 class Collection(BaseCollection, ABC):
     """Base class for all collections of data frames with a predefined schema.
@@ -367,7 +380,13 @@ class Collection(BaseCollection, ABC):
 
     @classmethod
     def validate(
-        cls, data: Mapping[str, FrameType], /, *, cast: bool = False, eager: bool = True
+        cls,
+        data: Mapping[str, FrameType],
+        /,
+        *,
+        cast: bool = False,
+        eager: bool = True,
+        **kwargs: Any,
     ) -> Self:
         """Validate that a set of data frames satisfy the collection's invariants.
 
@@ -387,6 +406,8 @@ class Collection(BaseCollection, ABC):
                 :meth:`~polars.LazyFrame.collect` on the individual member or
                 :meth:`collect_all` on the collection. Note that, in the latter case,
                 information from error messages is limited.
+            kwargs: Keyword arguments passed directly to :meth:`polars.collect_all` and
+                :meth:`polars.LazyFrame.collect` when `eager=True`.
 
         Raises:
             ValueError: If an insufficient set of input data frames is provided, i.e. if
@@ -408,7 +429,7 @@ class Collection(BaseCollection, ABC):
         if eager:
             # If we perform the validation eagerly, we call filter and check the failure
             # information to properly construct a useful error message.
-            filtered, failures = cls.filter(data, cast=cast, eager=True)
+            filtered, failures = cls.filter(data, cast=cast, eager=True, **kwargs)
             if any(len(failure) > 0 for failure in failures.values()):
                 errors: dict[str, str] = {}
                 for member, failure in failures.items():
@@ -476,7 +497,9 @@ class Collection(BaseCollection, ABC):
             return cls._init(members)
 
     @classmethod
-    def is_valid(cls, data: Mapping[str, FrameType], /, *, cast: bool = False) -> bool:
+    def is_valid(
+        cls, data: Mapping[str, FrameType], /, *, cast: bool = False, **kwargs: Any
+    ) -> bool:
         """Utility method to check whether :meth:`validate` raises an exception.
 
         Args:
@@ -485,6 +508,8 @@ class Collection(BaseCollection, ABC):
                 the member as key.
             cast: Whether columns with a wrong data type in the member data frame are
                 cast to their schemas' defined data types if possible.
+            kwargs: Keyword arguments passed directly to :meth:`polars.collect_all` and
+                :meth:`polars.LazyFrame.collect`.
 
         Returns:
             Whether the provided members satisfy the invariants of the collection.
@@ -499,7 +524,7 @@ class Collection(BaseCollection, ABC):
         members: dict[str, pl.LazyFrame] = {}
         for member, schema in cls.member_schemas().items():
             if member in data:
-                if not schema.is_valid(data[member], cast=cast):
+                if not schema.is_valid(data[member], cast=cast, **kwargs):
                     return False
                 members[member] = data[member].lazy()
 
@@ -510,9 +535,12 @@ class Collection(BaseCollection, ABC):
             keep = [filter.logic(result_cls).select(primary_key) for filter in filters]
             joined = _join_all(*keep, on=primary_key, how="inner")
             removed_rows = pl.collect_all(
-                data[member].lazy().join(joined, on=primary_key, how="anti")
-                for member in cls.members()
-                if member in data
+                (
+                    data[member].lazy().join(joined, on=primary_key, how="anti")
+                    for member in cls.members()
+                    if member in data
+                ),
+                **kwargs,
             )
             return all(df.is_empty() for df in removed_rows)
 
@@ -522,7 +550,13 @@ class Collection(BaseCollection, ABC):
 
     @classmethod
     def filter(
-        cls, data: Mapping[str, FrameType], /, *, cast: bool = False, eager: bool = True
+        cls,
+        data: Mapping[str, FrameType],
+        /,
+        *,
+        cast: bool = False,
+        eager: bool = True,
+        **kwargs: Any,
     ) -> CollectionFilterResult[Self]:
         """Filter the members data frame by their schemas and the collection's filters.
 
@@ -538,6 +572,8 @@ class Collection(BaseCollection, ABC):
             eager: Whether the filter operation should be performed eagerly.
                 Note that until https://github.com/pola-rs/polars/pull/24129 is
                 released, eagerly filtering can provide significant speedups.
+            kwargs: Keyword arguments passed directly to :meth:`polars.collect_all` and
+                :meth:`polars.LazyFrame.collect` when `eager=True`.
 
         Returns:
             A named tuple with fields `result` and `failure`. The `result` field
@@ -580,7 +616,7 @@ class Collection(BaseCollection, ABC):
                 continue
 
             member_result, failures[member_name] = member.schema.filter(
-                data[member_name].lazy(), cast=cast, eager=eager
+                data[member_name].lazy(), cast=cast, eager=eager, **kwargs
             )
             results[member_name] = member_result.lazy()
 
@@ -592,28 +628,25 @@ class Collection(BaseCollection, ABC):
             result_cls = cls._init(results)
             primary_key = cls.common_primary_key()
 
-            keep: dict[str, pl.LazyFrame] = {}
-            for name, filter in filters.items():
-                keep[name] = (
-                    filter.logic(result_cls)
-                    .select(primary_key)
-                    .pipe(collect_if, eager)
-                    .lazy()
-                )
+            keep = {
+                name: filter.logic(result_cls).select(primary_key)
+                for name, filter in filters.items()
+            }
+            keep = collect_all_if(keep, eager, **kwargs)
 
-            drop: dict[str, pl.LazyFrame] = {}
-            for failure_propagating_member in failure_propagating_members:
-                annotation_column = f"{failure_propagating_member}|failure_propagation"
-                drop[annotation_column] = (
+            drop: dict[str, pl.LazyFrame] = {
+                f"{failure_propagating_member}|failure_propagation": (
                     failures[failure_propagating_member]
                     ._lf.select(primary_key)
                     .unique()
-                    .pipe(collect_if, eager)
-                    .lazy()
                 )
+                for failure_propagating_member in failure_propagating_members
+            }
+            drop = collect_all_if(drop, eager, **kwargs)
 
             # Now we can iterate over the results and left-join onto each individual
-            # filter to obtain independent boolean indicators of whether to keep the row
+            # filter to obtain independent boolean indicators of whether to keep the row.
+            lfs_with_eval: dict[str, pl.LazyFrame] = {}
             for member_name, filtered in results.items():
                 member_info = cls.members()[member_name]
                 if member_info.ignored_in_filters:
@@ -635,7 +668,11 @@ class Collection(BaseCollection, ABC):
                         maintain_order="left",
                     ).with_columns(pl.col(name).fill_null(True))
 
-                lf_with_eval = lf_with_eval.pipe(collect_if, eager).lazy()
+                lfs_with_eval[member_name] = lf_with_eval
+
+            lfs_with_eval = collect_all_if(lfs_with_eval, eager, **kwargs)
+            for member_name, lf_with_eval in lfs_with_eval.items():
+                member_info = cls.members()[member_name]
 
                 # Filtering `lf_with_eval` by the rows for which all joins
                 # "succeeded", we can identify the rows that pass all the filters. We
@@ -700,7 +737,7 @@ class Collection(BaseCollection, ABC):
 
         result = CollectionFilterResult(cls._init(results), failures)
         if eager:
-            return result.collect_all()
+            return result.collect_all(**kwargs)
         return result
 
     def join(
@@ -795,22 +832,52 @@ class Collection(BaseCollection, ABC):
 
     # ---------------------------------- COLLECTION ---------------------------------- #
 
-    def collect_all(self) -> Self:
+    def collect_all(self, **kwargs: Any) -> Self:
         """Collect all members of the collection.
 
         This method collects all members in parallel for maximum efficiency. It is
         particularly useful when :meth:`filter` is called with lazy frame inputs.
 
+        Args:
+            kwargs: Keyword arguments passed directly to :meth:`polars.collect_all`.
+
         Returns:
             The same collection with all members collected once. Members annotated
             with :class:`~dataframely.DataFrame` are returned as DataFrames, while
             members annotated with :class:`~dataframely.LazyFrame` are returned as
-            "shallow-lazy" frames (obtained by calling ``.collect().lazy()``).
+            "shallow-lazy" frames (obtained by calling `.collect().lazy()`).
         """
         lazy_dict = self.to_dict()
-        dfs = pl.collect_all(lazy_dict.values())
+        dfs = pl.collect_all(lazy_dict.values(), **kwargs)
         return self._init(dict(zip(lazy_dict, dfs)))
 
+    def pipe(
+        self,
+        function: Callable[Concatenate[Self, P], T],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> T:
+        """Apply a function to this collection.
+
+        This method allows chaining operations on a collection in a fluent style,
+        analogously to :meth:`polars.LazyFrame.pipe`.
+
+        Args:
+            function: The callable to apply. It receives this collection as its first
+                argument, followed by any additional ``args`` and ``kwargs``.
+            args: Additional positional arguments to pass to ``function``.
+            kwargs: Additional keyword arguments to pass to ``function``.
+
+        Returns:
+            The return value of ``function`` when called as described.
+
+        Example:
+            >>> def add_prefix(collection: MyCollection, prefix: str) -> MyCollection:
+            ...     ...
+            >>> result = my_collection.pipe(add_prefix, prefix="foo")
+        """
+        return function(self, *args, **kwargs)
+
     # --------------------------------- SERIALIZATION -------------------------------- #
 
     @classmethod

{dataframely-2.11.0 → dataframely-2.12.0}/dataframely/functional.py

@@ -84,9 +84,9 @@ def require_relationship_one_to_at_least_one(
         columns, filtered to ensure a 1:{1,N} relationship.
     """
     if drop_duplicates:
-        return lhs.unique(on, keep="none").join(rhs.unique(on), on=on)
+        return lhs.unique(on, keep="none").join(rhs, on=on, how="semi")
 
-    return lhs.join(rhs.unique(on), on=on)
+    return lhs.join(rhs, on=on, how="semi")
 
 
 # ------------------------------------------------------------------------------------ #

{dataframely-2.11.0 → dataframely-2.12.0}/dataframely/schema.py

@@ -501,6 +501,7 @@ class Schema(BaseSchema, ABC):
         *,
         cast: bool = False,
         eager: Literal[True] = True,
+        **kwargs: Any,
     ) -> DataFrame[Self]: ...
 
     @overload
@@ -512,6 +513,7 @@ class Schema(BaseSchema, ABC):
         *,
         cast: bool = False,
         eager: Literal[False],
+        **kwargs: Any,
     ) -> LazyFrame[Self]: ...
 
     @overload
@@ -523,6 +525,7 @@ class Schema(BaseSchema, ABC):
         *,
         cast: bool = False,
         eager: bool,
+        **kwargs: Any,
     ) -> DataFrame[Self] | LazyFrame[Self]: ...
 
     @classmethod
@@ -533,6 +536,7 @@ class Schema(BaseSchema, ABC):
         *,
         cast: bool = False,
         eager: bool = True,
+        **kwargs: Any,
     ) -> DataFrame[Self] | LazyFrame[Self]:
         """Validate that a data frame satisfies the schema.
 
@@ -554,6 +558,8 @@ class Schema(BaseSchema, ABC):
                     not surface *all* validation issues as the validation is aborted
                     once the first failure is encountered. Likewise, the reported
                     validation failure can be non-deterministic.
+            kwargs: Keyword arguments passed directly to :meth:`polars.LazyFrame.collect`
+                when `eager=True`.
 
         Returns:
             The input eager or lazy frame, wrapped in a generic version of the
@@ -574,7 +580,7 @@ class Schema(BaseSchema, ABC):
                 `eager=False`.
         """
         if eager:
-            out, failure = cls.filter(df, cast=cast, eager=True)
+            out, failure = cls.filter(df, cast=cast, eager=True, **kwargs)
             if len(failure) > 0:
                 counts = failure.counts()
                 raise ValidationError(
@@ -608,7 +614,7 @@ class Schema(BaseSchema, ABC):
 
     @classmethod
     def is_valid(
-        cls, df: pl.DataFrame | pl.LazyFrame, /, *, cast: bool = False
+        cls, df: pl.DataFrame | pl.LazyFrame, /, *, cast: bool = False, **kwargs: Any
     ) -> bool:
         """Check whether a data frame satisfies the schema.
 
@@ -627,6 +633,7 @@ class Schema(BaseSchema, ABC):
                 cast to the schema's defined data type before running validation. If set
                 to `False`, a wrong data type will result in a return value of
                 `False`.
+            kwargs: Keyword arguments passed directly to :meth:`polars.LazyFrame.collect`.
 
         Returns:
             Whether the provided dataframe can be validated with this schema.
@@ -646,12 +653,12 @@ class Schema(BaseSchema, ABC):
                 return (
                     lf.pipe(with_evaluation_rules, rules)
                     .select(all_rules(rules.keys()))
-                    .collect()
+                    .collect(**kwargs)
                     .item()
                 )
             # NOTE: We cannot simply return `True` here as, otherwise, we wouldn't
             #  validate the schema.
-            return lf.select(pl.lit(True)).collect().item()
+            return lf.select(pl.lit(True)).collect(**kwargs).item()
         except SchemaError:
             # If we encounter a schema error, we gracefully handle this as 'invalid'
             return False
@@ -699,6 +706,7 @@ class Schema(BaseSchema, ABC):
         *,
         cast: bool = False,
         eager: bool = True,
+        **kwargs: Any,
     ) -> FilterResult[Self] | LazyFilterResult[Self]:
         """Filter the data frame by the rules of this schema, returning `(valid,
         failures)`.
@@ -709,16 +717,16 @@ class Schema(BaseSchema, ABC):
         succeeds.
 
         Args:
-            df: The data frame to filter for valid rows.
-                The data frame is collected within this method, regardless of whether
-                a :class:`~polars.DataFrame` or :class:`~polars.LazyFrame` is passed.
-            cast:
-                Whether columns with a wrong data type in the input data frame are
+            df: The data frame to filter for valid rows. The data frame is collected
+                within this method, regardless of whether a :class:`~polars.DataFrame`
+                or :class:`~polars.LazyFrame` is passed.
+            cast: Whether columns with a wrong data type in the input data frame are
                 cast to the schema's defined data type if possible. Rows for which the
                 cast fails for any column are filtered out.
             eager: Whether the filter operation should be performed eagerly. If `False`, the
-                returned lazy frame will
-                fail to collect if the validation does not pass.
+                returned lazy frame will fail to collect if the validation does not pass.
+            kwargs: Keyword arguments passed directly to :meth:`polars.LazyFrame.collect`
+                when `eager=True`.
 
         Returns:
             A tuple of the validated rows in the input data frame (potentially
@@ -754,8 +762,8 @@ class Schema(BaseSchema, ABC):
             match_to_schema, cls, casting=("lenient" if cast else "none")
         )
         if rules := cls._validation_rules(with_cast=cast):
-            evaluated = (
-                lf.pipe(cls._with_evaluated_rules, rules).pipe(collect_if, eager).lazy()
+            evaluated = lf.pipe(cls._with_evaluated_rules, rules).pipe(
+                collect_if, eager, **kwargs
             )
             filtered = evaluated.filter(pl.col(_COLUMN_VALID)).select(
                 cls.column_names()

{dataframely-2.11.0 → dataframely-2.12.0}/docs/api/collection/operations.rst

@@ -8,4 +8,5 @@ Operations
 
     Collection.collect_all
     Collection.join
+    Collection.pipe
     concat_collection_members

{dataframely-2.11.0 → dataframely-2.12.0}/pyproject.toml

@@ -27,7 +27,7 @@ description = "A declarative, polars-native data frame validation library"
 name = "dataframely"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "2.11.0"
+version = "2.12.0"
 
 [project.optional-dependencies]
 deltalake = ["deltalake"]

{dataframely-2.11.0 → dataframely-2.12.0}/tests/benches/test_collection.py

@@ -15,13 +15,13 @@ def partitioned_dataset(dataset: pl.DataFrame) -> dict[str, pl.DataFrame]:
             "elevation",
             "aspect",
             "slope",
-            idx=pl.int_range(pl.len(), dtype=pl.UInt32),
+            idx=pl.int_range(pl.len(), dtype=pl.UInt32).shuffle(),
         ),
         "second": dataset.select(
             "horizontal_distance_to_hydrology",
             "vertical_distance_to_hydrology",
             "horizontal_distance_to_roadways",
-            idx=pl.int_range(pl.len(), dtype=pl.UInt32),
+            idx=pl.int_range(pl.len(), dtype=pl.UInt32).shuffle(),
         ),
     }
 
@@ -95,18 +95,24 @@ class MultiFilterCollection(dy.Collection):
 
 
 @pytest.mark.benchmark(group="collection-filter-multi")
+@pytest.mark.parametrize("engine", ["in-memory", "streaming"])
 def test_multi_filter_validate(
-    benchmark: BenchmarkFixture, partitioned_dataset: dict[str, pl.DataFrame]
+    benchmark: BenchmarkFixture,
+    partitioned_dataset: dict[str, pl.DataFrame],
+    engine: str,
 ) -> None:
-    benchmark(MultiFilterCollection.validate, partitioned_dataset)
+    benchmark(MultiFilterCollection.validate, partitioned_dataset, engine=engine)
 
 
 @pytest.mark.benchmark(group="collection-filter-multi")
+@pytest.mark.parametrize("engine", ["in-memory", "streaming"])
 def test_multi_filter_filter(
-    benchmark: BenchmarkFixture, partitioned_dataset: dict[str, pl.DataFrame]
+    benchmark: BenchmarkFixture,
+    partitioned_dataset: dict[str, pl.DataFrame],
+    engine: str,
 ) -> None:
     def benchmark_fn() -> None:
-        _, failure = MultiFilterCollection.filter(partitioned_dataset)
+        _, failure = MultiFilterCollection.filter(partitioned_dataset, engine=engine)
         _ = [len(f) for f in failure.values()]
 
     benchmark(benchmark_fn)

dataframely-2.12.0/tests/collection/test_pipe.py

@@ -0,0 +1,55 @@
+# Copyright (c) QuantCo 2025-2026
+# SPDX-License-Identifier: BSD-3-Clause
+
+import dataframely as dy
+
+
+class SchemaOne(dy.Schema):
+    id = dy.Int64(primary_key=True)
+    name = dy.String(nullable=False)
+
+
+class SchemaTwo(dy.Schema):
+    id = dy.Int64(primary_key=True)
+    name = dy.String(nullable=False)
+
+
+class MyCollection(dy.Collection):
+    member_one: dy.LazyFrame[SchemaOne]
+    member_two: dy.LazyFrame[SchemaTwo]
+
+
+def test_pipe_passes_self() -> None:
+    # Arrange
+    collection = MyCollection.sample(overrides=[{"id": 1}, {"id": 2}])
+
+    # Act
+    result = collection.pipe(lambda c: c)
+
+    # Assert
+    assert result is collection
+
+
+def test_pipe_forwards_args_and_kwargs() -> None:
+    # Arrange
+    collection = MyCollection.sample(overrides=[{"id": 1}, {"id": 2}])
+
+    def combine(c: MyCollection, prefix: str, *, suffix: str) -> str:
+        return f"{prefix}{type(c).__name__}{suffix}"
+
+    # Act
+    result = collection.pipe(combine, "pre-", suffix="-post")
+
+    # Assert
+    assert result == "pre-MyCollection-post"
+
+
+def test_pipe_returns_arbitrary_type() -> None:
+    # Arrange
+    collection = MyCollection.sample(overrides=[{"id": 1}, {"id": 2}, {"id": 3}])
+
+    # Act
+    result = collection.pipe(lambda c: c.member_one.collect().height)
+
+    # Assert
+    assert result == 3