feat: Move shuffle block decompression and decoding to native code and add LZ4 & Snappy support

[andygrove]

Which issue does this PR close?

Part of #1123

Closes #1178

Rationale for this change

We currently perform encoding + compression in native code and decoding + decompression in JVM code. There are some downsides to this approach:

• We need compatible and efficient JVM and Rust compression libraries (this seems challenging for LZ4)
• We need compatible and efficient JVM and Rust encoding libraries (we use Arrow IPC currently)
• We cannot have unit tests for roundtrip encoding and compression; only integration tests (results in slow dev cycles)
• Makes it difficult to experiment with different encoding and compression libraries and techniques
• We are missing out on performance, potentially

What changes are included in this PR?

• Call native code for decompression + decoding
• Add metrics for decode + decompress
• Add support for LZ4
• Implement unit tests for round-trip

ZSTD

LZ4

Microbenchmarks

shuffle_writer/shuffle_writer: encode (no compression))
                        time:   [10.044 µs 10.371 µs 10.698 µs]
shuffle_writer/shuffle_writer: encode and compress (lz4)
                        time:   [132.60 µs 133.01 µs 133.51 µs]
shuffle_writer/shuffle_writer: encode and compress (zstd level 1)
                        time:   [217.81 µs 218.01 µs 218.25 µs]

TPC-H

edit: Results update on 1/1/2025 after making compression configurable for columnar shuffle as well as native shuffle.

How are these changes tested?

Existing tests

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 77.67857% with 25 lines in your changes missing coverage. Please review.

Project coverage is 34.71%. Comparing base (103f82f) to head (7dd2ff6).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...execution/shuffle/NativeBatchDecoderIterator.scala	71.60%	12 Missing and 11 partials ⚠️
...t/execution/shuffle/CometShuffleExchangeExec.scala	66.66%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #1192      +/-   ##
============================================
+ Coverage     34.06%   34.71%   +0.64%     
- Complexity      925      958      +33     
============================================
  Files           115      115              
  Lines         43569    43614      +45     
  Branches       9528     9534       +6     
============================================
+ Hits          14843    15141     +298     
+ Misses        25777    25507     -270     
- Partials       2949     2966      +17     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[andygrove]

The status is that this is working, but I am running into some executor OOMs when trying to run complete benchmarks. I will pick this up again after the holidays.

[andygrove]

@Dandandan you may be interested in the benchmark results

[Dandandan]

I think zstd should have faster negative levels as well (-4 or -5 might come close), would be interesting to see how it compares. Not sure if it is available in the rust bindings.

[andygrove]

Added Snappy.

shuffle_writer/shuffle_writer: encode and compress (snappy)
                        time:   [87.556 µs 88.328 µs 89.092 µs]

[andygrove]

@viirya @kazuyukitanimura @parthchandra @mbutrovich This is ready for review now

[mbutrovich]

This is a somewhat confusingly named feature flag. Can we leave a comment here that we're enabling unsafe encode and decode for performance?

[mbutrovich]

Could we modify this test so instead of being named round_trip_zstd it iterates through all the compression schemes?

[andygrove]

Good idea. I have implemented this in 1fc3d49 and also fixed the bug that it exposed 🙂

[mbutrovich]

I know this isn't your code, but it jumped out at me: can we replace this 8-byte write with a seek?

// seek past ipc_length placeholder
    output.seek_relative(8)?;

[comphead]

awesome, its really challenging to find a well maintained snappy Rust library.

[comphead]

probably dumb question, looks like the writers gets created for every input.

In real life example it will be recreated for every batch, probably the writer can be reused until the data stream has completed?

[andygrove]

Spark's approach is to write shuffled output data in blocks where each block is independent (contains schema + data for that batch), so it doesn't allow for a streaming approach. We accumulate data in the shuffle writer until we reach the configured batch size and then write that batch out to bytes. There is also no guarantee of the order in which the shuffle reader will read these blocks, so that is another reason why we cannot keep a writer open for multiple batches.

[comphead]

do you think it makes any sense if I play with reusable writers to check if it gives any benefit on multiple bacthes? I feel we can have a map of it where key is schema hash + compression algorithm. During real job runtime its probably not expected to have more than 50-100 writers

[andygrove]

Feel free to experiment, and let's see what else we can learn. I'm not sure exactly what you are planning on trying but it seems that we could just have one writer per PartitionBuffer and write multiple batches to it. It may complicate the spill logic but you can just disable that for testing the idea. I wonder how this would differ from just increasing the shuffle batch size so that we write larger batches?

[comphead]

yeah, thats another way to increase batch_size so the compression writer will be created less often.
The test idea was to find out how expensive is the writer creation, to see if this can be an issue.

Perhaps I can just try to create a bench and quantify it so it gives us some ideas.

[comphead]

what I found is

Time elapsed in writer_create is: 107.125µs
Time elapsed in total is: 1.699042ms

The average for creation is 7% of overall time. However to make writer a singleton is not a trivial task...

[andygrove]

In the PR that follows this one (#1190) I propose replacing ArrowWriter with a faster proprietary version.

[andygrove]

@viirya @kazuyukitanimura @mbutrovich @comphead Thanks for the reviews so far. I believe I have addressed all feedback now.

[kazuyukitanimura]

Thanks @andygrove

[kazuyukitanimura]

nit since the config name now has zstd, the constant name should ideally reflect it, but optional

[alamb]

I think @andygrove mentioned this PR on the sync up call today

I was thinking that a building block to quickly write arrow data to/from disk/SSD would be a valuable feature for many systems in addition to Comet -- for example I think it could significantly improve performance when data is spilled in sorting / grouping.

I wonder if any thought has gone into this? I would be happy to help organize a project upstream in DataFusion to creat this building block but didn't want to start that discussion until I asked over here first

[andygrove]

@alamb Here is the follow on PR for the custom encoding to replace Arrow IPC:

• feat: Implement custom RecordBatch serde for shuffle for improved performance #1190

It would definitely be good to collaborate on getting some version of this implemented upstream.

I also filed an issue to explore using interleave_record_batch (I did not know that this existed).

• Optimize repartitioning logic in ShuffleWriterExec using interleave_record_batch #1235

[alamb]

@alamb Here is the follow on PR for the custom encoding to replace Arrow IPC:

THank you -- to follow up here I also filed a ticket to consider adding a better spill format in datafusion:

• Optimized spill file format datafusion#14078