andygrove opened a new issue, #7496:
URL: https://github.com/apache/datafusion-comet/issues/7496
Part of #7267.
## Problem
The JVM columnar shuffle writes a string or binary column dictionary-encoded
when each distinct value repeats more than
`spark.comet.shuffle.jvm.preferDictionary.ratio` times in a batch (default 10;
`native/shuffle/src/spark_unsafe/row.rs:1771` and `:1790`). The block stays
small, because each distinct value is stored once.
The reader expands it back to plain values. In the direct read path, which
is on by default, `ShuffleScanExec` calls `unpack_dictionary`
(`native/core/src/execution/operators/shuffle_scan.rs:265-267`), which casts
with `.expect("failed to unpack dictionary array")`. When the expanded column
passes 2 GiB, the cast's `take` returns `OffsetOverflowError` and the task
fails with a native panic. That takes 8192 rows sharing at most 819 distinct
values that average more than 256 KiB, for example a wide attribute copied onto
many rows by a join that runs in Spark before the shuffle. Spark shuffles rows
and never expands anything.
A unit test with one 300 KiB value referenced by 8192 keys panics with
`failed to unpack dictionary array: OffsetOverflowError(2147635200)`.
The same expansion fails, with an error instead of a panic, wherever else
these blocks are read:
- `ScanExec` (`native/core/src/execution/operators/scan.rs:329`), for
batches the JVM decoded and passed to native code.
- The remote (Celeborn) shuffle read path, `unpack_dictionaries` in
`native/shuffle/src/remote_schema.rs`.
- The JVM, in `CometVectorUtils.withDecodedVectors`, where
`DictionaryEncoder.decode` throws `OversizedAllocationException`.
Line references are for main at 4e6b12290f.
## Workaround
Lower `spark.comet.shuffle.jvm.batchSize`, so each block holds fewer rows.
## Proposed fix
- Return an error from `unpack_dictionary` instead of panicking.
- Compute the expanded size from the keys and the value lengths before
casting. When it would pass `i32::MAX`, split the batch and emit several
batches from the block.
## Tests
Build a `DictionaryArray<Int32, Utf8>` with one 300 KiB value and 8192 keys.
`take` sizes its output before it allocates, so this runs in milliseconds
today. After the fix, run the same shape with a small size budget injected for
the test and check the split.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
