diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index 7e1babd58d..c9c68ac423 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 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: + 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])