diff --git a/dataframely/_polars.py b/dataframely/_polars.py index b19bb50..cd58edb 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 d85662f..7583465 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,7 +645,11 @@ 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[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] # Filtering `lf_with_eval` by the rows for which all joins # "succeeded", we can identify the rows that pass all the filters. We diff --git a/dataframely/schema.py b/dataframely/schema.py index c7fc92c..726aa6d 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()