From b770a04ca8e96a28ff9239db5183ea7da4f676a5 Mon Sep 17 00:00:00 2001 From: hedger9487 Date: Tue, 25 Aug 2026 13:04:52 +0800 Subject: [PATCH] PyArrow: Cast dictionary-encoded arrays to target schema during scan projection (#3260) --- pyiceberg/io/pyarrow.py | 11 +++++++++++ tests/io/test_pyarrow.py | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index c36f1639d9..b2a14a4d70 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,12 +1960,18 @@ 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) if field.field_type.is_primitive: if (target_type := schema_to_pyarrow(field.field_type, include_field_ids=self._include_field_ids)) != values.type: + if pa.types.is_dictionary(values.type): + if field.name not in self._dictionary_columns: + return values.cast(target_type) + return values + if field.field_type == TimestampType(): source_tz_compatible = values.type.tz is None or ( self._allow_timestamp_tz_mismatch and values.type.tz in UTC_ALIASES diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index b31c18949b..880d559a1f 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -5462,3 +5462,39 @@ def test_dictionary_columns_produces_dict_encoded_output(tmpdir: str) -> None: # Values must be identical assert result_plain.column("label").to_pylist() == result_dict.column("label").to_pylist() + + +def test_arrow_scan_mixed_dict_encoded_and_plain_strings(tmpdir: str) -> None: + schema_plain = pa.schema([pa.field("col", pa.string(), metadata={b"PARQUET:field_id": b"1"})]) + table_plain = pa.Table.from_pylist([{"col": "plain_a"}, {"col": "plain_b"}], schema=schema_plain) + file_plain = _write_table_to_data_file(f"{tmpdir}/plain.parquet", schema_plain, table_plain) + file_plain.spec_id = 0 + + schema_dict = pa.schema([pa.field("col", pa.dictionary(pa.int32(), pa.string()), metadata={b"PARQUET:field_id": b"1"})]) + table_dict = pa.Table.from_pylist([{"col": "dict_a"}, {"col": "dict_b"}], schema=schema_dict) + file_dict = _write_table_to_data_file(f"{tmpdir}/dict.parquet", schema_dict, table_dict) + file_dict.spec_id = 0 + + iceberg_schema = Schema( + NestedField(1, "col", StringType(), required=False), + ) + table_metadata = TableMetadataV2( + location=f"file://{tmpdir}", + last_column_id=1, + format_version=2, + schemas=[iceberg_schema], + partition_specs=[PartitionSpec()], + ) + io = PyArrowFileIO() + tasks = [FileScanTask(file_plain), FileScanTask(file_dict)] + + scan = ArrowScan( + table_metadata=table_metadata, + io=io, + projected_schema=iceberg_schema, + row_filter=AlwaysTrue(), + ) + + result = scan.to_table(tasks) + assert result.schema.field("col").type == schema_to_pyarrow(StringType()) + assert result.column("col").to_pylist() == ["plain_a", "plain_b", "dict_a", "dict_b"]