diff --git a/INSTALL-AND-REFERENCE.md b/INSTALL-AND-REFERENCE.md index a586dd07..c2d36c85 100644 --- a/INSTALL-AND-REFERENCE.md +++ b/INSTALL-AND-REFERENCE.md @@ -146,6 +146,15 @@ filters. Replace `pc.field("age") >= 18` with `table.column("age") >= 18`, and r `[a, b]` with `a & b`. Code using PyArrow for Arrow-side filtering can continue to use `pyarrow.compute.field()` outside the Tower table API. +Arrow schemas passed to `create()` or `create_if_not_exists()` are validated directly by +PyIceberg, which assigns field IDs and preserves nested nullability and `b"doc"` field +metadata. Timestamp units from seconds through microseconds, UTC-zoned microsecond +timestamps, `time64[us]`, `date32`, and Decimal128 values up to precision 38 are supported. +Nanosecond timestamps are rejected by default instead of being silently downcast, as are +`time32`, `time64[ns]`, `date64`, Float16, Decimal256, and non-UTC zoned timestamps. Convert +those fields explicitly before creating the table when the loss is acceptable. PyIceberg's +native validation exceptions propagate unchanged. + ### dbt Core support ```bash diff --git a/src/tower/_tables.py b/src/tower/_tables.py index a099085e..40d5b53b 100644 --- a/src/tower/_tables.py +++ b/src/tower/_tables.py @@ -39,7 +39,6 @@ ) from .exceptions import PyArrowFilterMigrationError from .tower_api_client.models import CatalogCredentials -from .utils.pyarrow import convert_pyarrow_schema from .utils.tables import ( make_table_name, namespace_or_default, @@ -650,7 +649,9 @@ def create(self, schema: pa.Schema) -> Table: Args: schema (pa.Schema): The PyArrow schema defining the structure of the table. - This will be converted to an Iceberg schema internally. + PyIceberg validates it and assigns Iceberg field IDs. Lossy or + unsupported types, including nanosecond timestamps by default, are + rejected. Returns: Table: A new Table instance wrapping the created Iceberg table. @@ -685,7 +686,7 @@ def create(self, schema: pa.Schema) -> Table: # along the way. table = catalog.create_table( identifier=table_name, - schema=convert_pyarrow_schema(schema), + schema=schema, ) return Table( @@ -712,8 +713,10 @@ def create_if_not_exists(self, schema: pa.Schema) -> Table: Args: schema (pa.Schema): The PyArrow schema defining the structure of the table. - This will be converted to an Iceberg schema internally. Note that this - schema is only used if the table needs to be created. + PyIceberg validates it and assigns Iceberg field IDs. Lossy or + unsupported types, including nanosecond timestamps by default, are + rejected. This schema is only used if the + table needs to be created. Returns: Table: A Table instance wrapping either the newly created or existing Iceberg table. @@ -747,7 +750,7 @@ def create_if_not_exists(self, schema: pa.Schema) -> Table: # exists. table = catalog.create_table_if_not_exists( identifier=table_name, - schema=convert_pyarrow_schema(schema), + schema=schema, ) return Table( diff --git a/src/tower/utils/pyarrow.py b/src/tower/utils/pyarrow.py deleted file mode 100644 index d3b46186..00000000 --- a/src/tower/utils/pyarrow.py +++ /dev/null @@ -1,178 +0,0 @@ -import pyarrow as pa - -from pyiceberg import types as iceberg_types -from pyiceberg.schema import Schema as IcebergSchema - - -class FieldIdManager: - """ - Manages the assignment of unique field IDs. - Field IDs in Iceberg start from 1. - """ - - def __init__(self, start_id=1): - # Initialize current_id to start_id - 1 so the first call to get_next_id() returns start_id - self.current_id = start_id - 1 - - def get_next_id(self) -> int: - """Returns the next available unique field ID.""" - self.current_id += 1 - return self.current_id - - -def arrow_to_iceberg_type_recursive( - arrow_type: pa.DataType, field_id_manager: FieldIdManager -) -> iceberg_types.IcebergType: - """ - Recursively convert a PyArrow DataType to a PyIceberg type, - managing field IDs for nested structures. - """ - # Primitive type mappings (most remain the same) - if pa.types.is_string(arrow_type) or pa.types.is_large_string(arrow_type): - return iceberg_types.StringType() - elif pa.types.is_integer(arrow_type): - if arrow_type.bit_width <= 32: # type: ignore - return iceberg_types.IntegerType() - else: - return iceberg_types.LongType() - elif pa.types.is_floating(arrow_type): - if arrow_type.bit_width <= 32: # type: ignore - return iceberg_types.FloatType() - else: - return iceberg_types.DoubleType() - elif pa.types.is_boolean(arrow_type): - return iceberg_types.BooleanType() - elif pa.types.is_date(arrow_type): - return iceberg_types.DateType() - elif pa.types.is_time(arrow_type): - return iceberg_types.TimeType() - elif pa.types.is_timestamp(arrow_type): - if arrow_type.tz is not None: # type: ignore - return iceberg_types.TimestamptzType() - else: - return iceberg_types.TimestampType() - elif pa.types.is_binary(arrow_type) or pa.types.is_large_binary(arrow_type): - return iceberg_types.BinaryType() - elif pa.types.is_fixed_size_binary(arrow_type): - return iceberg_types.FixedType(length=arrow_type.byte_width) # type: ignore - elif pa.types.is_decimal(arrow_type): - return iceberg_types.DecimalType(arrow_type.precision, arrow_type.scale) # type: ignore - - # Nested type mappings - elif ( - pa.types.is_list(arrow_type) - or pa.types.is_large_list(arrow_type) - or pa.types.is_fixed_size_list(arrow_type) - ): - # The element field itself in Iceberg needs an ID. - element_id = field_id_manager.get_next_id() - - # Recursively convert the list's element type. - # arrow_type.value_type is the DataType of the elements. - # arrow_type.value_field is the Field of the elements (contains name, type, nullability). - element_pyarrow_type = arrow_type.value_type # type: ignore - element_iceberg_type = arrow_to_iceberg_type_recursive( - element_pyarrow_type, field_id_manager - ) - - # Determine if the elements themselves are required (not nullable). - element_is_required = not arrow_type.value_field.nullable # type: ignore - - return iceberg_types.ListType( - element_id=element_id, - element_type=element_iceberg_type, - element_required=element_is_required, - ) - elif pa.types.is_struct(arrow_type): - struct_iceberg_fields = [] - # arrow_type is a StructType. Iterate through its fields. - for i in range(arrow_type.num_fields): # type: ignore - pyarrow_child_field = arrow_type.field(i) # This is a pyarrow.Field - - # Each field within the struct needs its own unique ID. - nested_field_id = field_id_manager.get_next_id() - nested_iceberg_type = arrow_to_iceberg_type_recursive( - pyarrow_child_field.type, field_id_manager - ) - - doc = None - if pyarrow_child_field.metadata and b"doc" in pyarrow_child_field.metadata: - doc = pyarrow_child_field.metadata[b"doc"].decode("utf-8") - - struct_iceberg_fields.append( - iceberg_types.NestedField( - field_id=nested_field_id, - name=pyarrow_child_field.name, - field_type=nested_iceberg_type, - required=not pyarrow_child_field.nullable, - doc=doc, - ) - ) - return iceberg_types.StructType(*struct_iceberg_fields) - elif pa.types.is_map(arrow_type): - # Iceberg MapType requires IDs for key and value fields. - key_id = field_id_manager.get_next_id() - value_id = field_id_manager.get_next_id() - - key_iceberg_type = arrow_to_iceberg_type_recursive( - arrow_type.key_type, field_id_manager - ) # type: ignore - value_iceberg_type = arrow_to_iceberg_type_recursive( - arrow_type.item_type, field_id_manager - ) # type: ignore - - # PyArrow map keys are always non-nullable by Arrow specification. - # Nullability of map values comes from the item_field. - value_is_required = not arrow_type.item_field.nullable # type: ignore - - return iceberg_types.MapType( - key_id=key_id, - key_type=key_iceberg_type, - value_id=value_id, - value_type=value_iceberg_type, - value_required=value_is_required, - ) - else: - raise ValueError(f"Unsupported Arrow type: {arrow_type}") - - -def convert_pyarrow_schema( - arrow_schema: pa.Schema, schema_id: int = 1, start_field_id: int = 1 -) -> IcebergSchema: - """ - Convert a PyArrow schema to a PyIceberg schema. - - Args: - arrow_schema: The input PyArrow.Schema. - schema_id: The schema ID for the Iceberg schema. - start_field_id: The starting ID for field ID assignment. - Returns: - An IcebergSchema object. - """ - field_id_manager = FieldIdManager(start_id=start_field_id) - iceberg_fields = [] - - for pyarrow_field in arrow_schema: # pyarrow_field is a pa.Field object - # Assign a unique ID for this top-level field. - top_level_field_id = field_id_manager.get_next_id() - - # Recursively convert the field's type. This will handle ID assignment - # for any nested structures using the same field_id_manager. - iceberg_field_type = arrow_to_iceberg_type_recursive( - pyarrow_field.type, field_id_manager - ) - - doc = None - if pyarrow_field.metadata and b"doc" in pyarrow_field.metadata: - doc = pyarrow_field.metadata[b"doc"].decode("utf-8") - - iceberg_fields.append( - iceberg_types.NestedField( - field_id=top_level_field_id, - name=pyarrow_field.name, - field_type=iceberg_field_type, - required=not pyarrow_field.nullable, # Top-level field nullability - doc=doc, - ) - ) - return IcebergSchema(*iceberg_fields, schema_id=schema_id) diff --git a/tests/tower/test_table_schemas.py b/tests/tower/test_table_schemas.py new file mode 100644 index 00000000..2683f05e --- /dev/null +++ b/tests/tower/test_table_schemas.py @@ -0,0 +1,229 @@ +import pyarrow as pa +import pytest +from pyiceberg import types as iceberg_types +from pyiceberg.catalog.memory import InMemoryCatalog +from pyiceberg.exceptions import ValidationError as IcebergValidationError +from pyiceberg.io.pyarrow import UnsupportedPyArrowTypeException + +import tower._tables as tables_module +from tower._context import TowerContext + + +class RecordingCatalog: + def __init__(self): + self.schemas = [] + + def create_namespace_if_not_exists(self, namespace): + pass + + def create_table(self, identifier, schema): + self.schemas.append(schema) + return object() + + def create_table_if_not_exists(self, identifier, schema): + self.schemas.append(schema) + return object() + + +def make_reference(catalog, name="events"): + context = TowerContext( + tower_url="https://api.example.com", + environment="production", + ) + return tables_module.TableReference( + context, + catalog, + name, + namespace="default", + ) + + +@pytest.fixture +def in_memory_schema_catalog(tmp_path, monkeypatch): + monkeypatch.setenv( + "PYICEBERG_DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE", + "false", + ) + catalog = InMemoryCatalog("schema-tests", warehouse=tmp_path.as_uri()) + catalog.create_namespace("default") + return catalog + + +@pytest.mark.parametrize("method", ["create", "create_if_not_exists"]) +def test_table_creation_passes_original_arrow_schema_to_catalog(method): + catalog = RecordingCatalog() + schema = pa.schema([pa.field("id", pa.int64(), nullable=False)]) + reference = make_reference(catalog) + + getattr(reference, method)(schema) + + assert catalog.schemas == [schema] + assert catalog.schemas[0] is schema + + +@pytest.mark.parametrize("catalog_type", ["s3-tables", "apache-polaris"]) +def test_external_string_catalog_creation_preserves_original_arrow_schema( + monkeypatch, catalog_type +): + context = TowerContext( + tower_url="https://api.example.com", + environment="production", + api_key="api-key", + ) + catalog = RecordingCatalog() + schema = pa.schema([pa.field("id", pa.int64(), nullable=False)]) + loaded_catalogs = [] + + def unexpected_call(*args, **kwargs): + raise AssertionError("external catalogs must not vend Tower credentials") + + def load_catalog(name): + loaded_catalogs.append(name) + return catalog + + monkeypatch.setattr( + tables_module.TowerContext, "build", staticmethod(lambda: context) + ) + monkeypatch.setattr( + tables_module, + "_describe_tower_catalog_type", + lambda ctx, name, environment: catalog_type, + ) + monkeypatch.setattr(tables_module, "_has_pyiceberg_catalog_config", unexpected_call) + monkeypatch.setattr(tables_module, "get_tower_catalog_credentials", unexpected_call) + monkeypatch.setattr(tables_module, "load_catalog", load_catalog) + + reference = tables_module.tables("events", catalog="external", namespace="default") + reference.create(schema) + + assert loaded_catalogs == ["external"] + assert reference._tower_vended is False + assert reference._catalog is catalog + assert catalog.schemas == [schema] + assert catalog.schemas[0] is schema + + +def test_pyiceberg_assigns_nested_field_ids_docs_and_nullability( + in_memory_schema_catalog, +): + schema = pa.schema( + [ + pa.field( + "id", + pa.int64(), + nullable=False, + metadata={b"doc": b"identifier"}, + ), + pa.field( + "profile", + pa.struct( + [ + pa.field( + "name", + pa.string(), + nullable=False, + metadata={b"doc": b"display name"}, + ), + pa.field( + "tags", + pa.list_(pa.field("element", pa.string(), nullable=True)), + nullable=True, + ), + ] + ), + nullable=True, + metadata={b"doc": b"profile doc"}, + ), + pa.field( + "attributes", + pa.map_( + pa.string(), + pa.field("value", pa.int32(), nullable=True), + ), + nullable=True, + ), + ] + ) + + make_reference(in_memory_schema_catalog, "nested").create(schema) + iceberg_schema = in_memory_schema_catalog.load_table("default.nested").schema() + + assert {field.name: field.field_id for field in iceberg_schema.fields} == { + "id": 1, + "profile": 2, + "attributes": 3, + } + + assert iceberg_schema.find_field("id").required is True + assert iceberg_schema.find_field("id").doc == "identifier" + assert iceberg_schema.find_field("profile").required is False + assert iceberg_schema.find_field("profile").doc == "profile doc" + assert iceberg_schema.find_field("profile.name").field_id == 4 + assert iceberg_schema.find_field("profile.name").required is True + assert iceberg_schema.find_field("profile.name").doc == "display name" + assert iceberg_schema.find_field("profile.tags").field_id == 5 + assert iceberg_schema.find_field("profile.tags.element").field_id == 6 + assert iceberg_schema.find_field("profile.tags.element").required is False + assert iceberg_schema.find_field("attributes.key").field_id == 7 + assert iceberg_schema.find_field("attributes.key").required is True + assert iceberg_schema.find_field("attributes.value").field_id == 8 + assert iceberg_schema.find_field("attributes.value").required is False + + +@pytest.mark.parametrize( + ("name", "arrow_type", "iceberg_type"), + [ + ("timestamp_s", pa.timestamp("s"), iceberg_types.TimestampType()), + ("timestamp_ms", pa.timestamp("ms"), iceberg_types.TimestampType()), + ("timestamp_us", pa.timestamp("us"), iceberg_types.TimestampType()), + ( + "timestamp_utc", + pa.timestamp("us", tz="UTC"), + iceberg_types.TimestamptzType(), + ), + ("time_us", pa.time64("us"), iceberg_types.TimeType()), + ("date", pa.date32(), iceberg_types.DateType()), + ( + "decimal", + pa.decimal128(38, 10), + iceberg_types.DecimalType(38, 10), + ), + ], +) +def test_pyiceberg_accepts_supported_arrow_precision( + in_memory_schema_catalog, name, arrow_type, iceberg_type +): + schema = pa.schema([pa.field("value", arrow_type)]) + + make_reference(in_memory_schema_catalog, name).create(schema) + + table = in_memory_schema_catalog.load_table(f"default.{name}") + assert table.schema().find_field("value").field_type == iceberg_type + + +@pytest.mark.parametrize( + ("name", "arrow_type"), + [ + ("timestamp_ns", pa.timestamp("ns")), + ("timestamp_non_utc", pa.timestamp("us", tz="Europe/Berlin")), + ("time32", pa.time32("s")), + ("time_ns", pa.time64("ns")), + ("float16", pa.float16()), + ("date64", pa.date64()), + ("decimal256", pa.decimal256(38, 10)), + ], +) +def test_pyiceberg_rejects_lossy_or_unsupported_arrow_types( + in_memory_schema_catalog, name, arrow_type +): + schema = pa.schema([pa.field("value", arrow_type)]) + + with pytest.raises(UnsupportedPyArrowTypeException): + make_reference(in_memory_schema_catalog, name).create(schema) + + +def test_pyiceberg_rejects_negative_decimal_scale(in_memory_schema_catalog): + schema = pa.schema([pa.field("value", pa.decimal128(10, -2))]) + + with pytest.raises(IcebergValidationError, match=r"decimal\(10, -2\)"): + make_reference(in_memory_schema_catalog, "negative_scale").create(schema)