[EPIC] Improve shuffle performance

[andygrove]

This epic is for improving shuffle / ScanExec performance.

Issues

• Implement faster single batch encoding/decoding for use in shuffle #1189
• Allow native shuffle batch size to be configured separately from comet default batch size #1187
• Add support for lz4 compression in shuffle #1178
• Comet native shuffle reader #1125
• Can we stop copying the Arrow schema over FFI for every batch? #1115
• Possible native shuffle optimization #977
• Optimize repartitioning logic in ShuffleWriterExec using interleave_record_batch #1235
• Native shuffle double allocates memory #1448
• Native shuffle inaccurate estimate of builder memory allocation #1449
• Re-implement memory management in native shuffle writer #1446
• Columnar shuffle uses wrong memory allocator in unified memory mode #1438
• Optimize native shuffle for single partition case #1453

Context

I have been comparing Comet and Ballista performance for TPC-H q3. Both execute similar native plans. I am using the comet-parquet-exec branch which uses DataFusion's ParquetExec.

Ballista is approximately 3x faster than Comet. Given that they are executing similar DataFusion native plans, I would expect performance to be similar.

The main difference between Comet and Ballista is that Comet transfers batches between JVM and native code during shuffle operations.

Most of the native execution time in Comet is spent in ScanExec which is reading Arrow batches from the JVM using Arrow FFI. This time was not included in our metrics prior to #1128 and #1111.

[viirya]

Most of the native execution time in Comet is spent in ScanExec which is reading Arrow batches from the JVM using Arrow FFI.

Do you mean when ScanExec is used as pseudo scan node to read shuffled data into next native execution?

If so, it makes some sense because Comet shuffle reader is still JVM based. We should make it native eventually to boost shuffle reading performance. It is on the early roadmap when we started on Comet shuffle. Although it is not super urgent and high priority at that moment. But now I think it is the time we can begin to work on this.

Opened: #1125

[andygrove]

Do you mean when ScanExec is used as pseudo scan node to read shuffled data into next native execution?

Yes, exactly.

[andygrove]

If so, it makes some sense because Comet shuffle reader is still JVM based.

This is also an issue for shuffle writes. The child node of ShuffleWriterExec is always a ScanNode reading the output from a native plan.

We pay the FFI cost twice - once to import from native plan to JVM, then again to export to the shuffle write native plan. We have the cost of schema serde in both directions. Perhaps there is a way to shortcut this and avoid a full serde because we do not need to read the batch in the JVM in this case, just pass it from one native plan to another.

[andygrove]

I created a Google document for collaborating on ideas around improving shuffle performance:

https://docs.google.com/document/d/1rx1ue7UZ4ljzic9Rc2kT-v35bfLB7Rhhe5FW1d0Sw4I/edit?usp=sharing

[andygrove]

Update: The ScanExec time is from calling CometBatchIterator.next() so includes the time for it to fetch input batches as well as the FFI cost. It looks like the cost of fetching the batches is much more than the FFI cost. Perhaps this is really measuring the execution time of the input query, so could be misleading