From 7137163e1bb88ab62a5192d4772f2d61c461e0ec Mon Sep 17 00:00:00 2001 From: Oliver Borchert Date: Mon, 6 Jul 2026 18:51:16 +0200 Subject: [PATCH 1/5] perf: Use semi-join instead of `unique` for checking 1:N relationships --- dataframely/functional.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dataframely/functional.py b/dataframely/functional.py index b099fe54..baabc713 100644 --- a/dataframely/functional.py +++ b/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") # ------------------------------------------------------------------------------------ # From 3574fd7ef87e8da92a496f950e9bd1ddc68bb4ee Mon Sep 17 00:00:00 2001 From: Oliver Borchert Date: Mon, 6 Jul 2026 22:41:25 +0200 Subject: [PATCH 2/5] Fix benchmark --- tests/benches/test_collection.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/benches/test_collection.py b/tests/benches/test_collection.py index f3a1aa3b..0bf0cb5f 100644 --- a/tests/benches/test_collection.py +++ b/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(), ), } From a5c02029c12b0c4b2c770b48ecf65bc51f883c78 Mon Sep 17 00:00:00 2001 From: Oliver Borchert Date: Mon, 6 Jul 2026 23:15:29 +0200 Subject: [PATCH 3/5] perf: Leverage `collect_all` for eager collection validation & filtering --- dataframely/_polars.py | 17 ++++++++++---- dataframely/collection/collection.py | 33 +++++++++++++--------------- dataframely/schema.py | 4 ++-- 3 files changed, 30 insertions(+), 24 deletions(-) diff --git a/dataframely/_polars.py b/dataframely/_polars.py index b19bb50f..cd58edb9 100644 --- a/dataframely/_polars.py +++ b/dataframely/_polars.py @@ -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) -> pl.LazyFrame: + """Collect a lazy frame based on `condition`.""" if condition: - return lf.collect() + return lf.collect().lazy() return lf + + +def collect_all_if( + lfs: dict[str, pl.LazyFrame], condition: bool +) -> dict[str, pl.LazyFrame]: + """Collect the lazy frames in the dictionary based on `condition`.""" + if condition: + dfs = pl.collect_all(lfs.values()) + return {k: v.lazy() for k, v in zip(lfs.keys(), dfs)} + return lfs diff --git a/dataframely/collection/collection.py b/dataframely/collection/collection.py index 059d52b6..83f7dfb9 100644 --- a/dataframely/collection/collection.py +++ b/dataframely/collection/collection.py @@ -31,7 +31,7 @@ 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, @@ -605,28 +605,25 @@ class HospitalInvoiceData(dy.Collection): 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) - 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) # 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: @@ -648,8 +645,8 @@ class HospitalInvoiceData(dy.Collection): 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 = collect_all_if(lfs_with_eval, eager) + for member_name, lf_with_eval in lfs_with_eval.items(): # Filtering `lf_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. diff --git a/dataframely/schema.py b/dataframely/schema.py index c7fc92c4..726aa6dc 100644 --- a/dataframely/schema.py +++ b/dataframely/schema.py @@ -754,8 +754,8 @@ def filter( 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 ) filtered = evaluated.filter(pl.col(_COLUMN_VALID)).select( cls.column_names() From 86109d20ccd4f5c0b56cf66ac50297201ad0c828 Mon Sep 17 00:00:00 2001 From: Oliver Borchert Date: Mon, 6 Jul 2026 23:33:05 +0200 Subject: [PATCH 4/5] Fix --- dataframely/collection/collection.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/dataframely/collection/collection.py b/dataframely/collection/collection.py index 83f7dfb9..99c5db8b 100644 --- a/dataframely/collection/collection.py +++ b/dataframely/collection/collection.py @@ -647,6 +647,8 @@ class HospitalInvoiceData(dy.Collection): lfs_with_eval = collect_all_if(lfs_with_eval, eager) 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 # keep these rows for the result. From f7b80cececb11c2795ea1cbe7905c0a76d2909cf Mon Sep 17 00:00:00 2001 From: Oliver Borchert Date: Mon, 6 Jul 2026 23:47:17 +0200 Subject: [PATCH 5/5] Fix --- dataframely/collection/collection.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/dataframely/collection/collection.py b/dataframely/collection/collection.py index 99c5db8b..2c8a5dc7 100644 --- a/dataframely/collection/collection.py +++ b/dataframely/collection/collection.py @@ -645,6 +645,8 @@ class HospitalInvoiceData(dy.Collection): maintain_order="left", ).with_columns(pl.col(name).fill_null(True)) + lfs_with_eval[member_name] = lf_with_eval + lfs_with_eval = collect_all_if(lfs_with_eval, eager) for member_name, lf_with_eval in lfs_with_eval.items(): member_info = cls.members()[member_name]