Skip to content

fix: preserve dictionary encoding in to_arrow_batch_reader#3595

Open
avison9 wants to merge 2 commits into
apache:mainfrom
avison9:fix/strict-metrics-not-equal-not-in
Open

fix: preserve dictionary encoding in to_arrow_batch_reader#3595
avison9 wants to merge 2 commits into
apache:mainfrom
avison9:fix/strict-metrics-not-equal-not-in

Conversation

@avison9

@avison9 avison9 commented Jul 2, 2026

Copy link
Copy Markdown

Rationale for this change

DataScan.to_arrow_batch_reader(dictionary_columns=("col",)) silently dropped dictionary
encoding — callers got plain string arrays instead of dictionary<values=string, indices=int32>,
with no error or warning. The eager equivalent to_arrow(dictionary_columns=...) worked
correctly. This fixes the streaming path to match.

Before:

result = table.scan().to_arrow_batch_reader(dictionary_columns=("label",)).read_all()
print(result.schema.field("label").type)
# string  ← dictionary encoding silently lost

After:

result = table.scan().to_arrow_batch_reader(dictionary_columns=("label",)).read_all()
print(result.schema.field("label").type)
# dictionary<values=string, indices=int32>  ✓

Root cause: _to_arrow_batch_reader_via_file_scan_tasks built target_schema via
schema_to_pyarrow (which returns plain types), then called .cast(target_schema) on the
RecordBatchReader. ArrowScan.to_record_batches correctly yielded dictionary-encoded
batches, but the cast converted them back to plain types. The fix rebuilds target_schema
with pa.dictionary(pa.int32(), field.type) for each column in dictionary_columns before
the cast, so the encoding is preserved end-to-end.

Are these changes tested?

Yes. Added test_to_arrow_batch_reader_preserves_dictionary_columns in
tests/io/test_pyarrow.py — writes a two-column Parquet file and asserts that
_to_arrow_batch_reader_via_file_scan_tasks returns a dictionary-encoded label column
while leaving id as plain int32.

Are there any user-facing changes?

Yes — to_arrow_batch_reader(dictionary_columns=...) now correctly returns dictionary-encoded
arrays for the requested columns, matching the documented behaviour and the existing to_arrow
path. Please apply the changelog label to this PR — the project has no changelog file/newsfragment
requirement (no CONTRIBUTING.md); release notes are generated from labeled, merged PR titles.

@Fokko Fokko left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR @avison9. I've left some small comments that we should resolve first

Comment thread pyiceberg/table/__init__.py Outdated

target_schema = schema_to_pyarrow(projected_schema)
if dictionary_columns:
"""schema_to_pyarrow returns plain types. ArrowScan yields

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's avoid multi-line strings for comments

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Noted, that is updated also

Comment thread pyiceberg/table/__init__.py Outdated
dict_col_set = set(dictionary_columns)
target_schema = pa.schema(
[
field.with_type(pa.dictionary(pa.int32(), field.type)) if field.name in dict_col_set else field

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we hardcode the int32 for the key? What happens if this is something different than int32?

@avison9 avison9 Jul 6, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Fokko, presently for every code path that actually reaches this line, it will always be int32. But that's because of an internal choice in Arrow's C++ Parquet reader, not something PyArrow documents or guarantees anyway.
If a future PyArrow version ever changed that reader's behavior (e.g., to shrink to int8/int16 for low-cardinality columns, matching dictionary_encode()'s behavior), the hardcoded int32 would silently be wrong again, in exactly the same class of bug this PR is fixing.
The best solution will be to derive the actual dictionary field from the first real batch ArrowScan produces, peek one batch, read its actual index type off first_batch.schema.field(name), and prepend that batch back onto the stream before building target_schema.
I have added this in a new commit, if you can please take a look

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants