From db4918fce7ed1928f045fe6703d06011ceb0af7c Mon Sep 17 00:00:00 2001 From: Gayathri Srividya Rajavarapu Date: Fri, 17 Jul 2026 18:38:11 +0530 Subject: [PATCH 1/2] fix: handle mixed dictionary/plain ArrowScan batches --- pyiceberg/io/pyarrow.py | 16 +++++++++++++++- tests/io/test_pyarrow.py | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index 7e1babd58d..b5d894d745 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -1802,10 +1802,24 @@ def to_table(self, tasks: Iterable[FileScanTask]) -> pa.Table: # Empty return arrow_schema.empty_table() + def _normalize_dictionary_columns(batch: pa.RecordBatch) -> pa.RecordBatch: + arrays = [] + has_dictionary = False + + for array in batch.columns: + if pa.types.is_dictionary(array.type): + arrays.append(array.dictionary_decode()) + has_dictionary = True + else: + arrays.append(array) + + return pa.record_batch(arrays, names=batch.schema.names) if has_dictionary else batch + # Note: cannot use pa.Table.from_batches(itertools.chain([first_batch], batches))) # as different batches can use different schema's (due to large_ types) result = pa.concat_tables( - (pa.Table.from_batches([batch]) for batch in itertools.chain([first_batch], batches)), promote_options="permissive" + (pa.Table.from_batches([_normalize_dictionary_columns(batch)]) for batch in itertools.chain([first_batch], batches)), + promote_options="permissive", ) return result diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index 532311899d..ac18d02764 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -1209,6 +1209,42 @@ def test_projection_add_column(file_int: str) -> None: ) +def test_arrowscan_to_table_mixed_dictionary_and_plain_string_batches() -> None: + schema = Schema(NestedField(field_id=1, name="col", field_type=StringType(), required=False)) + + scan = ArrowScan( + table_metadata=TableMetadataV2( + location="file://a/b/", + last_column_id=1, + format_version=2, + schemas=[schema], + partition_specs=[PartitionSpec()], + ), + io=PyArrowFileIO(), + projected_schema=schema, + row_filter=AlwaysTrue(), + case_sensitive=True, + ) + + plain_string = pa.array(["a"], type=pa.string()) + dictionary_string = plain_string.dictionary_encode() + + with patch.object( + ArrowScan, + "to_record_batches", + return_value=iter( + [ + pa.record_batch([plain_string], names=["col"]), + pa.record_batch([dictionary_string], names=["col"]), + ] + ), + ): + result = scan.to_table(tasks=[]) + + assert result.schema.field("col").type == pa.string() + assert result.column("col").to_pylist() == ["a", "a"] + + def test_read_list(schema_list: Schema, file_list: str) -> None: result_table = project(schema_list, [file_list]) From b88b709d1a07710de0964a5b91d0204b00f52c51 Mon Sep 17 00:00:00 2001 From: Gayathri Srividya Rajavarapu Date: Fri, 17 Jul 2026 19:08:58 +0530 Subject: [PATCH 2/2] fix: preserve requested dictionary columns --- pyiceberg/io/pyarrow.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index b5d894d745..c9c68ac423 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -1806,8 +1806,8 @@ def _normalize_dictionary_columns(batch: pa.RecordBatch) -> pa.RecordBatch: arrays = [] has_dictionary = False - for array in batch.columns: - if pa.types.is_dictionary(array.type): + for name, array in zip(batch.schema.names, batch.columns, strict=True): + if pa.types.is_dictionary(array.type) and name not in self._dictionary_columns: arrays.append(array.dictionary_decode()) has_dictionary = True else: