andygrove opened a new issue, #5198:
URL: https://github.com/apache/datafusion-comet/issues/5198
## Summary
I did a read-through audit of the entire shuffle surface looking for
performance
opportunities, and want to record the candidates somewhere before they're
lost.
**Nothing here has been benchmarked.** These are code-reading findings:
places where the
implementation does per-row, per-batch, or per-partition work that looks
avoidable. Each one
needs to be measured before it's worth acting on, and some may turn out to
be noise next to
the surrounding I/O and encode costs. Treat this as a work list for
investigation, not a list
of agreed changes.
Scope covered:
- `native/shuffle/` — partitioners, writers, IPC block encoding,
row-to-Arrow conversion
- `native/core/src/execution/operators/shuffle_scan.rs` — native shuffle read
- `spark/src/main/java/org/apache/spark/shuffle/` and
`spark/src/main/{java,scala}/org/apache/spark/sql/comet/execution/shuffle/` —
JVM shuffle
read and write paths
The native multi-partition writer has had recent attention
(apache/datafusion-comet#5006,
apache/datafusion-comet#5003, apache/datafusion-comet#5004), so most of
what's left is
concentrated in the JVM columnar path and the read path.
## Candidates
### 1. `process_sorted_row_partition` rebuilds the schema and IPC writer
once per batch
`native/shuffle/src/spark_unsafe/row.rs:1422`
```rust
let block_writer = ShuffleBlockWriter::try_new(batch.schema().as_ref(),
codec.clone())?;
written += block_writer.write_batch(&batch, &mut cursor, &ipc_time)?;
```
This sits inside the `while current_row < row_num` loop, so every batch pays
for:
- `make_batch` (`row.rs:1490`): one `format!("c{i}")` String allocation per
column, plus a
fresh `Schema`
- `ShuffleBlockWriter::try_new`: a `flattened_fields()` walk and a full IPC
schema flatbuffer
encode
This is the same per-block cost that apache/datafusion-comet#5006 removed
from the native
writer, still present on the JVM-columnar encode path.
Complication: `builder_to_array` (`row.rs:1448-1485`) decides per batch
whether a Utf8/Binary
column is emitted as `Dictionary` or plain, so the schema genuinely can vary
between batches
when `preferDictionaryRatio > 1.0`. A plain hoist is only safe when no
column is
dictionary-capable; otherwise the writer would need caching keyed on the
produced schema.
Measurable with `native/shuffle/benches/row_columnar.rs`.
### 2. The shuffle reader discards its thread-local direct buffer once per
block
`NativeBatchDecoderIterator.scala:193-197`, called from `:178`
`close()` calls `resetDataBuf()`, which throws away the thread-local
`ByteBuffer` whenever it
has grown past 128 KB. `CometBlockStoreShuffleReader.read()` (`:109-118`)
constructs and
closes one `NativeBatchDecoderIterator` **per shuffle block**, so for any
shuffle whose
compressed blocks exceed 128 KB we appear to get two `allocateDirect` calls
per block: the
discard, then the regrow at `:143`. `allocateDirect` zeroes the region and
goes through
`Bits.reserveMemory`, which can force a `System.gc()` under direct-memory
pressure.
If that reading is right, the thread-local isn't buying anything on this
path. Worth noting
that the sibling class on the ShuffleScan path,
`CometShuffleBlockIterator.java:103-106`, keeps its buffer across blocks and
never resets, so
the two paths are inconsistent.
Needs an end-to-end shuffle-read benchmark; the existing benches don't cover
this.
### 3. `RowPartition` boxes two objects per shuffled row
`spark/src/main/scala/org/apache/spark/shuffle/sort/RowPartition.scala:25-31`
```scala
private var rowAddresses: ArrayBuffer[Long] = new
ArrayBuffer[Long](initialSize)
private var rowSizes: ArrayBuffer[Int] = new ArrayBuffer[Int](initialSize)
```
Scala's `ArrayBuffer` is not specialized — it's backed by `Array[AnyRef]`.
So each `addRow`
should box a `java.lang.Long` (addresses are always outside the `Long` cache
range, so always
an allocation) and a `java.lang.Integer` (row sizes are usually >127, so
usually an
allocation). `getRowAddresses`/`getRowSizes` then unbox element-by-element
into primitive
arrays for the JNI call.
That's on the order of two allocations plus an unboxing pass per row, across
the whole JVM
columnar shuffle write path. Manually grown `Array[Long]`/`Array[Int]` with
a `count` would
remove the allocations and reduce `toArray` to an `Arrays.copyOf` memcpy —
or eliminate it, if
the JNI signature took an explicit length.
Of everything here this looks like the largest steady-state win, but it
touches a JNI
boundary, so it needs the most care.
### 4. Per-partition setup cost in the bypass-merge writer
`CometBypassMergeSortShuffleWriter.java:180-199` allocates one
`CometDiskBlockWriter` per
partition, and each one's `ArrowIPCWriter` constructor calls
`serializeSchema(schema)`
(`CometDiskBlockWriter.java:298` → `SpillWriter.java:88-104`). That's
`numPartitions × numColumns` protobuf serializations per task, with
`numPartitions` up to
`spark.shuffle.sort.bypassMergeThreshold` (200 by default), all producing
the identical
result. Same shape for the per-partition `new Native()` (`:135`) and the
four `CometConf`
lookups in the constructor (`:144-152`).
Separately, `partitioner.getPartition(key)` is called twice per record on
that path:
```java
int partition_id = partitioner.getPartition(key);
partitionWriters[partitioner.getPartition(key)].insertRow((UnsafeRow)
record._2(), partition_id);
```
`CometBypassMergeSortShuffleWriter.java:214-216` — `partition_id` is already
in hand.
### 5. Index-file parsing allocates two objects per partition
`CometNativeShuffleWriter.scala:158-171`
```scala
Files.readAllBytes(tempIndexFilePath).grouped(OFFSET_LENGTH).drop(1).map(indexBytes
=>
ByteBuffer.wrap(indexBytes).order(ByteOrder.LITTLE_ENDIAN).getLong ...)
```
`grouped` allocates an intermediate array per partition and the `map`
allocates a `ByteBuffer`
per partition. Wrapping the whole byte array in a single little-endian
`ByteBuffer` and
looping `getLong` into a pre-sized `Array[Long]` would drop `2 ×
numPartitions` allocations
per map task. Only likely to matter at high partition counts.
### 6. `PartitionedBatchIterator` widens the whole partition's index array
up front
`native/shuffle/src/partitioners/partitioned_batch_iterator.rs:85-89`
```rust
let record_batches = buffered_batches.iter().collect::<Vec<_>>();
let current_indices = indices.iter().map(|(i_batch, i_row)| (*i_batch as
usize, *i_row as usize)).collect::<Vec<_>>();
```
`interleave_record_batch` takes `&[(usize, usize)]` (arrow-select 58.4,
`interleave.rs:562`),
so the `u32 → usize` widening is unavoidable — but it doesn't have to happen
for the whole
partition at once. As written, every spill and every `shuffle_write`
allocates
`rows_in_partition × 16` bytes summed across all partitions, roughly a
second full copy of the
index array, and `record_batches` is rebuilt per partition.
Building `record_batches` once in `PartitionedBatchesProducer` and widening
indices
`batch_size` at a time into a single reusable scratch buffer would bound
this to one
allocation per spill.
Related, in the caller: `multi_partition.rs:497-500` replaces
`partition_indices` with
`vec![vec![]; num_output_partitions]` on every spill, discarding all
per-partition `Vec`
capacity so each inner `Vec` re-grows from zero after each spill. Recycling
the drained `Vec`s
would avoid that.
Measurable with `native/shuffle/benches/shuffle_writer.rs` /
`src/bin/shuffle_bench.rs`.
### 7. Config lookup inside the spill loop
`SpillWriter.java:199`
```java
int batchSize = (int) CometConf.COMET_SHUFFLE_JVM_BATCH_SIZE().get();
```
`doSpilling` runs once per partition per spill, so
`SpillSorter.writeSortedFileNative`
(`:247`, `:283`) hits this `numPartitions` times per spill file.
`CometDiskBlockWriter`
already caches the identical value in `columnarBatchSize` (`:144`) and then
doesn't use it.
### 8. Smaller items
- **`ArrowIPCWriter.spill` comparator recomputes memory usage.**
`CometDiskBlockWriter.java:382-391` sorts the static `currentWriters` list
with a comparator
calling `getActiveMemoryUsage()` twice per comparison; each call is
`synchronized` and walks
the `allocatedPages` linked list (`SpillWriter.java:249-258`). Precompute
the sort keys.
- **Two JNI round-trips per block on the ShuffleScan path.**
`shuffle_scan.rs:151-171` calls
`has_next()` then `get_buffer()`. The returned `DirectByteBuffer` is the
same object unless
it was regrown (`CometShuffleBlockIterator.java:105`), so caching a global
ref and
re-fetching only on capacity change would halve the JNI calls per block.
Probably small next
to decode cost.
- **`count_new_buffers` calls `to_data()` per column per batch.**
`multi_partition.rs:166`
allocates a `Vec<Buffer>` per array on the insert hot path just to read
buffer pointers and
capacities.
- **Dead initial fetch in the decoder.**
`NativeBatchDecoderIterator.scala:48` runs
`fetchNext()` before `channel` is assigned at `:52`, so it always returns
`None` via the
`channel == null` guard at `:93`. Harmless today because `hasNext`
re-fetches, but it
silently depends on field declaration order.
## Suggested next steps
Roughly in order of effort-to-impact:
1. Benchmark and fix item 1 and item 2 — both are localized, and both re-do
per-item work the
codebase already avoids elsewhere.
2. Items 4 through 7 are mostly hoisting and should be quick to confirm.
3. Item 3 is likely the biggest win but needs care at the JNI boundary;
worth its own issue and
its own before/after numbers.
Benchmark coverage: `native/shuffle/benches/shuffle_writer.rs`,
`native/shuffle/benches/row_columnar.rs`, and
`native/shuffle/src/bin/shuffle_bench.rs` cover
items 1, 6, and indirectly 3. Item 2 has no coverage and would need an
end-to-end shuffle-read
benchmark.
Happy for anyone to pick individual items off into their own issues/PRs.
Please post numbers
either way, including for the ones that turn out not to matter — a measured
"no effect" is
useful information to record here.
--
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]
