feat: Implement custom RecordBatch serde for shuffle for improved performance

[andygrove]

Which issue does this PR close?

Closes #1189

Rationale for this change

Arrow IPC is a good general purpose serde framework but we can get better performance by implementing specialized code optimized for Comet, which encodes single batches to shuffle blocks.

This PR implements a new BatchWriter and BatchReader and updates shuffle writer to use them when possible (when all data types are supported), falling back to Arrow IPC for other cases.

Specializations include:

• The schema gets encoded to bytes just once per shuffle operation rather than once per batch. The encoded schema bytes are then written out directly with each shuffle block, avoiding the schema serde cost per batch.
• Raw data, offset, and null buffers are written out directly with no flatbuffer encoding, no alignment, and no metadata

TPC-H

Based on median of 5 runs I am seeing a 2-3% speedup.

Microbenchmarks

Without compression:

shuffle_writer/shuffle_writer: write encoded (enable_fast_encoding=true, compression=None)
                        time:   [6.7237 µs 6.7697 µs 6.8052 µs]
shuffle_writer/shuffle_writer: write encoded (enable_fast_encoding=false, compression=None)
                        time:   [16.679 µs 16.803 µs 16.922 µs]

With LZ4 compression:

shuffle_writer/shuffle_writer: write encoded (enable_fast_encoding=true, compression=Lz4Frame)
                        time:   [173.49 µs 176.76 µs 180.52 µs]
shuffle_writer/shuffle_writer: write encoded (enable_fast_encoding=false, compression=Lz4Frame)
                        time:   [202.86 µs 204.07 µs 205.25 µs]

With ZSTD compression:

shuffle_writer/shuffle_writer: write encoded (enable_fast_encoding=true, compression=Zstd(1))
                        time:   [331.70 µs 333.24 µs 334.73 µs]
shuffle_writer/shuffle_writer: write encoded (enable_fast_encoding=false, compression=Zstd(1))
                        time:   [339.05 µs 340.44 µs 341.91 µs]

Benchmark Results

What changes are included in this PR?

How are these changes tested?

Existing tests + updated unit tests.

[parthchandra]

Are you going to end up implementing a form of arrow (stream) IPC?

[andygrove]

I discovered that we may be able to just use https://docs.rs/arrow-ipc/latest/arrow_ipc/writer/struct.IpcDataGenerator.html#method.encoded_batch and am going to look into that next

[andygrove]

Are you going to end up implementing a form of arrow (stream) IPC?

Yes, but without using flatbuffers to align and encode anything, just the raw bytes, and without the metadata messages.

[comphead]

Thanks @andygrove interesting if other datatypes keep the same performance benefit

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 57.50%. Comparing base (be48839) to head (7666fd9).
Report is 4 commits behind head on main.

Additional details and impacted files

@@              Coverage Diff              @@
##               main    #1190       +/-   ##
=============================================
+ Coverage     34.69%   57.50%   +22.80%     
+ Complexity      991      961       -30     
=============================================
  Files           116      113        -3     
  Lines         44891    11165    -33726     
  Branches       9864     2137     -7727     
=============================================
- Hits          15574     6420     -9154     
+ Misses        26168     3633    -22535     
+ Partials       3149     1112     -2037     

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

[comphead]

Thanks @andygrove thats amazing, I suppose it is extensible to support nested types in future?

[andygrove]

Thanks @andygrove thats amazing, I suppose it is extensible to support nested types in future?

Yes, I don't foresee any issues with that.

[alamb]

FYI I filed a ticket in DataFusion to consider adding this code (or something similar) upstream:

• Optimized spill file format datafusion#14078

I think it would help other serialization usecases (not just the shuffle reader/writer) -- for example sorting and grouping

[kazuyukitanimura]

Thank you @andygrove

[kazuyukitanimura]

nit/optional: snappy and zstd(6) here as well?

[viirya]

To be safer, maybe add Int32Type to the key type for this pattern.

[tustvold]

FWIW this encoding format is almost identical to the IPC format AFAICT with only some minor changes to the metadata encoding. The exact same validation is done when reading an array, and both preserve the buffers as is^1..

I suspect that much of the overhead is fixed overheads in StreamWriter, e.g. encoding the schema, and that these could be optimised and/or eliminated by using the lower-level APIs such as write_message and root_as_message. The benchmarks at least appear to agree with this, with a relatively fixed performance delta on the order of 10s of microseconds between the two encoders.

Just thought I'd put it out there, there is likely low hanging fruit in the arrow-rs IPC implementation and improvements there would definitely be welcome not just by myself.

^1. This is actually not entirely true, when dealing with sliced arrays, the arrow IPC implementation has additional logic to avoid encoding data that isn't referenced

[tustvold]

FYI if the array is sliced this will write the full data buffer, as the slice is not "materialized" into the buffer

[andygrove]

FWIW this encoding format is almost identical to the IPC format AFAICT with only some minor changes to the metadata encoding. The exact same validation is done when reading an array, and both preserve the buffers as is1..

I suspect that much of the overhead is fixed overheads in StreamWriter, e.g. encoding the schema, and that these could be optimised and/or eliminated by using the lower-level APIs such as write_message and root_as_message. The benchmarks at least appear to agree with this, with a relatively fixed performance delta on the order of 10s of microseconds between the two encoders.

Just thought I'd put it out there, there is likely low hanging fruit in the arrow-rs IPC implementation and improvements there would definitely be welcome not just by myself.

Thanks for the feedback @tustvold. I am following the current efforts to optimize Arrow IPC and will try to help.

1. This is actually not entirely true, when dealing with sliced arrays, the arrow IPC implementation has additional logic to avoid encoding data that isn't referenced

For the Comet use case, the batches that we are writing have always been freshly created and are guaranteed to not have been sliced, so we can get away with this approach.

[andygrove]

Thanks for the reviews @kazuyukitanimura and @viirya. I plan on merging this once CI is green.