diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index fef33a25c7..3a6bfa4143 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=protected-access,unused-argument,redefined-outer-name +import json import logging import os import tempfile @@ -63,6 +64,7 @@ ) from pyiceberg.expressions.literals import literal from pyiceberg.io import S3_RETRY_STRATEGY_IMPL, InputStream, OutputStream, load_file_io +from pyiceberg.io.fileformat import FileFormatFactory from pyiceberg.io.pyarrow import ( ICEBERG_SCHEMA, PYARROW_PARQUET_FIELD_ID_KEY, @@ -114,8 +116,11 @@ StructType, TimestampNanoType, TimestampType, + TimestamptzNanoType, TimestamptzType, TimeType, + UnknownType, + UUIDType, ) from tests.catalog.test_base import InMemoryCatalog from tests.conftest import UNIFIED_AWS_SESSION_PROPERTIES @@ -2408,6 +2413,49 @@ def test_stats_aggregator_physical_type_does_not_match_expected_raise_error( StatsAggregator(iceberg_type, physical_type_string) +def test_iceberg_types_write_spec_compliant_parquet_types(tmp_path: Path) -> None: + """Every Iceberg primitive type must be written with the Parquet types mandated by the spec.""" + expected_parquet_types: list[tuple[PrimitiveType, str, dict[str, Any]]] = [ + (BooleanType(), "BOOLEAN", {"Type": "None"}), + (IntegerType(), "INT32", {"Type": "None"}), + (LongType(), "INT64", {"Type": "None"}), + (FloatType(), "FLOAT", {"Type": "None"}), + (DoubleType(), "DOUBLE", {"Type": "None"}), + (DecimalType(9, 2), "INT32", {"Type": "Decimal", "precision": 9, "scale": 2}), + (DecimalType(18, 2), "INT64", {"Type": "Decimal", "precision": 18, "scale": 2}), + (DecimalType(38, 2), "FIXED_LEN_BYTE_ARRAY", {"Type": "Decimal", "precision": 38, "scale": 2}), + (DateType(), "INT32", {"Type": "Date"}), + (TimeType(), "INT64", {"Type": "Time", "isAdjustedToUTC": False, "timeUnit": "microseconds"}), + (TimestampType(), "INT64", {"Type": "Timestamp", "isAdjustedToUTC": False, "timeUnit": "microseconds"}), + (TimestamptzType(), "INT64", {"Type": "Timestamp", "isAdjustedToUTC": True, "timeUnit": "microseconds"}), + (TimestampNanoType(), "INT64", {"Type": "Timestamp", "isAdjustedToUTC": False, "timeUnit": "nanoseconds"}), + (TimestamptzNanoType(), "INT64", {"Type": "Timestamp", "isAdjustedToUTC": True, "timeUnit": "nanoseconds"}), + (StringType(), "BYTE_ARRAY", {"Type": "String"}), + (UUIDType(), "FIXED_LEN_BYTE_ARRAY", {"Type": "UUID"}), + (FixedType(16), "FIXED_LEN_BYTE_ARRAY", {"Type": "None"}), + (BinaryType(), "BYTE_ARRAY", {"Type": "None"}), + (UnknownType(), "INT32", {"Type": "Null"}), + ] + schema = Schema( + *[ + NestedField(field_id=field_id, name=str(iceberg_type), field_type=iceberg_type) + for field_id, (iceberg_type, _, _) in enumerate(expected_parquet_types, 1) + ] + ) + + file_path = str(tmp_path / "all_types.parquet") + writer = FileFormatFactory.get(FileFormat.PARQUET).create_writer(PyArrowFileIO().new_output(file_path), schema, {}) + writer.write(schema_to_pyarrow(schema).empty_table()) + writer.close() + + parquet_schema = pq.ParquetFile(file_path).schema + for pos, (iceberg_type, physical_type, logical_type) in enumerate(expected_parquet_types): + column = parquet_schema.column(pos) + assert column.physical_type == physical_type, f"{iceberg_type}: {column.physical_type} != {physical_type}" + column_logical_type = json.loads(column.logical_type.to_json()) + assert {key: column_logical_type.get(key) for key in logical_type} == logical_type, str(iceberg_type) + + def test_bin_pack_arrow_table(arrow_table_with_null: pa.Table) -> None: # default packs to 1 bin since the table is small bin_packed = bin_pack_arrow_table( diff --git a/tests/io/test_pyarrow_visitor.py b/tests/io/test_pyarrow_visitor.py index e98d76e262..5fa3ecbd1f 100644 --- a/tests/io/test_pyarrow_visitor.py +++ b/tests/io/test_pyarrow_visitor.py @@ -65,8 +65,11 @@ StructType, TimestampNanoType, TimestampType, + TimestamptzNanoType, TimestamptzType, TimeType, + UnknownType, + UUIDType, ) @@ -863,3 +866,110 @@ def test_expression_to_complementary_pyarrow( "or is_null(string_field, {nan_is_null=false})) or is_nan(double_field))>" ) assert repr(result) == expected_repr + + +PYARROW_TO_ICEBERG_SUPPORTED: list[tuple[pa.DataType, IcebergType]] = [ + (pa.null(), UnknownType()), + (pa.bool_(), BooleanType()), + (pa.int8(), IntegerType()), + (pa.int16(), IntegerType()), + (pa.int32(), IntegerType()), + (pa.int64(), LongType()), + (pa.uint8(), IntegerType()), + (pa.uint16(), IntegerType()), + (pa.uint32(), IntegerType()), + (pa.uint64(), LongType()), + (pa.float16(), FloatType()), + (pa.float32(), FloatType()), + (pa.float64(), DoubleType()), + (pa.decimal128(24, 2), DecimalType(24, 2)), + (pa.string(), StringType()), + (pa.large_string(), StringType()), + (pa.string_view(), StringType()), + (pa.binary(), BinaryType()), + (pa.large_binary(), BinaryType()), + (pa.binary_view(), BinaryType()), + (pa.binary(16), FixedType(16)), + (pa.date32(), DateType()), + (pa.time64("us"), TimeType()), + (pa.timestamp("s"), TimestampType()), + (pa.timestamp("ms"), TimestampType()), + (pa.timestamp("us"), TimestampType()), + (pa.timestamp("ns"), TimestampNanoType()), + (pa.timestamp("us", tz="UTC"), TimestamptzType()), + (pa.timestamp("ns", tz="UTC"), TimestamptzNanoType()), + (pa.uuid(), UUIDType()), + (pa.dictionary(pa.int32(), pa.string()), StringType()), + ( + pa.list_(pa.field("element", pa.int32(), nullable=True, metadata={"PARQUET:field_id": "1"})), + ListType(element_id=1, element_type=IntegerType(), element_required=False), + ), + ( + pa.large_list(pa.field("element", pa.int32(), nullable=True, metadata={"PARQUET:field_id": "1"})), + ListType(element_id=1, element_type=IntegerType(), element_required=False), + ), + ( + pa.list_(pa.field("element", pa.int32(), nullable=True, metadata={"PARQUET:field_id": "1"}), 3), + ListType(element_id=1, element_type=IntegerType(), element_required=False), + ), + ( + pa.struct([pa.field("x", pa.int32(), nullable=True, metadata={"PARQUET:field_id": "1"})]), + StructType(NestedField(1, "x", IntegerType(), required=False)), + ), + ( + pa.map_( + pa.field("key", pa.string(), nullable=False, metadata={"PARQUET:field_id": "1"}), + pa.field("value", pa.int32(), nullable=True, metadata={"PARQUET:field_id": "2"}), + ), + MapType(key_id=1, key_type=StringType(), value_id=2, value_type=IntegerType(), value_required=False), + ), +] + +PYARROW_TO_ICEBERG_UNSUPPORTED: list[tuple[pa.DataType, str]] = [ + (pa.decimal32(4, 2), "Unsupported type: decimal32(4, 2)"), + (pa.decimal64(10, 2), "Unsupported type: decimal64(10, 2)"), + (pa.decimal256(60, 2), "Unsupported type: decimal256(60, 2)"), + (pa.date64(), "Unsupported type: date64[ms]"), + (pa.time32("s"), "Unsupported type: time32[s]"), + (pa.time32("ms"), "Unsupported type: time32[ms]"), + (pa.time64("ns"), "Unsupported type: time64[ns]"), + (pa.timestamp("us", tz="America/New_York"), "Unsupported type: timestamp[us, tz=America/New_York]"), + (pa.duration("us"), "Unsupported type: duration[us]"), + (pa.month_day_nano_interval(), "Unsupported type: month_day_nano_interval"), + (pa.list_view(pa.int32()), "Expected primitive type, got: "), + (pa.large_list_view(pa.int32()), "Expected primitive type, got: "), + (pa.run_end_encoded(pa.int32(), pa.string()), "Unsupported type: run_end_encoded"), + (pa.dense_union([pa.field("a", pa.int32())]), "Expected primitive type, got: "), + (pa.sparse_union([pa.field("a", pa.int32())]), "Expected primitive type, got: "), + *([(pa.json_(), "Unsupported type: extension")] if hasattr(pa, "json_") else []), +] + + +@pytest.mark.parametrize("pyarrow_type, expected", PYARROW_TO_ICEBERG_SUPPORTED, ids=str) +def test_pyarrow_type_enumeration_to_iceberg(pyarrow_type: pa.DataType, expected: IcebergType) -> None: + assert visit_pyarrow(pyarrow_type, _ConvertToIceberg(format_version=3)) == expected + + +@pytest.mark.parametrize( + "pyarrow_type, expected_error", + PYARROW_TO_ICEBERG_UNSUPPORTED, + ids=[str(pyarrow_type) for pyarrow_type, _ in PYARROW_TO_ICEBERG_UNSUPPORTED], +) +def test_pyarrow_type_enumeration_unsupported(pyarrow_type: pa.DataType, expected_error: str) -> None: + with pytest.raises(TypeError, match=re.escape(expected_error)): + visit_pyarrow(pyarrow_type, _ConvertToIceberg(format_version=3)) + + +def test_pyarrow_type_enumeration_is_exhaustive() -> None: + """Fail when a PyArrow upgrade introduces a type that the enumeration above does not classify.""" + covered_type_ids = {pyarrow_type.id for pyarrow_type, _ in PYARROW_TO_ICEBERG_SUPPORTED} | { + pyarrow_type.id for pyarrow_type, _ in PYARROW_TO_ICEBERG_UNSUPPORTED + } + # The legacy interval types cannot be constructed from the public Python API + unconstructible = {"Type_INTERVAL_MONTHS", "Type_INTERVAL_DAY_TIME"} + missing = { + name + for name in dir(pa.lib) + if name.startswith("Type_") and name not in unconstructible and getattr(pa.lib, name) not in covered_type_ids + } + assert not missing, f"PyArrow types missing from the enumeration: {missing}"