Skip to content

Commit

Permalink
Test transform function consistency for all transforms (#1573)
Browse files Browse the repository at this point in the history
I like this test from #1562, lets expand it to include all transforms
  • Loading branch information
kevinjqliu authored Jan 30, 2025
1 parent 5e4815a commit 361a407
Showing 1 changed file with 28 additions and 4 deletions.
32 changes: 28 additions & 4 deletions tests/table/test_partitioning.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,15 @@

from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionField, PartitionSpec
from pyiceberg.schema import Schema
from pyiceberg.transforms import BucketTransform, IdentityTransform, TruncateTransform
from pyiceberg.transforms import (
BucketTransform,
DayTransform,
HourTransform,
IdentityTransform,
MonthTransform,
TruncateTransform,
YearTransform,
)
from pyiceberg.typedef import Record
from pyiceberg.types import (
BinaryType,
Expand Down Expand Up @@ -186,11 +194,27 @@ def test_partition_type(table_schema_simple: Schema) -> None:
(BinaryType(), b"\x8e\xd1\x87\x01"),
],
)
def test_bucketing_function(source_type: PrimitiveType, value: Any) -> None:
bucket = BucketTransform(2) # type: ignore
def test_transform_consistency_with_pyarrow_transform(source_type: PrimitiveType, value: Any) -> None:
import pyarrow as pa

assert bucket.transform(source_type)(value) == bucket.pyarrow_transform(source_type)(pa.array([value])).to_pylist()[0]
all_transforms = [ # type: ignore
IdentityTransform(),
BucketTransform(10),
TruncateTransform(10),
YearTransform(),
MonthTransform(),
DayTransform(),
HourTransform(),
]
for t in all_transforms:
if t.can_transform(source_type):
try:
assert t.transform(source_type)(value) == t.pyarrow_transform(source_type)(pa.array([value])).to_pylist()[0]
except ValueError as e:
# Skipping unsupported feature
if "FeatureUnsupported => Unsupported data type for truncate transform" in str(e):
continue
raise


def test_deserialize_partition_field_v2() -> None:
Expand Down

0 comments on commit 361a407

Please sign in to comment.