diff --git a/pyiceberg/partitioning.py b/pyiceberg/partitioning.py index 3074c30ea1..a7887e0a44 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,36 @@ 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") + 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) + 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})" + 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})" class PartitionSpec(IcebergBaseModel): diff --git a/pyiceberg/table/sorting.py b/pyiceberg/table/sorting.py index 61f34c4780..e6e7b219b2 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,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") + 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) + 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 cd0d7cebcb..88e69d48ff 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: @@ -1001,6 +1007,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..168e11853c 100644 --- a/tests/table/test_partitioning.py +++ b/tests/table/test_partitioning.py @@ -273,6 +273,50 @@ 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 + # 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: + 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())) @@ -304,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 5f7f5d016e..8388d0bcd7 100644 --- a/tests/table/test_sorting.py +++ b/tests/table/test_sorting.py @@ -175,3 +175,26 @@ 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]" + + +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) 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')"