From fe8d646d273169f69422a2a506fdb281cbbab503 Mon Sep 17 00:00:00 2001 From: ElisabethBrockhausQC Date: Mon, 6 Jul 2026 18:33:32 +0200 Subject: [PATCH 1/3] feat: Add kwargs to Collection.collect_all() --- dataframely/collection/collection.py | 83 +++++++--------------------- 1 file changed, 21 insertions(+), 62 deletions(-) diff --git a/dataframely/collection/collection.py b/dataframely/collection/collection.py index 059d52b..16d41fc 100644 --- a/dataframely/collection/collection.py +++ b/dataframely/collection/collection.py @@ -112,10 +112,7 @@ def create_empty(cls) -> Self: An instance of this collection. """ return cls._init( - { - name: member.schema.create_empty() - for name, member in cls.members().items() - } + {name: member.schema.create_empty() for name, member in cls.members().items()} ) @classmethod @@ -189,11 +186,7 @@ def sample( :meth:`_preprocess_sample` method appropriately. """ # Preconditions - if ( - num_rows is not None - and overrides is not None - and len(overrides) != num_rows - ): + if num_rows is not None and overrides is not None and len(overrides) != num_rows: raise ValueError("`num_rows` mismatches the length of `overrides`.") if num_rows is None and overrides is None: num_rows = 1 @@ -205,22 +198,15 @@ def sample( # 1) Preprocess all samples to make sampling efficient and ensure shared primary # keys. - samples = ( - overrides - if overrides is not None - else [{} for _ in range(cast(int, num_rows))] - ) + samples = overrides if overrides is not None else [{} for _ in range(cast(int, num_rows))] processed_samples = [ - cls._preprocess_sample(dict(sample.items()), i, g) - for i, sample in enumerate(samples) + cls._preprocess_sample(dict(sample.items()), i, g) for i, sample in enumerate(samples) ] # 2) Ensure that all samples have primary keys assigned to ensure that we # can properly sample members. if requires_dependent_sampling: - if not all( - all(k in sample for k in primary_key) for sample in processed_samples - ): + if not all(all(k in sample for k in primary_key) for sample in processed_samples): raise ValueError("All samples must contain the common primary keys.") # 3) Sample all members independently. If we have a common primary key, we need @@ -438,13 +424,10 @@ def validate( ) details = [ - f" > Member '{member}' failed validation:\n" - + textwrap.indent(error, " ") + f" > Member '{member}' failed validation:\n" + textwrap.indent(error, " ") for member, error in errors.items() ] - message = "\n".join( - [f"{len(errors)} members failed validation:"] + details - ) + message = "\n".join([f"{len(errors)} members failed validation:"] + details) raise ValidationError(message) return filtered else: @@ -462,16 +445,12 @@ def validate( primary_key = cls.common_primary_key() filter_names = list(filters.keys()) keep = [ - filter.logic(result_cls).select( - *primary_key, pl.lit(True).alias(name) - ) + filter.logic(result_cls).select(*primary_key, pl.lit(True).alias(name)) for name, filter in filters.items() ] members = { name: ( - _join_all( - lf, *keep, on=primary_key, how="left", maintain_order="left" - ) + _join_all(lf, *keep, on=primary_key, how="left", maintain_order="left") .filter( all_rules_required( filter_names, @@ -608,10 +587,7 @@ class HospitalInvoiceData(dy.Collection): 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() + filter.logic(result_cls).select(primary_key).pipe(collect_if, eager).lazy() ) drop: dict[str, pl.LazyFrame] = {} @@ -695,9 +671,7 @@ class HospitalInvoiceData(dy.Collection): failure_lf = failure_lf.with_columns( pl.coalesce( name, - pl.col(f"{_FILTER_COLUMN_PREFIX}{name}").cast( - pl.dtype_of(name) - ), + pl.col(f"{_FILTER_COLUMN_PREFIX}{name}").cast(pl.dtype_of(name)), ) for name in member_info.schema.column_names() ).drop( @@ -808,7 +782,7 @@ def cast(cls, data: Mapping[str, FrameType], /) -> Self: # ---------------------------------- 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 @@ -821,7 +795,7 @@ def collect_all(self) -> Self: "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( @@ -900,8 +874,7 @@ def serialize(cls) -> str: for name, info in cls.members().items() }, "filters": { - name: filter.logic(cls.create_empty()) - for name, filter in cls._filters().items() + name: filter.logic(cls.create_empty()) for name, filter in cls._filters().items() }, } return json.dumps(result, cls=SchemaJSONEncoder) @@ -1062,9 +1035,7 @@ def scan_parquet( **kwargs, ) - def write_delta( - self, target: str | Path | deltalake.DeltaTable, **kwargs: Any - ) -> None: + def write_delta(self, target: str | Path | deltalake.DeltaTable, **kwargs: Any) -> None: """Write the members of this collection to Delta Lake tables. This method writes each member to a Delta Lake table at the provided target location. @@ -1262,9 +1233,7 @@ def _read( # Use strict=False when validation is "allow", "warn" or "skip" to tolerate # missing or broken collection metadata. strict = validation == "forbid" - collection_types = _deserialize_types( - serialized_collection_types, strict=strict - ) + collection_types = _deserialize_types(serialized_collection_types, strict=strict) collection_type = _reconcile_collection_types(collection_types) if cls._requires_validation_for_reading_parquets(collection_type, validation): @@ -1291,9 +1260,7 @@ def _requires_validation_for_reading_parquets( else "no collection schema to check validity can be read from the source" ) if validation == "forbid": - raise ValidationRequiredError( - f"Cannot read collection without validation: {msg}." - ) + raise ValidationRequiredError(f"Cannot read collection without validation: {msg}.") if validation == "warn": warnings.warn(f"Reading parquet file requires validation: {msg}.") return True @@ -1330,15 +1297,11 @@ def read_parquet_metadata_collection( @overload -def deserialize_collection( - data: str, strict: Literal[True] = True -) -> type[Collection]: ... +def deserialize_collection(data: str, strict: Literal[True] = True) -> type[Collection]: ... @overload -def deserialize_collection( - data: str, strict: Literal[False] -) -> type[Collection] | None: ... +def deserialize_collection(data: str, strict: Literal[False]) -> type[Collection] | None: ... @overload @@ -1409,9 +1372,7 @@ def deserialize_collection(data: str, strict: bool = True) -> type[Collection] | ) except (ValueError, TypeError, JSONDecodeError, plexc.ComputeError) as e: if strict: - raise DeserializationError( - "The Collection metadata could not be deserialized" - ) from e + raise DeserializationError("The Collection metadata could not be deserialized") from e return None @@ -1430,9 +1391,7 @@ def _join_all( return result -def _extract_keys_if_exist( - data: Mapping[str, Any], keys: Sequence[str] -) -> dict[str, Any]: +def _extract_keys_if_exist(data: Mapping[str, Any], keys: Sequence[str]) -> dict[str, Any]: return {key: data[key] for key in keys if key in data} From 2d9d37d54ec1fd519f8b4a083590a58048b3935c Mon Sep 17 00:00:00 2001 From: ElisabethBrockhausQC Date: Mon, 6 Jul 2026 18:37:32 +0200 Subject: [PATCH 2/3] lint --- dataframely/collection/collection.py | 79 +++++++++++++++++++++------- 1 file changed, 60 insertions(+), 19 deletions(-) diff --git a/dataframely/collection/collection.py b/dataframely/collection/collection.py index 16d41fc..7fa9ef5 100644 --- a/dataframely/collection/collection.py +++ b/dataframely/collection/collection.py @@ -112,7 +112,10 @@ def create_empty(cls) -> Self: An instance of this collection. """ return cls._init( - {name: member.schema.create_empty() for name, member in cls.members().items()} + { + name: member.schema.create_empty() + for name, member in cls.members().items() + } ) @classmethod @@ -186,7 +189,11 @@ def sample( :meth:`_preprocess_sample` method appropriately. """ # Preconditions - if num_rows is not None and overrides is not None and len(overrides) != num_rows: + if ( + num_rows is not None + and overrides is not None + and len(overrides) != num_rows + ): raise ValueError("`num_rows` mismatches the length of `overrides`.") if num_rows is None and overrides is None: num_rows = 1 @@ -198,15 +205,22 @@ def sample( # 1) Preprocess all samples to make sampling efficient and ensure shared primary # keys. - samples = overrides if overrides is not None else [{} for _ in range(cast(int, num_rows))] + samples = ( + overrides + if overrides is not None + else [{} for _ in range(cast(int, num_rows))] + ) processed_samples = [ - cls._preprocess_sample(dict(sample.items()), i, g) for i, sample in enumerate(samples) + cls._preprocess_sample(dict(sample.items()), i, g) + for i, sample in enumerate(samples) ] # 2) Ensure that all samples have primary keys assigned to ensure that we # can properly sample members. if requires_dependent_sampling: - if not all(all(k in sample for k in primary_key) for sample in processed_samples): + if not all( + all(k in sample for k in primary_key) for sample in processed_samples + ): raise ValueError("All samples must contain the common primary keys.") # 3) Sample all members independently. If we have a common primary key, we need @@ -424,10 +438,13 @@ def validate( ) details = [ - f" > Member '{member}' failed validation:\n" + textwrap.indent(error, " ") + f" > Member '{member}' failed validation:\n" + + textwrap.indent(error, " ") for member, error in errors.items() ] - message = "\n".join([f"{len(errors)} members failed validation:"] + details) + message = "\n".join( + [f"{len(errors)} members failed validation:"] + details + ) raise ValidationError(message) return filtered else: @@ -445,12 +462,16 @@ def validate( primary_key = cls.common_primary_key() filter_names = list(filters.keys()) keep = [ - filter.logic(result_cls).select(*primary_key, pl.lit(True).alias(name)) + filter.logic(result_cls).select( + *primary_key, pl.lit(True).alias(name) + ) for name, filter in filters.items() ] members = { name: ( - _join_all(lf, *keep, on=primary_key, how="left", maintain_order="left") + _join_all( + lf, *keep, on=primary_key, how="left", maintain_order="left" + ) .filter( all_rules_required( filter_names, @@ -587,7 +608,10 @@ class HospitalInvoiceData(dy.Collection): 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() + filter.logic(result_cls) + .select(primary_key) + .pipe(collect_if, eager) + .lazy() ) drop: dict[str, pl.LazyFrame] = {} @@ -671,7 +695,9 @@ class HospitalInvoiceData(dy.Collection): failure_lf = failure_lf.with_columns( pl.coalesce( name, - pl.col(f"{_FILTER_COLUMN_PREFIX}{name}").cast(pl.dtype_of(name)), + pl.col(f"{_FILTER_COLUMN_PREFIX}{name}").cast( + pl.dtype_of(name) + ), ) for name in member_info.schema.column_names() ).drop( @@ -874,7 +900,8 @@ def serialize(cls) -> str: for name, info in cls.members().items() }, "filters": { - name: filter.logic(cls.create_empty()) for name, filter in cls._filters().items() + name: filter.logic(cls.create_empty()) + for name, filter in cls._filters().items() }, } return json.dumps(result, cls=SchemaJSONEncoder) @@ -1035,7 +1062,9 @@ def scan_parquet( **kwargs, ) - def write_delta(self, target: str | Path | deltalake.DeltaTable, **kwargs: Any) -> None: + def write_delta( + self, target: str | Path | deltalake.DeltaTable, **kwargs: Any + ) -> None: """Write the members of this collection to Delta Lake tables. This method writes each member to a Delta Lake table at the provided target location. @@ -1233,7 +1262,9 @@ def _read( # Use strict=False when validation is "allow", "warn" or "skip" to tolerate # missing or broken collection metadata. strict = validation == "forbid" - collection_types = _deserialize_types(serialized_collection_types, strict=strict) + collection_types = _deserialize_types( + serialized_collection_types, strict=strict + ) collection_type = _reconcile_collection_types(collection_types) if cls._requires_validation_for_reading_parquets(collection_type, validation): @@ -1260,7 +1291,9 @@ def _requires_validation_for_reading_parquets( else "no collection schema to check validity can be read from the source" ) if validation == "forbid": - raise ValidationRequiredError(f"Cannot read collection without validation: {msg}.") + raise ValidationRequiredError( + f"Cannot read collection without validation: {msg}." + ) if validation == "warn": warnings.warn(f"Reading parquet file requires validation: {msg}.") return True @@ -1297,11 +1330,15 @@ def read_parquet_metadata_collection( @overload -def deserialize_collection(data: str, strict: Literal[True] = True) -> type[Collection]: ... +def deserialize_collection( + data: str, strict: Literal[True] = True +) -> type[Collection]: ... @overload -def deserialize_collection(data: str, strict: Literal[False]) -> type[Collection] | None: ... +def deserialize_collection( + data: str, strict: Literal[False] +) -> type[Collection] | None: ... @overload @@ -1372,7 +1409,9 @@ def deserialize_collection(data: str, strict: bool = True) -> type[Collection] | ) except (ValueError, TypeError, JSONDecodeError, plexc.ComputeError) as e: if strict: - raise DeserializationError("The Collection metadata could not be deserialized") from e + raise DeserializationError( + "The Collection metadata could not be deserialized" + ) from e return None @@ -1391,7 +1430,9 @@ def _join_all( return result -def _extract_keys_if_exist(data: Mapping[str, Any], keys: Sequence[str]) -> dict[str, Any]: +def _extract_keys_if_exist( + data: Mapping[str, Any], keys: Sequence[str] +) -> dict[str, Any]: return {key: data[key] for key in keys if key in data} From 4ae558df942a03ad1df7e6939311fd544a484513 Mon Sep 17 00:00:00 2001 From: ElisabethBrockhausQC Date: Mon, 6 Jul 2026 18:39:01 +0200 Subject: [PATCH 3/3] docstring --- dataframely/collection/collection.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/dataframely/collection/collection.py b/dataframely/collection/collection.py index 7fa9ef5..d85662f 100644 --- a/dataframely/collection/collection.py +++ b/dataframely/collection/collection.py @@ -814,6 +814,9 @@ def collect_all(self, **kwargs: Any) -> Self: 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