Skip to content

dictionary_columns is not properly propagated in ‎DataScan.to_arrow_batch_reader‎ #3540

Description

@kevinjqliu

I found one propagation gap: dictionary_columns reaches ArrowScan, but to_arrow_batch_reader() wraps the batches with the plain projected schema and casts back to it:

pa.RecordBatchReader.from_batches(target_schema, batches).cast(target_schema)

That erases the dictionary type, so to_arrow(dictionary_columns=...) preserves dictionary encoding, but to_arrow_batch_reader(dictionary_columns=...).read_all() returns the normal string type.

def to_arrow_batch_reader(self, dictionary_columns: tuple[str, ...] = ()) -> pa.RecordBatchReader:
"""Return an Arrow RecordBatchReader from this DataScan.
For large results, using a RecordBatchReader requires less memory than
loading an Arrow Table for the same DataScan, because a RecordBatch
is read one at a time.
Args:
dictionary_columns:
A tuple of column names that PyArrow should read as
dictionary-encoded (``pa.DictionaryArray``). Dictionary
encoding can substantially reduce memory usage for columns
with low-cardinality repeated string values.
Only applies to Parquet files; silently ignored for ORC.
Returns:
pa.RecordBatchReader: Arrow RecordBatchReader from the Iceberg table's DataScan
which can be used to read a stream of record batches one by one.
"""
import pyarrow as pa
from pyiceberg.io.pyarrow import ArrowScan, schema_to_pyarrow
target_schema = schema_to_pyarrow(self.projection())
batches = ArrowScan(
self.table_metadata,
self.io,
self.projection(),
self.row_filter,
self.case_sensitive,
self.limit,
dictionary_columns=dictionary_columns,
).to_record_batches(self.plan_files())
return pa.RecordBatchReader.from_batches(
target_schema,
batches,
).cast(target_schema)

Example public-path regression test:

def test_to_arrow_batch_reader_preserves_dictionary_columns(catalog: Catalog) -> None:
    arrow_table = pa.table(
        {
            "id": pa.array([1, 2, 3, 4], type=pa.int32()),
            "label": pa.array(["a", "b", "a", "b"], type=pa.string()),
        }
    )
    catalog.create_namespace_if_not_exists("default")
    table = catalog.create_table("default.dict_test", schema=arrow_table.schema)
    table.append(arrow_table)

    result = table.scan().to_arrow_batch_reader(dictionary_columns=("label",)).read_all()

    assert pa.types.is_dictionary(result.schema.field("label").type)
    assert result.column("label").to_pylist() == ["a", "b", "a", "b"]

I didn’t see another user-facing path where the option is accepted and dropped. Minor adjacent cleanup: update the abstract TableScan.to_arrow signature so type checkers see the new keyword.

Originally posted by @kevinjqliu in #3461 (comment)

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions