dataframely 1.8.2__tar.gz → 1.9.0__tar.gz

{dataframely-1.8.2 → dataframely-1.9.0}/.gitignore

@@ -388,3 +388,6 @@ $RECYCLE.BIN/
 
 # interactive testing
 /*.ipynb
+
+# pytest-benchmark
+.benchmarks/

{dataframely-1.8.2 → dataframely-1.9.0}/PKG-INFO

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

{dataframely-1.8.2 → dataframely-1.9.0}/dataframely/_base_collection.py

@@ -66,8 +66,8 @@ class CollectionMember:
 # --------------------------------------- UTILS -------------------------------------- #
 
 
-def _common_primary_keys(columns: Iterable[type[Schema]]) -> set[str]:
-    return set.intersection(*[set(schema.primary_keys()) for schema in columns])
+def _common_primary_keys(schemas: Iterable[type[Schema]]) -> set[str]:
+    return set.intersection(*[set(schema.primary_keys()) for schema in schemas])
 
 
 # ------------------------------------------------------------------------------------ #
@@ -137,7 +137,7 @@ class CollectionMeta(ABCMeta):
                     for member in result.members.values()
                 ),
                 *(
-                    set(member.schema._validation_rules())
+                    set(member.schema._validation_rules(with_cast=True))
                     for member in result.members.values()
                 ),
             )

{dataframely-1.8.2 → dataframely-1.9.0}/dataframely/_base_schema.py

@@ -12,7 +12,7 @@ from typing import Any
 
 import polars as pl
 
-from ._rule import GroupRule, Rule
+from ._rule import DtypeCastRule, GroupRule, Rule
 from .columns import Column
 from .exc import ImplementationError
 
@@ -24,11 +24,13 @@ else:
 _COLUMN_ATTR = "__dataframely_columns__"
 _RULE_ATTR = "__dataframely_rules__"
 
+ORIGINAL_COLUMN_PREFIX = "__DATAFRAMELY_ORIGINAL__"
+
 # --------------------------------------- UTILS -------------------------------------- #
 
 
 def _build_rules(
-    custom: dict[str, Rule], columns: dict[str, Column]
+    custom: dict[str, Rule], columns: dict[str, Column], *, with_cast: bool
 ) -> dict[str, Rule]:
     # NOTE: Copy here to prevent in-place modification of the custom rules
     rules: dict[str, Rule] = copy(custom)
@@ -46,6 +48,21 @@ def _build_rules(
     }
     rules.update(column_rules)
 
+    # Add casting rules if requested. Here, we can simply check whether the nullability
+    # property of a column changes due to lenient dtype casting. Whenever casting fails,
+    # the value is set to `null`, mismatching the previous nullability.
+    # NOTE: This check assumes that both the original and cast column are present in the
+    #  data frame.
+    if with_cast:
+        casting_rules = {
+            f"{col_name}|dtype": DtypeCastRule(
+                pl.col(col_name).is_null()
+                == pl.col(f"{ORIGINAL_COLUMN_PREFIX}{col_name}").is_null()
+            )
+            for col_name in columns
+        }
+        rules.update(casting_rules)
+
     return rules
 
 
@@ -93,9 +110,7 @@ class SchemaMeta(ABCMeta):
         # we assume that users cast dtypes, i.e. additional rules for dtype casting
         # are also checked.
         all_column_names = set(result.columns)
-        all_rule_names = set(_build_rules(result.rules, result.columns).keys()) | set(
-            f"{col}|dtype" for col in result.columns
-        )
+        all_rule_names = set(_build_rules(result.rules, result.columns, with_cast=True))
         common_names = all_column_names & all_rule_names
         if len(common_names) > 0:
             common_list = ", ".join(sorted(f"'{col}'" for col in common_names))
@@ -189,8 +204,10 @@ class BaseSchema(metaclass=SchemaMeta):
         return _primary_keys(cls.columns())
 
     @classmethod
-    def _validation_rules(cls) -> dict[str, Rule]:
-        return _build_rules(cls._schema_validation_rules(), cls.columns())
+    def _validation_rules(cls, *, with_cast: bool) -> dict[str, Rule]:
+        return _build_rules(
+            cls._schema_validation_rules(), cls.columns(), with_cast=with_cast
+        )
 
     @classmethod
     def _schema_validation_rules(cls) -> dict[str, Rule]:

{dataframely-1.8.2 → dataframely-1.9.0}/dataframely/_polars.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import datetime as dt
-from collections.abc import Iterable
 from typing import TypeVar
 
 import polars as pl
@@ -14,27 +13,6 @@ FrameType = TypeVar("FrameType", pl.DataFrame, pl.LazyFrame)
 EPOCH_DATETIME = dt.datetime(1970, 1, 1)
 SECONDS_PER_DAY = 86400
 
-# --------------------------------------- JOINS -------------------------------------- #
-
-
-def join_all_inner(dfs: Iterable[FrameType], on: str | list[str]) -> FrameType:
-    it = iter(dfs)
-    result = next(it)
-    while (df := next(it, None)) is not None:
-        result = result.join(df, on=on)
-    return result
-
-
-def join_all_outer(dfs: Iterable[FrameType], on: str | list[str]) -> FrameType:
-    it = iter(dfs)
-    result = next(it)
-    while (df := next(it, None)) is not None:
-        result = result.join(df, on=on, how="full", coalesce=True)
-    return result
-
-
-# ------------------------------------- DATETIMES ------------------------------------ #
-
 
 def date_matches_resolution(t: dt.date, resolution: str) -> bool:
     return pl.Series([t], dtype=pl.Date).dt.truncate(resolution).item() == t

{dataframely-1.8.2 → dataframely-1.9.0}/dataframely/_rule.py

@@ -59,6 +59,14 @@ class Rule:
         return str(self.expr)
 
 
+class DtypeCastRule(Rule):
+    """Rule that evaluates whether casting a column to another dtype is successful.
+
+    The only purpose of this rule is to provide a runtime type to distinguish it from
+    other rules.
+    """
+
+
 class GroupRule(Rule):
     """Rule that is evaluated on a group of columns."""
 
@@ -147,7 +155,8 @@ def with_evaluation_rules(lf: pl.LazyFrame, rules: dict[str, Rule]) -> pl.LazyFr
         while ``False`` indicates an issue.
     """
     # Rules must be distinguished into two types of rules:
-    #  1. Simple rules can simply be selected on the data frame
+    #  1. Simple rules can simply be selected on the data frame (this includes rules
+    #     that check whether dtype casts succeeded)
     #  2. "Group" rules require a `group_by` and a subsequent join
     simple_exprs = {
         name: rule.expr
@@ -160,7 +169,7 @@ def with_evaluation_rules(lf: pl.LazyFrame, rules: dict[str, Rule]) -> pl.LazyFr
 
     # Before we can select all of the simple expressions, we need to turn the
     # group rules into something to use in a `select` statement as well.
-    return (
+    result = (
         # NOTE: A value of `null` always validates successfully as nullability should
         #  already be checked via dedicated rules.
         _with_group_rules(lf, group_rules).with_columns(
@@ -168,6 +177,27 @@ def with_evaluation_rules(lf: pl.LazyFrame, rules: dict[str, Rule]) -> pl.LazyFr
         )
     )
 
+    # If there is at least one rule that checks for successful dtype casting, we need
+    # to take an extra step: rules other than the "dtype rules" might not be reliable
+    # if casting failed, i.e. if any of the "dtype rules" evaluated to `False`. For
+    # this reason, we set all other rule evaluations to `null` in the case of dtype
+    # casting failure.
+    dtype_rule_names = [
+        name for name, rule in rules.items() if isinstance(rule, DtypeCastRule)
+    ]
+    if len(dtype_rule_names) > 0:
+        non_dtype_rule_names = [
+            name for name, rule in rules.items() if not isinstance(rule, DtypeCastRule)
+        ]
+        all_dtype_casts_valid = pl.all_horizontal(dtype_rule_names)
+        return result.with_columns(
+            pl.when(all_dtype_casts_valid)
+            .then(pl.col(non_dtype_rule_names))
+            .otherwise(pl.lit(None, dtype=pl.Boolean))
+        )
+
+    return result
+
 
 def _with_group_rules(lf: pl.LazyFrame, rules: dict[str, GroupRule]) -> pl.LazyFrame:
     # First, we partition the rules by group columns. This will minimize the number
@@ -184,16 +214,12 @@ def _with_group_rules(lf: pl.LazyFrame, rules: dict[str, GroupRule]) -> pl.LazyF
         # We group by the group columns and apply all expressions
         group_evaluations[group_columns] = lf.group_by(group_columns).agg(**group_rules)
 
-    # Eventually, we apply the rule evaluations onto the input data frame. For this,
-    # we're using left-joins. This has two effects:
-    #  - We're essentially "broadcasting" the results within each group across rows
-    #    in the same group.
-    #  - While an inner-join would be semantically more accurate, the left-join
-    #    preserves the order of the left data frame.
+    # Eventually, we apply the rule evaluations onto the input data frame. For this, we
+    # "broadcast" the results within each group across rows in the same group.
     result = lf
     for group_columns, frame in group_evaluations.items():
         result = result.join(
-            frame, on=list(group_columns), how="left", nulls_equal=True
+            frame, on=list(group_columns), nulls_equal=True, maintain_order="left"
         )
     return result
 

{dataframely-1.8.2 → dataframely-1.9.0}/dataframely/collection.py

@@ -10,14 +10,14 @@ from collections.abc import Iterable, Mapping, Sequence
 from dataclasses import asdict
 from json import JSONDecodeError
 from pathlib import Path
-from typing import IO, Annotated, Any, cast
+from typing import IO, Annotated, Any, Literal, cast
 
 import polars as pl
 import polars.exceptions as plexc
 
 from ._base_collection import BaseCollection, CollectionMember
 from ._filter import Filter
-from ._polars import FrameType, join_all_inner, join_all_outer
+from ._polars import FrameType
 from ._serialization import (
     COLLECTION_METADATA_KEY,
     SERIALIZATION_FORMAT_VERSION,
@@ -41,6 +41,8 @@ if sys.version_info >= (3, 11):
 else:
     from typing_extensions import Self
 
+_FILTER_COLUMN_PREFIX = "__DATAFRAMELY_FILTER_COLUMN__"
+
 
 class Collection(BaseCollection, ABC):
     """Base class for all collections of data frames with a predefined schema.
@@ -376,7 +378,8 @@ class Collection(BaseCollection, ABC):
         Returns:
             An instance of the collection. All members of the collection are guaranteed
             to be valid with respect to their respective schemas and the filters on this
-            collection did not remove rows from any member.
+            collection did not remove rows from any member. The input order of each
+            member is maintained.
         """
         out, failure = cls.filter(data, cast=cast)
         if any(len(fail) > 0 for fail in failure.values()):
@@ -439,7 +442,9 @@ class Collection(BaseCollection, ABC):
               filtered out by any of the collection's filters. While collection members
               are always instances of :class:`~polars.LazyFrame`, the members of the
               returned collection are essentially eager as they are constructed by
-              calling ``.lazy()`` on eager data frames.
+              calling ``.lazy()`` on eager data frames. Just like in polars' native
+              :meth:`~polars.DataFrame.filter`, the order of rows is maintained in all
+              returned data frames.
             - A mapping from member name to a :class:`FailureInfo` object which provides
               details on why individual rows had been removed. Optional members are only
               included in this dictionary if they had been provided in the input.
@@ -465,7 +470,7 @@ class Collection(BaseCollection, ABC):
                 data[member_name], cast=cast
             )
 
-        # Once we're done that, we can apply the filters on this collection. To this end,
+        # Once we've done that, we can apply the filters on this collection. To this end,
         # we iterate over all filters and store the filter results.
         filters = cls._filters()
         if len(filters) > 0:
@@ -476,41 +481,35 @@ class Collection(BaseCollection, ABC):
             for name, filter in filters.items():
                 keep[name] = filter.logic(result_cls).select(primary_keys).collect()
 
-            # Using the filter results, we can define a joint data frame that we use to filter
-            # the input.
-            all_keep = join_all_inner(
-                [df.lazy() for df in keep.values()], on=primary_keys
-            ).collect()
-
-            # Now we can iterate over the results where we do the following:
-            # - Join the current result onto `all_keep` to get rid of the rows we do not
-            #   want to keep.
-            # - Anti-join onto each individual filter to extend the failure reasons
+            # 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
             for member_name, filtered in results.items():
-                if cls.members()[member_name].ignored_in_filters:
+                member_info = cls.members()[member_name]
+                if member_info.ignored_in_filters:
                     continue
-                results[member_name] = filtered.join(all_keep, on=primary_keys)
-
-                new_failure_names = list(filters.keys())
-                new_failure_pks = [
-                    filtered.select(primary_keys)
-                    .lazy()
-                    .unique()
-                    .join(filter_keep.lazy(), on=primary_keys, how="anti")
-                    .with_columns(pl.lit(False).alias(name))
-                    for name, filter_keep in keep.items()
-                ]
-                # NOTE: The outer join might generate NULL values if a primary key is not
-                #  filtered out by all filters. In this case, we want to assign a validation
-                #  value of `True`.
-                all_new_failure_pks = join_all_outer(
-                    new_failure_pks, on=primary_keys
-                ).with_columns(pl.col(new_failure_names).fill_null(True))
-
-                # At this point, we have a data frame with the primary keys of the *excluded*
-                # rows of the member result along with the reasons. We join on the result again
-                # which will give us the same shape as the current failure info lf. Thus, we can
-                # simply concatenate diagonally, resulting in a failure info object as follows:
+
+                lf_with_eval = filtered.lazy()
+                for name, filter_keep in keep.items():
+                    lf_with_eval = lf_with_eval.join(
+                        filter_keep.lazy().with_columns(pl.lit(True).alias(name)),
+                        on=primary_keys,
+                        how="left",
+                        maintain_order="left",
+                    ).with_columns(pl.col(name).fill_null(False))
+
+                result_with_eval = lf_with_eval.collect()
+
+                # Filtering `result_with_eval` by the rows for which all joins
+                # "succeeded", we can identify the rows that pass all the filters. We
+                # keep these rows for the result.
+                results[member_name] = result_with_eval.filter(
+                    pl.all_horizontal(keep.keys())
+                ).drop(keep.keys())
+
+                # Filtering `result_with_eval` with the inverse condition, we find all
+                # the problematic rows. We can build a single failure info object by
+                # simply concatenating diagonally with the already existing failure. The
+                # resulting failure info looks as follows:
                 #
                 #  | Source Data | Rule Columns (schema) | Filter Name Columns (collection) |
                 #  | ----------- | --------------------- | -------------------------------- |
@@ -518,21 +517,98 @@ class Collection(BaseCollection, ABC):
                 #  | ...         | NULL                  | <filled>                         |
                 #
                 failure = failures[member_name]
-                new_failure = FailureInfo(
-                    lf=pl.concat(
-                        [
-                            failure._lf,
-                            filtered.lazy().join(all_new_failure_pks, on=primary_keys),
-                        ],
-                        how="diagonal",
-                    ),
-                    rule_columns=failure._rule_columns + new_failure_names,
+                filtered_failure = result_with_eval.filter(
+                    ~pl.all_horizontal(keep.keys())
+                ).lazy()
+
+                # If we cast previously, `failure` and `filtered_failure` have different
+                # dtypes for the source data: `failure` keeps the original dtypes while
+                # `filtered_failure` has the target dtypes. Hence, we need to cast
+                # `filtered_failure` to the original dtypes. This is safe because any
+                # row in `filtered_failure` must have already been successfully cast and
+                # a "roundtrip cast" is always possible.
+                # Doing this in a fully lazy way is not trivial: we do a diagonal
+                # concatenation where we duplicate each column of the source data. We
+                # then coalesce the two versions into the original column dtype.
+                if cast:
+                    filtered_failure = filtered_failure.rename(
+                        {
+                            name: f"{_FILTER_COLUMN_PREFIX}{name}"
+                            for name in member_info.schema.column_names()
+                        }
+                    )
+
+                failure_lf = pl.concat([failure._lf, filtered_failure], how="diagonal")
+                if cast:
+                    failure_lf = failure_lf.with_columns(
+                        pl.coalesce(
+                            name,
+                            pl.col(f"{_FILTER_COLUMN_PREFIX}{name}").cast(
+                                pl.dtype_of(name)
+                            ),
+                        )
+                        for name in member_info.schema.column_names()
+                    ).drop(
+                        f"{_FILTER_COLUMN_PREFIX}{name}"
+                        for name in member_info.schema.column_names()
+                    )
+
+                failures[member_name] = FailureInfo(
+                    lf=failure_lf,
+                    rule_columns=failure._rule_columns + list(keep.keys()),
                     schema=failure.schema,
                 )
-                failures[member_name] = new_failure
 
         return cls._init(results), failures
 
+    def join(
+        self,
+        primary_keys: pl.LazyFrame,
+        how: Literal["semi", "anti"] = "semi",
+        maintain_order: Literal["none", "left"] = "none",
+    ) -> Self:
+        """Filter the collection by joining onto a data frame containing entries for the
+        common primary key columns whose respective rows should be kept or removed in
+        the collection members.
+
+        Args:
+            primary_keys: The data frame to join on. Must contain the common primary key
+                columns of the collection.
+            how: The join strategy to use. Like in polars, `semi` will keep all rows
+                that can be found in `primary_keys`, `anti` will remove them.
+            maintain_order: The `maintain_order` option to use for the polars join.
+
+        Returns:
+            The collection, with members potentially reduced in length.
+
+        Raises:
+            ValueError: If the collection contains any member that is annotated with
+                `ignored_in_filters=True`.
+
+        Attention:
+            This method does not validate the resulting collection. Ensure to only use
+            this if the resulting collection still satisfies the filters of the
+            collection. The joins are not evaluated eagerly. Therefore, a downstream
+            call to :meth:`collect` might fail, especially if `primary_keys` does not
+            contain all columns for all common primary keys.
+        """
+        if any(member.ignored_in_filters for member in self.members().values()):
+            raise ValueError(
+                "The join operation is not supported for collections with members that are ignored in filters."
+            )
+
+        return self.cast(
+            {
+                key: lf.join(
+                    primary_keys,
+                    on=self.common_primary_keys(),
+                    how=how,
+                    maintain_order=maintain_order,
+                )
+                for key, lf in self.to_dict().items()
+            }
+        )
+
     # ------------------------------------ CASTING ----------------------------------- #
 
     @classmethod

{dataframely-1.8.2 → dataframely-1.9.0}/dataframely/failure.py

@@ -19,6 +19,7 @@ if TYPE_CHECKING:  # pragma: no cover
     from .schema import Schema
 
 RULE_METADATA_KEY = "dataframely_rule_columns"
+UNKNOWN_SCHEMA_NAME = "__DATAFRAMELY_UNKNOWN__"
 
 S = TypeVar("S", bound=BaseSchema)
 
@@ -176,7 +177,7 @@ class FailureInfo(Generic[S]):
     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 @@ class FailureInfo(Generic[S]):
             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,
         )