From a19ff11c1ae75e3c045915a76f3ba42221b69dc6 Mon Sep 17 00:00:00 2001 From: Noritaka Sekiyama Date: Thu, 9 Jul 2026 08:14:16 +0900 Subject: [PATCH 1/4] Read tables with multi-argument transforms as unknown transforms Per the V3 spec, readers must read tables with unknown transforms, ignoring them. PyIceberg raised on partition or sort fields with more than one entry in source-ids, so such tables failed to load. Model source-ids on PartitionField and SortField, treat multi-argument transforms as UnknownTransform (null partition values, always-true projection), and serialize per spec: source-id for single-argument transforms, source-ids otherwise. Also fix two latent tolerance bugs: unknown transform names sharing a prefix with known ones (e.g. bucketv2[4]) failed to parse, and str(UnknownTransform) returned "unknown" instead of the original name, corrupting metadata on rewrite. Closes #3628 --- pyiceberg/partitioning.py | 24 ++++++++++++++++-- pyiceberg/table/sorting.py | 23 +++++++++++++++-- pyiceberg/transforms.py | 16 +++++++++--- tests/table/test_partitioning.py | 42 ++++++++++++++++++++++++++++++++ tests/table/test_sorting.py | 17 +++++++++++++ tests/test_transforms.py | 16 ++++++++++++ 6 files changed, 131 insertions(+), 7 deletions(-) diff --git a/pyiceberg/partitioning.py b/pyiceberg/partitioning.py index 4279ff5f09..d9917931c6 100644 --- a/pyiceberg/partitioning.py +++ b/pyiceberg/partitioning.py @@ -29,6 +29,7 @@ Field, PlainSerializer, WithJsonSchema, + model_serializer, model_validator, ) @@ -77,6 +78,7 @@ class PartitionField(IcebergBaseModel): """ source_id: int = Field(alias="source-id") + source_ids: list[int] | None = Field(alias="source-ids", default=None, repr=False) field_id: int = Field(alias="field-id") transform: Annotated[ # type: ignore Transform, @@ -115,13 +117,31 @@ def map_source_ids_onto_source_id(cls, data: Any) -> Any: if len(source_ids) == 0: raise ValueError("Empty source-ids is not allowed") if len(source_ids) > 1: - raise ValueError("Multi argument transforms are not yet supported") + # Multi-argument transforms cannot be evaluated; per the spec, v3 readers + # must read tables with such transforms, ignoring them + data["transform"] = UnknownTransform(transform=str(data.get("transform"))) + else: + data.pop("source-ids", None) data["source-id"] = source_ids[0] return data + @model_serializer(mode="wrap") + def _serialize_source_ids(self, handler: Any) -> Any: + serialized = handler(self) + # Per the spec, single-argument transforms write only source-id and + # multi-argument transforms write only source-ids + if self.source_ids is not None and len(self.source_ids) > 1: + serialized.pop("source-id", None) + serialized.pop("source_id", None) + else: + serialized.pop("source-ids", None) + serialized.pop("source_ids", None) + return serialized + def __str__(self) -> str: """Return the string representation of the PartitionField class.""" - return f"{self.field_id}: {self.name}: {self.transform}({self.source_id})" + sources = ", ".join(str(s) for s in self.source_ids) if self.source_ids else self.source_id + return f"{self.field_id}: {self.name}: {self.transform}({sources})" class PartitionSpec(IcebergBaseModel): diff --git a/pyiceberg/table/sorting.py b/pyiceberg/table/sorting.py index 61f34c4780..08aeec016e 100644 --- a/pyiceberg/table/sorting.py +++ b/pyiceberg/table/sorting.py @@ -24,12 +24,13 @@ Field, PlainSerializer, WithJsonSchema, + model_serializer, model_validator, ) from pyiceberg.exceptions import ValidationError from pyiceberg.schema import Schema -from pyiceberg.transforms import IdentityTransform, Transform, parse_transform +from pyiceberg.transforms import IdentityTransform, Transform, UnknownTransform, parse_transform from pyiceberg.typedef import IcebergBaseModel from pyiceberg.types import IcebergType @@ -107,11 +108,29 @@ def map_source_ids_onto_source_id(cls, data: Any) -> Any: if len(source_ids) == 0: raise ValueError("Empty source-ids is not allowed") if len(source_ids) > 1: - raise ValueError("Multi argument transforms are not yet supported") + # Multi-argument transforms cannot be evaluated; per the spec, v3 readers + # must read tables with such transforms, ignoring them + data["transform"] = UnknownTransform(transform=str(data.get("transform"))) + else: + data.pop("source-ids", None) data["source-id"] = source_ids[0] return data + @model_serializer(mode="wrap") + def _serialize_source_ids(self, handler: Any) -> Any: + serialized = handler(self) + # Per the spec, single-argument transforms write only source-id and + # multi-argument transforms write only source-ids + if self.source_ids is not None and len(self.source_ids) > 1: + serialized.pop("source-id", None) + serialized.pop("source_id", None) + else: + serialized.pop("source-ids", None) + serialized.pop("source_ids", None) + return serialized + source_id: int = Field(alias="source-id") + source_ids: list[int] | None = Field(alias="source-ids", default=None, repr=False) transform: Annotated[ # type: ignore Transform, BeforeValidator(parse_transform), diff --git a/pyiceberg/transforms.py b/pyiceberg/transforms.py index 5e0027a829..ab475ab51c 100644 --- a/pyiceberg/transforms.py +++ b/pyiceberg/transforms.py @@ -31,7 +31,7 @@ import mmh3 from pydantic import Field, PositiveInt, PrivateAttr -from pyiceberg.exceptions import NotInstalledError +from pyiceberg.exceptions import NotInstalledError, ValidationError from pyiceberg.expressions import ( BoundEqualTo, BoundGreaterThan, @@ -226,9 +226,15 @@ def parse_transform(v: Any) -> Transform[Any, Any]: elif v == VOID: return VoidTransform() elif v.startswith(BUCKET): - return BucketTransform(num_buckets=BUCKET_PARSER.match(v)) + try: + return BucketTransform(num_buckets=BUCKET_PARSER.match(v)) + except ValidationError: + return UnknownTransform(transform=v) elif v.startswith(TRUNCATE): - return TruncateTransform(width=TRUNCATE_PARSER.match(v)) + try: + return TruncateTransform(width=TRUNCATE_PARSER.match(v)) + except ValidationError: + return UnknownTransform(transform=v) elif v == YEAR: return YearTransform() elif v == MONTH: @@ -1000,6 +1006,10 @@ def project(self, name: str, pred: BoundPredicate) -> UnboundPredicate | None: def strict_project(self, name: str, pred: BoundPredicate) -> UnboundPredicate | None: return None + def __str__(self) -> str: + """Return the original transform name so it round-trips through serialization.""" + return self._transform + def __repr__(self) -> str: """Return the string representation of the UnknownTransform class.""" return f"UnknownTransform(transform={repr(self._transform)})" diff --git a/tests/table/test_partitioning.py b/tests/table/test_partitioning.py index b150fc2f67..45099a58d8 100644 --- a/tests/table/test_partitioning.py +++ b/tests/table/test_partitioning.py @@ -273,6 +273,48 @@ def test_deserialize_partition_field_empty_source_ids_rejected() -> None: PartitionField.model_validate_json(json_partition_spec) +def test_deserialize_partition_field_multi_arg() -> None: + import json as json_lib + + from pyiceberg.transforms import UnknownTransform + + json_partition_spec = """{"source-ids": [1, 2], "field-id": 1000, "transform": "bucket[4]", "name": "multi_bucket"}""" + field = PartitionField.model_validate_json(json_partition_spec) + + # v3 readers must read tables with multi-argument transforms, treating them as unknown + assert isinstance(field.transform, UnknownTransform) + assert field.source_id == 1 + assert field.source_ids == [1, 2] + + # the field must round-trip: source-ids only, with the original transform name + serialized = json_lib.loads(field.model_dump_json()) + assert serialized["source-ids"] == [1, 2] + assert "source-id" not in serialized + assert serialized["transform"] == "bucket[4]" + + +def test_serialize_partition_field_single_source_id_only() -> None: + import json as json_lib + + json_partition_spec = """{"source-ids": [1], "field-id": 1000, "transform": "truncate[19]", "name": "str_truncate"}""" + field = PartitionField.model_validate_json(json_partition_spec) + serialized = json_lib.loads(field.model_dump_json()) + assert serialized["source-id"] == 1 + assert "source-ids" not in serialized + + +def test_partition_type_with_multi_arg_field() -> None: + from pyiceberg.types import StringType + + schema = Schema(NestedField(1, "a", IntegerType()), NestedField(2, "b", IntegerType())) + field = PartitionField.model_validate_json( + """{"source-ids": [1, 2], "field-id": 1000, "transform": "bucket[4]", "name": "m"}""" + ) + spec = PartitionSpec(field) + struct = spec.partition_type(schema) + assert struct.fields[0].field_type == StringType() + + def test_incompatible_source_column_not_found() -> None: schema = Schema(NestedField(1, "foo", IntegerType()), NestedField(2, "bar", IntegerType())) diff --git a/tests/table/test_sorting.py b/tests/table/test_sorting.py index 5f7f5d016e..26f7e1949a 100644 --- a/tests/table/test_sorting.py +++ b/tests/table/test_sorting.py @@ -175,3 +175,20 @@ def test_incompatible_transform_source_type() -> None: sort_order.check_compatible(schema) assert "Invalid source field foo with type int for transform: year" in str(exc.value) + + +def test_deserialize_sort_field_multi_arg() -> None: + from pyiceberg.transforms import UnknownTransform + + payload = '{"source-ids":[19,20],"transform":"bucket[4]","direction":"asc","null-order":"nulls-first"}' + field = SortField.model_validate_json(payload) + + # v3 readers must read tables with multi-argument transforms, treating them as unknown + assert isinstance(field.transform, UnknownTransform) + assert field.source_id == 19 + assert field.source_ids == [19, 20] + + serialized = json.loads(field.model_dump_json()) + assert serialized["source-ids"] == [19, 20] + assert "source-id" not in serialized + assert serialized["transform"] == "bucket[4]" diff --git a/tests/test_transforms.py b/tests/test_transforms.py index d296fcdb21..03f3ba39d3 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -574,6 +574,22 @@ def test_unknown_transform_str() -> None: assert str(UnknownTransform("unknown")) == "unknown" +def test_unknown_transform_str_preserves_original_name() -> None: + # serializing metadata with an unknown transform must not rewrite its name + assert str(UnknownTransform("zorder")) == "zorder" + assert str(UnknownTransform("bucketv2[4]")) == "bucketv2[4]" + + +def test_parse_transform_unknown_with_known_prefix() -> None: + # unknown transforms that share a prefix with known ones must not fail parsing + from pyiceberg.transforms import parse_transform + + for name in ("bucketv2[4]", "truncatev2[8]", "bucket", "truncate[x]"): + transform = parse_transform(name) + assert isinstance(transform, UnknownTransform), name + assert str(transform) == name + + def test_unknown_transform_repr() -> None: assert repr(UnknownTransform("unknown")) == "UnknownTransform(transform='unknown')" From 5b55e16ba19e37a2ba284038744d434d6bfaa99c Mon Sep 17 00:00:00 2001 From: Noritaka Sekiyama Date: Thu, 9 Jul 2026 09:08:27 +0900 Subject: [PATCH 2/4] Address review: require transform for multi-argument fields A multi-argument field without a transform previously fabricated UnknownTransform('None') and masked the missing required field; raise a clear error instead. Assert explicitly that a single-element source-ids is normalized onto source-id, and only use the list form of __str__ for genuinely multi-argument fields. --- pyiceberg/partitioning.py | 9 +++++++-- pyiceberg/table/sorting.py | 4 +++- tests/table/test_partitioning.py | 8 ++++++++ tests/table/test_sorting.py | 6 ++++++ 4 files changed, 24 insertions(+), 3 deletions(-) diff --git a/pyiceberg/partitioning.py b/pyiceberg/partitioning.py index d9917931c6..7ac59a9ba1 100644 --- a/pyiceberg/partitioning.py +++ b/pyiceberg/partitioning.py @@ -117,9 +117,11 @@ def map_source_ids_onto_source_id(cls, data: Any) -> Any: if len(source_ids) == 0: raise ValueError("Empty source-ids is not allowed") if len(source_ids) > 1: + if data.get("transform") is None: + raise ValueError("Transform is required for a multi-argument field") # Multi-argument transforms cannot be evaluated; per the spec, v3 readers # must read tables with such transforms, ignoring them - data["transform"] = UnknownTransform(transform=str(data.get("transform"))) + data["transform"] = UnknownTransform(transform=str(data["transform"])) else: data.pop("source-ids", None) data["source-id"] = source_ids[0] @@ -140,7 +142,10 @@ def _serialize_source_ids(self, handler: Any) -> Any: def __str__(self) -> str: """Return the string representation of the PartitionField class.""" - sources = ", ".join(str(s) for s in self.source_ids) if self.source_ids else self.source_id + if self.source_ids is not None and len(self.source_ids) > 1: + sources = ", ".join(str(s) for s in self.source_ids) + else: + sources = str(self.source_id) return f"{self.field_id}: {self.name}: {self.transform}({sources})" diff --git a/pyiceberg/table/sorting.py b/pyiceberg/table/sorting.py index 08aeec016e..e6e7b219b2 100644 --- a/pyiceberg/table/sorting.py +++ b/pyiceberg/table/sorting.py @@ -108,9 +108,11 @@ def map_source_ids_onto_source_id(cls, data: Any) -> Any: if len(source_ids) == 0: raise ValueError("Empty source-ids is not allowed") if len(source_ids) > 1: + if data.get("transform") is None: + raise ValueError("Transform is required for a multi-argument field") # Multi-argument transforms cannot be evaluated; per the spec, v3 readers # must read tables with such transforms, ignoring them - data["transform"] = UnknownTransform(transform=str(data.get("transform"))) + data["transform"] = UnknownTransform(transform=str(data["transform"])) else: data.pop("source-ids", None) data["source-id"] = source_ids[0] diff --git a/tests/table/test_partitioning.py b/tests/table/test_partitioning.py index 45099a58d8..168e11853c 100644 --- a/tests/table/test_partitioning.py +++ b/tests/table/test_partitioning.py @@ -301,6 +301,8 @@ def test_serialize_partition_field_single_source_id_only() -> None: serialized = json_lib.loads(field.model_dump_json()) assert serialized["source-id"] == 1 assert "source-ids" not in serialized + # a single-element source-ids is normalized onto source-id + assert field.source_ids is None def test_partition_type_with_multi_arg_field() -> None: @@ -346,3 +348,9 @@ def test_incompatible_transform_source_type() -> None: spec.check_compatible(schema) assert "Invalid source field foo with type int for transform: year" in str(exc.value) + + +def test_deserialize_partition_field_multi_arg_requires_transform() -> None: + json_partition_spec = """{"source-ids": [1, 2], "field-id": 1000, "name": "m"}""" + with pytest.raises(Exception, match="Transform is required for a multi-argument field"): + PartitionField.model_validate_json(json_partition_spec) diff --git a/tests/table/test_sorting.py b/tests/table/test_sorting.py index 26f7e1949a..8388d0bcd7 100644 --- a/tests/table/test_sorting.py +++ b/tests/table/test_sorting.py @@ -192,3 +192,9 @@ def test_deserialize_sort_field_multi_arg() -> None: assert serialized["source-ids"] == [19, 20] assert "source-id" not in serialized assert serialized["transform"] == "bucket[4]" + + +def test_deserialize_sort_field_multi_arg_requires_transform() -> None: + payload = '{"source-ids":[19,20],"direction":"asc","null-order":"nulls-first"}' + with pytest.raises(Exception, match="Transform is required for a multi-argument field"): + SortField.model_validate_json(payload) From a4a7360a10768c7157a67750ba88b3f1135ab3ff Mon Sep 17 00:00:00 2001 From: Noritaka Sekiyama Date: Sat, 11 Jul 2026 08:56:31 +0900 Subject: [PATCH 3/4] Address review: update SortField.__str__ and drop underscore keys from serializers --- pyiceberg/partitioning.py | 2 -- pyiceberg/table/sorting.py | 7 ++++--- tests/table/test_partitioning.py | 2 ++ tests/table/test_sorting.py | 2 ++ 4 files changed, 8 insertions(+), 5 deletions(-) diff --git a/pyiceberg/partitioning.py b/pyiceberg/partitioning.py index 7ac59a9ba1..1e09f38cc3 100644 --- a/pyiceberg/partitioning.py +++ b/pyiceberg/partitioning.py @@ -134,10 +134,8 @@ def _serialize_source_ids(self, handler: Any) -> Any: # multi-argument transforms write only source-ids if self.source_ids is not None and len(self.source_ids) > 1: serialized.pop("source-id", None) - serialized.pop("source_id", None) else: serialized.pop("source-ids", None) - serialized.pop("source_ids", None) return serialized def __str__(self) -> str: diff --git a/pyiceberg/table/sorting.py b/pyiceberg/table/sorting.py index e6e7b219b2..79d5a09f94 100644 --- a/pyiceberg/table/sorting.py +++ b/pyiceberg/table/sorting.py @@ -125,10 +125,8 @@ def _serialize_source_ids(self, handler: Any) -> Any: # multi-argument transforms write only source-ids if self.source_ids is not None and len(self.source_ids) > 1: serialized.pop("source-id", None) - serialized.pop("source_id", None) else: serialized.pop("source-ids", None) - serialized.pop("source_ids", None) return serialized source_id: int = Field(alias="source-id") @@ -147,8 +145,11 @@ def __str__(self) -> str: if isinstance(self.transform, IdentityTransform): # In the case of an identity transform, we can omit the transform return f"{self.source_id} {self.direction} {self.null_order}" + if self.source_ids is not None and len(self.source_ids) > 1: + sources = ", ".join(str(s) for s in self.source_ids) else: - return f"{self.transform}({self.source_id}) {self.direction} {self.null_order}" + sources = str(self.source_id) + return f"{self.transform}({sources}) {self.direction} {self.null_order}" INITIAL_SORT_ORDER_ID = 1 diff --git a/tests/table/test_partitioning.py b/tests/table/test_partitioning.py index 168e11853c..316ceda01b 100644 --- a/tests/table/test_partitioning.py +++ b/tests/table/test_partitioning.py @@ -292,6 +292,8 @@ def test_deserialize_partition_field_multi_arg() -> None: assert "source-id" not in serialized assert serialized["transform"] == "bucket[4]" + assert str(field) == "1000: multi_bucket: bucket[4](1, 2)" + def test_serialize_partition_field_single_source_id_only() -> None: import json as json_lib diff --git a/tests/table/test_sorting.py b/tests/table/test_sorting.py index 8388d0bcd7..4569147f22 100644 --- a/tests/table/test_sorting.py +++ b/tests/table/test_sorting.py @@ -193,6 +193,8 @@ def test_deserialize_sort_field_multi_arg() -> None: assert "source-id" not in serialized assert serialized["transform"] == "bucket[4]" + assert str(field) == "bucket[4](19, 20) ASC NULLS FIRST" + def test_deserialize_sort_field_multi_arg_requires_transform() -> None: payload = '{"source-ids":[19,20],"direction":"asc","null-order":"nulls-first"}' From 1981aac5ba3d5e504e4d1028f7b1a3f2089f8a93 Mon Sep 17 00:00:00 2001 From: Noritaka Sekiyama Date: Fri, 21 Aug 2026 13:46:14 +0900 Subject: [PATCH 4/4] Address review: consolidate source-id handling into a shared mixin Extract the source-id/source-ids pair, its before-validator, its serializer and a transform_arguments accessor into TransformSourceMixin, inherited by PartitionField and SortField, replacing the near-identical copies in both. The validator is now flat and treats source-ids as authoritative, so a field carrying both keys (which the spec never writes) is still normalized. Before, the whole block was skipped for such input and a multi-argument transform stayed evaluable instead of becoming UnknownTransform. --- pyiceberg/partitioning.py | 46 +++------------------ pyiceberg/table/sorting.py | 46 +++------------------ pyiceberg/transforms.py | 71 +++++++++++++++++++++++++++++++- tests/table/test_partitioning.py | 54 ++++++++++++++++++++++++ tests/table/test_sorting.py | 40 ++++++++++++++++++ 5 files changed, 173 insertions(+), 84 deletions(-) diff --git a/pyiceberg/partitioning.py b/pyiceberg/partitioning.py index 1e09f38cc3..e2dc35897c 100644 --- a/pyiceberg/partitioning.py +++ b/pyiceberg/partitioning.py @@ -29,8 +29,6 @@ Field, PlainSerializer, WithJsonSchema, - model_serializer, - model_validator, ) from pyiceberg.exceptions import ValidationError @@ -42,6 +40,7 @@ IdentityTransform, MonthTransform, Transform, + TransformSourceMixin, TruncateTransform, UnknownTransform, VoidTransform, @@ -67,18 +66,17 @@ PARTITION_FIELD_ID_START: int = 1000 -class PartitionField(IcebergBaseModel): +class PartitionField(TransformSourceMixin): """PartitionField represents how one partition value is derived from the source column via transformation. Attributes: - source_id(int): The source column id of table's schema. field_id(int): The partition field id across all the table partition specs. transform(Transform): The transform used to produce partition values from source column. name(str): The name of this partition field. + + The source columns are carried by `TransformSourceMixin`. """ - source_id: int = Field(alias="source-id") - source_ids: list[int] | None = Field(alias="source-ids", default=None, repr=False) field_id: int = Field(alias="field-id") transform: Annotated[ # type: ignore Transform, @@ -107,43 +105,9 @@ def __init__( super().__init__(**data) - @model_validator(mode="before") - @classmethod - def map_source_ids_onto_source_id(cls, data: Any) -> Any: - if isinstance(data, dict): - if "source-id" not in data and "source-ids" in data: - source_ids = data["source-ids"] - if isinstance(source_ids, list): - if len(source_ids) == 0: - raise ValueError("Empty source-ids is not allowed") - if len(source_ids) > 1: - if data.get("transform") is None: - raise ValueError("Transform is required for a multi-argument field") - # Multi-argument transforms cannot be evaluated; per the spec, v3 readers - # must read tables with such transforms, ignoring them - data["transform"] = UnknownTransform(transform=str(data["transform"])) - else: - data.pop("source-ids", None) - data["source-id"] = source_ids[0] - return data - - @model_serializer(mode="wrap") - def _serialize_source_ids(self, handler: Any) -> Any: - serialized = handler(self) - # Per the spec, single-argument transforms write only source-id and - # multi-argument transforms write only source-ids - if self.source_ids is not None and len(self.source_ids) > 1: - serialized.pop("source-id", None) - else: - serialized.pop("source-ids", None) - return serialized - def __str__(self) -> str: """Return the string representation of the PartitionField class.""" - if self.source_ids is not None and len(self.source_ids) > 1: - sources = ", ".join(str(s) for s in self.source_ids) - else: - sources = str(self.source_id) + sources = ", ".join(str(source_id) for source_id in self.transform_arguments) return f"{self.field_id}: {self.name}: {self.transform}({sources})" diff --git a/pyiceberg/table/sorting.py b/pyiceberg/table/sorting.py index 79d5a09f94..f7b59649c6 100644 --- a/pyiceberg/table/sorting.py +++ b/pyiceberg/table/sorting.py @@ -24,13 +24,12 @@ Field, PlainSerializer, WithJsonSchema, - model_serializer, model_validator, ) from pyiceberg.exceptions import ValidationError from pyiceberg.schema import Schema -from pyiceberg.transforms import IdentityTransform, Transform, UnknownTransform, parse_transform +from pyiceberg.transforms import IdentityTransform, Transform, TransformSourceMixin, parse_transform from pyiceberg.typedef import IcebergBaseModel from pyiceberg.types import IcebergType @@ -61,11 +60,12 @@ def __repr__(self) -> str: return f"NullOrder.{self.name}" -class SortField(IcebergBaseModel): +class SortField(TransformSourceMixin): """Sort order field. + The source columns are carried by `TransformSourceMixin`. + Args: - source_id (int): Source column id from the table’s schema. transform (str): Transform that is used to produce values to be sorted on from the source column. This is the same transform as described in partition transforms. direction (SortDirection): Sort direction, that can only be either asc or desc. @@ -98,39 +98,6 @@ def set_null_order(cls, values: dict[str, Any]) -> dict[str, Any]: values["null-order"] = NullOrder.NULLS_FIRST if values["direction"] == SortDirection.ASC else NullOrder.NULLS_LAST return values - @model_validator(mode="before") - @classmethod - def map_source_ids_onto_source_id(cls, data: Any) -> Any: - if isinstance(data, dict): - if "source-id" not in data and "source-ids" in data: - source_ids = data["source-ids"] - if isinstance(source_ids, list): - if len(source_ids) == 0: - raise ValueError("Empty source-ids is not allowed") - if len(source_ids) > 1: - if data.get("transform") is None: - raise ValueError("Transform is required for a multi-argument field") - # Multi-argument transforms cannot be evaluated; per the spec, v3 readers - # must read tables with such transforms, ignoring them - data["transform"] = UnknownTransform(transform=str(data["transform"])) - else: - data.pop("source-ids", None) - data["source-id"] = source_ids[0] - return data - - @model_serializer(mode="wrap") - def _serialize_source_ids(self, handler: Any) -> Any: - serialized = handler(self) - # Per the spec, single-argument transforms write only source-id and - # multi-argument transforms write only source-ids - if self.source_ids is not None and len(self.source_ids) > 1: - serialized.pop("source-id", None) - else: - serialized.pop("source-ids", None) - return serialized - - source_id: int = Field(alias="source-id") - source_ids: list[int] | None = Field(alias="source-ids", default=None, repr=False) transform: Annotated[ # type: ignore Transform, BeforeValidator(parse_transform), @@ -145,10 +112,7 @@ def __str__(self) -> str: if isinstance(self.transform, IdentityTransform): # In the case of an identity transform, we can omit the transform return f"{self.source_id} {self.direction} {self.null_order}" - if self.source_ids is not None and len(self.source_ids) > 1: - sources = ", ".join(str(s) for s in self.source_ids) - else: - sources = str(self.source_id) + sources = ", ".join(str(source_id) for source_id in self.transform_arguments) return f"{self.transform}({sources}) {self.direction} {self.null_order}" diff --git a/pyiceberg/transforms.py b/pyiceberg/transforms.py index ab475ab51c..f0c7f4c05a 100644 --- a/pyiceberg/transforms.py +++ b/pyiceberg/transforms.py @@ -29,7 +29,7 @@ from uuid import UUID import mmh3 -from pydantic import Field, PositiveInt, PrivateAttr +from pydantic import Field, PositiveInt, PrivateAttr, model_serializer, model_validator from pyiceberg.exceptions import NotInstalledError, ValidationError from pyiceberg.expressions import ( @@ -67,7 +67,7 @@ TimestampLiteral, literal, ) -from pyiceberg.typedef import IcebergRootModel, L +from pyiceberg.typedef import IcebergBaseModel, IcebergRootModel, L from pyiceberg.types import ( BinaryType, DateType, @@ -1054,6 +1054,73 @@ def pyarrow_transform(self, source: IcebergType) -> "Callable[[pa.Array], pa.Arr return lambda arr: pa.nulls(len(arr), type=arr.type) +class TransformSourceMixin(IcebergBaseModel): + """Shared `source-id` and `source-ids` handling for fields that apply a transform to source columns. + + Both partition fields and sort fields carry this pair, and the spec writes only one of the + two: `source-id` for a transform with a single argument, `source-ids` for a multi-argument + transform. This mixin owns reading, serializing and reporting them, so that neither field + type has to reach into the raw keys. + + Attributes: + source_id(int): The source column id of the table's schema. + source_ids(list[int] | None): The source column ids of a multi-argument transform. + """ + + source_id: int = Field(alias="source-id") + source_ids: list[int] | None = Field(alias="source-ids", default=None, repr=False) + + @property + def transform_arguments(self) -> list[int]: + """Return the source column ids that the transform is applied to.""" + source_ids = self.source_ids + if source_ids is not None and len(source_ids) > 1: + return list(source_ids) + return [self.source_id] + + @property + def is_multi_argument(self) -> bool: + """Return True if the transform takes more than one source column.""" + return len(self.transform_arguments) > 1 + + @model_validator(mode="before") + @classmethod + def map_source_ids_onto_source_id(cls, data: Any) -> Any: + if not isinstance(data, dict) or "source-ids" not in data: + return data + + source_ids = data["source-ids"] + if not isinstance(source_ids, list): + return data + if len(source_ids) == 0: + raise ValueError("Empty source-ids is not allowed") + + # The spec writes only one of the two keys, so a source-id next to source-ids comes + # from a non-conformant writer; source-ids is the one that carries the arity + data["source-id"] = source_ids[0] + if len(source_ids) == 1: + data.pop("source-ids", None) + return data + + if data.get("transform") is None: + raise ValueError("Transform is required for a multi-argument field") + # Multi-argument transforms cannot be evaluated; per the spec, v3 readers + # must read tables with such transforms, ignoring them + data["transform"] = UnknownTransform(transform=str(data["transform"])) + return data + + @model_serializer(mode="wrap") + def _serialize_source_ids(self, handler: Any) -> Any: + serialized = handler(self) + # Per the spec, single-argument transforms write only source-id and + # multi-argument transforms write only source-ids + if self.is_multi_argument: + serialized.pop("source-id", None) + else: + serialized.pop("source-ids", None) + return serialized + + def _truncate_number( name: str, pred: BoundLiteralPredicate, transform: Callable[[Any | None], Any | None] ) -> UnboundPredicate | None: diff --git a/tests/table/test_partitioning.py b/tests/table/test_partitioning.py index 316ceda01b..9da56d7589 100644 --- a/tests/table/test_partitioning.py +++ b/tests/table/test_partitioning.py @@ -356,3 +356,57 @@ def test_deserialize_partition_field_multi_arg_requires_transform() -> None: json_partition_spec = """{"source-ids": [1, 2], "field-id": 1000, "name": "m"}""" with pytest.raises(Exception, match="Transform is required for a multi-argument field"): PartitionField.model_validate_json(json_partition_spec) + + +def test_deserialize_partition_field_both_source_id_and_source_ids_multi_arg() -> None: + import json as json_lib + + from pyiceberg.transforms import UnknownTransform + + # a non-conformant writer emitted both keys; source-ids carries the arity, so it wins + json_partition_spec = ( + """{"source-id": 9, "source-ids": [1, 2], "field-id": 1000, "transform": "bucket[4]", "name": "multi_bucket"}""" + ) + field = PartitionField.model_validate_json(json_partition_spec) + + assert isinstance(field.transform, UnknownTransform) + assert field.source_id == 1 + assert field.source_ids == [1, 2] + + serialized = json_lib.loads(field.model_dump_json()) + assert serialized["source-ids"] == [1, 2] + assert "source-id" not in serialized + + +def test_deserialize_partition_field_both_source_id_and_source_ids_single_arg() -> None: + import json as json_lib + + json_partition_spec = ( + """{"source-id": 9, "source-ids": [1], "field-id": 1000, "transform": "truncate[19]", "name": "str_truncate"}""" + ) + field = PartitionField.model_validate_json(json_partition_spec) + + assert field.source_id == 1 + assert field.source_ids is None + + serialized = json_lib.loads(field.model_dump_json()) + assert serialized["source-id"] == 1 + assert "source-ids" not in serialized + + +def test_deserialize_partition_field_empty_source_ids_rejected_next_to_source_id() -> None: + json_partition_spec = """{"source-id": 1, "source-ids": [], "field-id": 1000, "transform": "identity", "name": "x"}""" + with pytest.raises(Exception, match="Empty source-ids is not allowed"): + PartitionField.model_validate_json(json_partition_spec) + + +def test_partition_field_transform_arguments() -> None: + single = PartitionField(source_id=1, field_id=1000, transform=TruncateTransform(width=19), name="str_truncate") + assert single.transform_arguments == [1] + assert single.is_multi_argument is False + + multi = PartitionField.model_validate_json( + """{"source-ids": [1, 2], "field-id": 1000, "transform": "bucket[4]", "name": "m"}""" + ) + assert multi.transform_arguments == [1, 2] + assert multi.is_multi_argument is True diff --git a/tests/table/test_sorting.py b/tests/table/test_sorting.py index 4569147f22..8e6c6164b5 100644 --- a/tests/table/test_sorting.py +++ b/tests/table/test_sorting.py @@ -200,3 +200,43 @@ def test_deserialize_sort_field_multi_arg_requires_transform() -> None: payload = '{"source-ids":[19,20],"direction":"asc","null-order":"nulls-first"}' with pytest.raises(Exception, match="Transform is required for a multi-argument field"): SortField.model_validate_json(payload) + + +def test_deserialize_sort_field_both_source_id_and_source_ids_multi_arg() -> None: + from pyiceberg.transforms import UnknownTransform + + # a non-conformant writer emitted both keys; source-ids carries the arity, so it wins + payload = '{"source-id":9,"source-ids":[19,20],"transform":"bucket[4]","direction":"asc","null-order":"nulls-first"}' + field = SortField.model_validate_json(payload) + + assert isinstance(field.transform, UnknownTransform) + assert field.source_id == 19 + assert field.source_ids == [19, 20] + + serialized = json.loads(field.model_dump_json()) + assert serialized["source-ids"] == [19, 20] + assert "source-id" not in serialized + + +def test_deserialize_sort_field_both_source_id_and_source_ids_single_arg() -> None: + payload = '{"source-id":9,"source-ids":[19],"transform":"identity","direction":"asc","null-order":"nulls-first"}' + field = SortField.model_validate_json(payload) + + assert field.source_id == 19 + assert field.source_ids is None + + serialized = json.loads(field.model_dump_json()) + assert serialized["source-id"] == 19 + assert "source-ids" not in serialized + + +def test_sort_field_transform_arguments() -> None: + single = SortField(source_id=19, transform=BucketTransform(num_buckets=4), null_order=NullOrder.NULLS_FIRST) + assert single.transform_arguments == [19] + assert single.is_multi_argument is False + + multi = SortField.model_validate_json( + '{"source-ids":[19,20],"transform":"bucket[4]","direction":"asc","null-order":"nulls-first"}' + ) + assert multi.transform_arguments == [19, 20] + assert multi.is_multi_argument is True