Read tables with multi-argument transforms as unknown transforms

pyiceberg/partitioning.py

@@ -29,7 +29,6 @@
     Field,
     PlainSerializer,
     WithJsonSchema,
-    model_validator,
 )
 
 from pyiceberg.exceptions import ValidationError
@@ -41,6 +40,7 @@
     IdentityTransform,
     MonthTransform,
     Transform,
+    TransformSourceMixin,
     TruncateTransform,
     UnknownTransform,
     VoidTransform,
@@ -68,17 +68,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")
     field_id: int = Field(alias="field-id")
     transform: Annotated[  # type: ignore
         Transform,
@@ -107,25 +107,10 @@ 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-ids" in data:
-                if "source-id" in data:
-                    raise ValueError("source-id and source-ids are mutually exclusive")
-                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:
-                        raise ValueError("Multi argument transforms are not yet supported")
-                    data["source-id"] = source_ids[0]
-        return data
-
     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(source_id) for source_id in self.transform_arguments)
+        return f"{self.field_id}: {self.name}: {self.transform}({sources})"
 
 
 class PartitionSpec(IcebergBaseModel):
@@ -206,7 +191,7 @@ def compatible_with(self, other: PartitionSpec) -> bool:
         if len(self.fields) != len(other.fields):
             return False
         return all(
-            this_field.source_id == that_field.source_id
+            this_field.transform_arguments == that_field.transform_arguments
             and this_field.transform == that_field.transform
             and this_field.name == that_field.name
             for this_field, that_field in zip(self.fields, other.fields, strict=True)

pyiceberg/table/sorting.py

@@ -29,7 +29,7 @@
 
 from pyiceberg.exceptions import ValidationError
 from pyiceberg.schema import Schema
-from pyiceberg.transforms import IdentityTransform, Transform, parse_transform
+from pyiceberg.transforms import IdentityTransform, Transform, TransformSourceMixin, parse_transform
 from pyiceberg.typedef import IcebergBaseModel
 from pyiceberg.types import IcebergType
 
@@ -60,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.
@@ -97,23 +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-ids" in data:
-                if "source-id" in data:
-                    raise ValueError("source-id and source-ids are mutually exclusive")
-                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:
-                        raise ValueError("Multi argument transforms are not yet supported")
-                    data["source-id"] = source_ids[0]
-        return data
-
-    source_id: int = Field(alias="source-id")
     transform: Annotated[  # type: ignore
         Transform,
         BeforeValidator(parse_transform),
@@ -128,8 +112,8 @@ 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}"
-        else:
-            return f"{self.transform}({self.source_id}) {self.direction} {self.null_order}"
+        sources = ", ".join(str(source_id) for source_id in self.transform_arguments)
+        return f"{self.transform}({sources}) {self.direction} {self.null_order}"
 
 
 INITIAL_SORT_ORDER_ID = 1

pyiceberg/transforms.py

@@ -29,9 +29,9 @@
 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
+from pyiceberg.exceptions import NotInstalledError, ValidationError
 from pyiceberg.expressions import (
     BoundEqualTo,
     BoundGreaterThan,
@@ -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,
@@ -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,10 +1006,24 @@ 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)})"
 
+    def __eq__(self, other: Any) -> bool:
+        """Compare the preserved transform name, since every unknown transform shares one root value."""
+        if isinstance(other, UnknownTransform):
+            return self._transform == other._transform
+        return False
+
+    def __hash__(self) -> int:
+        """Hash the preserved transform name so distinct unknown transforms do not collide."""
+        return hash((self.root, self._transform))
+
     def pyarrow_transform(self, source: IcebergType) -> "Callable[[pa.Array], pa.Array]":
         raise NotImplementedError()
 
@@ -1044,6 +1064,74 @@ 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
+
+        if "source-id" in data:
+            raise ValueError("source-id and source-ids are mutually exclusive")
+
+        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")
+
+        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:

pyiceberg/utils/parsing.py

@@ -28,7 +28,8 @@ class ParseNumberFromBrackets:
 
     def __init__(self, prefix: str):
         self.prefix = prefix
-        self.regex = re.compile(rf"{prefix}\[(\d+)\]")
+        # anchored: a name such as truncate[8]v2 is a different transform, not truncate[8]
+        self.regex = re.compile(rf"^{prefix}\[(\d+)\]$")
 
     def match(self, str_repr: str) -> int:
         matches = self.regex.search(str_repr)

tests/table/test_partitioning.py

@@ -105,6 +105,31 @@ def test_partition_compatible_with() -> None:
     assert not lhs.compatible_with(rhs)
 
 
+def test_partition_compatible_with_compares_every_source_id() -> None:
+    # multi-argument fields differ in their later source ids, which the first one cannot show
+    lhs = PartitionSpec.model_validate(
+        {"spec-id": 0, "fields": [{"field-id": 1000, "name": "p", "source-ids": [1, 2], "transform": "zorder(a,b)"}]}
+    )
+    rhs = PartitionSpec.model_validate(
+        {"spec-id": 0, "fields": [{"field-id": 1000, "name": "p", "source-ids": [1, 3], "transform": "zorder(a,c)"}]}
+    )
+
+    assert not lhs.compatible_with(rhs)
+    assert lhs.compatible_with(lhs)
+
+
+def test_partition_compatible_with_compares_unknown_transform_names() -> None:
+    # same source ids, different transform: the names are all that separates them
+    lhs = PartitionSpec.model_validate(
+        {"spec-id": 0, "fields": [{"field-id": 1000, "name": "p", "source-ids": [1, 2], "transform": "zorder(a,b)"}]}
+    )
+    rhs = PartitionSpec.model_validate(
+        {"spec-id": 0, "fields": [{"field-id": 1000, "name": "p", "source-ids": [1, 2], "transform": "somethingelse(a,b)"}]}
+    )
+
+    assert not lhs.compatible_with(rhs)
+
+
 def test_unpartitioned() -> None:
     assert len(UNPARTITIONED_PARTITION_SPEC.fields) == 0
     assert UNPARTITIONED_PARTITION_SPEC.is_unpartitioned()
@@ -279,6 +304,52 @@ def test_deserialize_partition_field_source_id_and_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]"
+
+    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
+
+    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()))
 
@@ -310,3 +381,21 @@ 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)
+
+
+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