fix: preserve dictionary encoding in to_arrow_batch_reader#3595
fix: preserve dictionary encoding in to_arrow_batch_reader#3595avison9 wants to merge 2 commits into
Conversation
|
|
||
| target_schema = schema_to_pyarrow(projected_schema) | ||
| if dictionary_columns: | ||
| """schema_to_pyarrow returns plain types. ArrowScan yields |
There was a problem hiding this comment.
Let's avoid multi-line strings for comments
| 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 |
There was a problem hiding this comment.
Can we hardcode the int32 for the key? What happens if this is something different than int32?
There was a problem hiding this comment.
@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
Rationale for this change
DataScan.to_arrow_batch_reader(dictionary_columns=("col",))silently dropped dictionaryencoding — callers got plain string arrays instead of
dictionary<values=string, indices=int32>,with no error or warning. The eager equivalent
to_arrow(dictionary_columns=...)workedcorrectly. This fixes the streaming path to match.
Before:
After:
Root cause:
_to_arrow_batch_reader_via_file_scan_tasksbuilttarget_schemaviaschema_to_pyarrow(which returns plain types), then called.cast(target_schema)on theRecordBatchReader.ArrowScan.to_record_batchescorrectly yielded dictionary-encodedbatches, but the cast converted them back to plain types. The fix rebuilds
target_schemawith
pa.dictionary(pa.int32(), field.type)for each column indictionary_columnsbeforethe cast, so the encoding is preserved end-to-end.
Are these changes tested?
Yes. Added
test_to_arrow_batch_reader_preserves_dictionary_columnsintests/io/test_pyarrow.py— writes a two-column Parquet file and asserts that_to_arrow_batch_reader_via_file_scan_tasksreturns a dictionary-encodedlabelcolumnwhile leaving
idas plainint32.Are there any user-facing changes?
Yes —
to_arrow_batch_reader(dictionary_columns=...)now correctly returns dictionary-encodedarrays for the requested columns, matching the documented behaviour and the existing
to_arrowpath. Please apply the
changeloglabel to this PR — the project has no changelog file/newsfragmentrequirement (no
CONTRIBUTING.md); release notes are generated from labeled, merged PR titles.