diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index c36f1639d9..b7ab84cfbb 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -1716,6 +1716,7 @@ def _task_to_record_batches( downcast_ns_timestamp_to_us=downcast_ns_timestamp_to_us, projected_missing_fields=projected_missing_fields, allow_timestamp_tz_mismatch=True, + dictionary_columns=dictionary_columns, ) @@ -1910,6 +1911,7 @@ def _to_requested_schema( projected_missing_fields: dict[int, Any] = EMPTY_DICT, allow_timestamp_tz_mismatch: bool = False, format_model: FileFormatModel | None = None, + dictionary_columns: tuple[str, ...] = (), ) -> pa.RecordBatch: # We could reuse some of these visitors struct_array = visit_with_partner( @@ -1922,6 +1924,7 @@ def _to_requested_schema( projected_missing_fields=projected_missing_fields, allow_timestamp_tz_mismatch=allow_timestamp_tz_mismatch, format_model=format_model, + dictionary_columns=dictionary_columns, ), ArrowAccessor(file_schema), ) @@ -1935,6 +1938,7 @@ class ArrowProjectionVisitor(SchemaWithPartnerVisitor[pa.Array, pa.Array | None] _projected_missing_fields: dict[int, Any] _allow_timestamp_tz_mismatch: bool _format_model: FileFormatModel | None + _dictionary_columns: tuple[str, ...] def __init__( self, @@ -1944,6 +1948,7 @@ def __init__( projected_missing_fields: dict[int, Any] = EMPTY_DICT, allow_timestamp_tz_mismatch: bool = False, format_model: FileFormatModel | None = None, + dictionary_columns: tuple[str, ...] = (), ) -> None: if include_field_ids and format_model is None: raise ValueError("format_model is required when include_field_ids=True") @@ -1955,6 +1960,7 @@ def __init__( # Allowed for reading (aligns with Spark); disallowed for writing to enforce Iceberg spec's strict typing. self._allow_timestamp_tz_mismatch = allow_timestamp_tz_mismatch self._format_model = format_model + self._dictionary_columns = dictionary_columns def _cast_if_needed(self, field: NestedField, values: pa.Array) -> pa.Array: file_field = self._file_schema.find_field(field.field_id) @@ -2014,6 +2020,9 @@ def _cast_if_needed(self, field: NestedField, values: pa.Array) -> pa.Array: promote(file_field.field_type, field.field_type), include_field_ids=self._include_field_ids ) return values.cast(target_schema) + elif pa.types.is_dictionary(values.type) and not pa.types.is_dictionary(target_type): + if not self._dictionary_columns or field.name not in self._dictionary_columns: + return values.cast(target_type) return values diff --git a/tests/table/test_upsert.py b/tests/table/test_upsert.py index 78ddbc7c5c..60df6bd3d5 100644 --- a/tests/table/test_upsert.py +++ b/tests/table/test_upsert.py @@ -927,3 +927,53 @@ def test_upsert_snapshot_properties(catalog: Catalog) -> None: for snapshot in snapshots[initial_snapshot_count:]: assert snapshot.summary is not None assert snapshot.summary.additional_properties.get("test_prop") == "test_value" + + +def test_upsert_dictionary_encoded_columns(catalog: Catalog) -> None: + identifier = "default.test_upsert_dictionary_encoded_columns" + _drop_table(catalog, identifier) + + schema = Schema( + NestedField(1, "id", IntegerType(), required=True), + NestedField(2, "name", StringType(), required=False), + identifier_field_ids=[1], + ) + tbl = catalog.create_table(identifier, schema=schema) + + arrow_schema = pa.schema( + [ + pa.field("id", pa.int32(), nullable=False), + pa.field("name", pa.string(), nullable=True), + ] + ) + arrow_dict_schema = pa.schema( + [ + pa.field("id", pa.int32(), nullable=False), + pa.field("name", pa.dictionary(pa.int32(), pa.string()), nullable=True), + ] + ) + + # Initial append with plain string + initial_df = pa.Table.from_arrays( + [pa.array([1, 2], type=pa.int32()), pa.array(["alice", "bob"])], + schema=arrow_schema, + ) + tbl.append(initial_df) + + # Upsert with dictionary-encoded column + dict_name = pa.DictionaryArray.from_arrays(pa.array([0, 1], type=pa.int32()), pa.array(["alice_updated", "charlie"])) + dict_df = pa.Table.from_arrays( + [pa.array([1, 3], type=pa.int32()), dict_name], + schema=arrow_dict_schema, + ) + result = tbl.upsert(dict_df) + assert result.rows_updated == 1 + assert result.rows_inserted == 1 + + # Verify table contents can be read back cleanly as PyArrow table + scanned_records = tbl.scan().to_arrow().to_pydict() + assert set(zip(scanned_records["id"], scanned_records["name"], strict=True)) == { + (1, "alice_updated"), + (2, "bob"), + (3, "charlie"), + }