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 34.84%. Comparing base (5c389d1) to head (96a87bd).
Report is 2 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #1190      +/-   ##
============================================
+ Coverage     34.81%   34.84%   +0.02%     
- Complexity      983      991       +8     
============================================
  Files           116      116              
  Lines         43767    43842      +75     
  Branches       9554     9562       +8     
============================================
+ Hits          15238    15275      +37     
- Misses        25561    25598      +37     
- Partials       2968     2969       +1     

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