|
14 | 14 | # KIND, either express or implied. See the License for the
|
15 | 15 | # specific language governing permissions and limitations
|
16 | 16 | # under the License.
|
| 17 | +import datetime |
| 18 | +from decimal import Decimal |
| 19 | +from typing import Any |
| 20 | +from uuid import UUID |
| 21 | + |
| 22 | +import pytest |
| 23 | + |
17 | 24 | from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionField, PartitionSpec
|
18 | 25 | from pyiceberg.schema import Schema
|
19 | 26 | from pyiceberg.transforms import BucketTransform, IdentityTransform, TruncateTransform
|
20 | 27 | from pyiceberg.typedef import Record
|
21 | 28 | from pyiceberg.types import (
|
| 29 | + BinaryType, |
| 30 | + DateType, |
| 31 | + DecimalType, |
| 32 | + FixedType, |
22 | 33 | IntegerType,
|
| 34 | + LongType, |
23 | 35 | NestedField,
|
| 36 | + PrimitiveType, |
24 | 37 | StringType,
|
25 | 38 | StructType,
|
| 39 | + TimestampType, |
| 40 | + TimestamptzType, |
| 41 | + TimeType, |
| 42 | + UUIDType, |
26 | 43 | )
|
27 | 44 |
|
28 | 45 |
|
@@ -153,6 +170,29 @@ def test_partition_type(table_schema_simple: Schema) -> None:
|
153 | 170 | )
|
154 | 171 |
|
155 | 172 |
|
| 173 | +@pytest.mark.parametrize( |
| 174 | + "source_type, value", |
| 175 | + [ |
| 176 | + (IntegerType(), 22), |
| 177 | + (LongType(), 22), |
| 178 | + (DecimalType(5, 9), Decimal(19.25)), |
| 179 | + (DateType(), datetime.date(1925, 5, 22)), |
| 180 | + (TimeType(), datetime.time(19, 25, 00)), |
| 181 | + (TimestampType(), datetime.datetime(19, 5, 1, 22, 1, 1)), |
| 182 | + (TimestamptzType(), datetime.datetime(19, 5, 1, 22, 1, 1, tzinfo=datetime.timezone.utc)), |
| 183 | + (StringType(), "abc"), |
| 184 | + (UUIDType(), UUID("12345678-1234-5678-1234-567812345678").bytes), |
| 185 | + (FixedType(5), 'b"\x8e\xd1\x87\x01"'), |
| 186 | + (BinaryType(), b"\x8e\xd1\x87\x01"), |
| 187 | + ], |
| 188 | +) |
| 189 | +def test_bucketing_function(source_type: PrimitiveType, value: Any) -> None: |
| 190 | + bucket = BucketTransform(2) # type: ignore |
| 191 | + import pyarrow as pa |
| 192 | + |
| 193 | + assert bucket.transform(source_type)(value) == bucket.pyarrow_transform(source_type)(pa.array([value])).to_pylist()[0] |
| 194 | + |
| 195 | + |
156 | 196 | def test_deserialize_partition_field_v2() -> None:
|
157 | 197 | json_partition_spec = """{"source-id": 1, "field-id": 1000, "transform": "truncate[19]", "name": "str_truncate"}"""
|
158 | 198 |
|
|
0 commit comments