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

andygrove · 2024-12-20T16:51:09Z

Which issue does this PR close?

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.

native/core/src/execution/shuffle/batch_serde.rs

parthchandra · 2024-12-20T21:46:58Z

native/core/src/execution/shuffle/batch_serde.rs

+use std::io::Write;
+use std::sync::Arc;
+
+pub fn write_batch_fast(


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

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

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 · 2024-12-20T23:02:01Z

native/core/benches/batch_serde.rs

+
+fn create_batch() -> RecordBatch {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("c0", DataType::Utf8, true),


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

codecov-commenter · 2024-12-30T14:13:09Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 34.67%. Comparing base (be48839) to head (7c93806).
Report is 2 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #1190      +/-   ##
============================================
- Coverage     34.69%   34.67%   -0.02%     
+ Complexity      991      987       -4     
============================================
  Files           116      116              
  Lines         44891    44883       -8     
  Branches       9864     9884      +20     
============================================
- Hits          15574    15565       -9     
+ Misses        26168    26155      -13     
- Partials       3149     3163      +14

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

native/core/benches/shuffle_writer.rs

native/core/src/execution/shuffle/codec.rs

…ch-serde

comphead · 2025-01-10T00:18:24Z

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

andygrove · 2025-01-10T01:03:24Z

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 · 2025-01-10T14:00:32Z

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 · 2025-01-10T22:58:01Z

native/core/benches/shuffle_writer.rs

+    for compression_codec in [
+        CompressionCodec::None,
+        CompressionCodec::Lz4Frame,
+        CompressionCodec::Zstd(1),


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

viirya · 2025-01-10T23:30:04Z

native/core/src/execution/shuffle/codec.rs

+                let null_buffer = self.read_null_buffer();
+                Arc::new(BinaryArray::try_new(offset_buffer, buffer, null_buffer)?)
+            }
+            DataType::Dictionary(_, _) => {


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

native/core/src/execution/shuffle/codec.rs

Co-authored-by: Liang-Chi Hsieh <[email protected]>

tustvold · 2025-01-11T09:36:32Z

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 · 2025-01-11T09:42:57Z

native/core/src/execution/shuffle/codec.rs

+            DataType::Utf8 => {
+                let arr = col.as_any().downcast_ref::<StringArray>().unwrap();
+                // write data buffer
+                self.write_buffer(arr.values())?;


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

andygrove force-pushed the experimental-fast-batch-serde branch from 49d0c27 to f7d8cce Compare December 20, 2024 17:07

kazuyukitanimura reviewed Dec 20, 2024

View reviewed changes

native/core/src/execution/shuffle/batch_serde.rs Outdated Show resolved Hide resolved

parthchandra reviewed Dec 20, 2024

View reviewed changes

comphead reviewed Dec 20, 2024

View reviewed changes

andygrove changed the title ~~feat: Implement fast serde for single record batches~~ [do not review] feat: Implement fast serde for single record batches Dec 21, 2024

andygrove force-pushed the experimental-fast-batch-serde branch from e32db8d to c2b9e5c Compare December 28, 2024 15:35

andygrove changed the title ~~[do not review] feat: Implement fast serde for single record batches~~ feat: Implement fast serde for single record batches Dec 28, 2024

andygrove changed the title ~~feat: Implement fast serde for single record batches~~ feat: Implement custom RecordBatch serde for shuffle for improved performance Dec 30, 2024

andygrove mentioned this pull request Jan 3, 2025

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

Merged

andygrove marked this pull request as ready for review January 8, 2025 01:02

andygrove added 2 commits January 7, 2025 18:05

Implement faster encoder for shuffle blocks

eda771a

make code more concise

b9b3d9d

andygrove force-pushed the experimental-fast-batch-serde branch from 96a87bd to b9b3d9d Compare January 8, 2025 01:12

andygrove requested review from comphead, kazuyukitanimura, parthchandra and viirya and removed request for parthchandra January 8, 2025 01:14

andygrove added 2 commits January 7, 2025 18:38

enable fast encoding for columnar shuffle

d0fe67a

update benches

629598a

kazuyukitanimura reviewed Jan 8, 2025

View reviewed changes

native/core/benches/shuffle_writer.rs Show resolved Hide resolved

native/core/src/execution/shuffle/codec.rs Show resolved Hide resolved

Merge remote-tracking branch 'apache/main' into experimental-fast-bat…

498bf70

…ch-serde

andygrove added 4 commits January 9, 2025 17:42

test all int types

98846c9

test float

a373aaf

remaining types

d03c957

add Snappy and Zstd(6) back to benchmark

5397ae7

andygrove mentioned this pull request Jan 10, 2025

docs: Update TPC-H benchmark results #1257

Draft

alamb mentioned this pull request Jan 10, 2025

Optimized spill file format apache/datafusion#14078

Open

3 tasks

fix regression

09448e5

andygrove requested a review from kazuyukitanimura January 10, 2025 21:06

kazuyukitanimura approved these changes Jan 10, 2025

View reviewed changes

viirya reviewed Jan 10, 2025

View reviewed changes

native/core/src/execution/shuffle/codec.rs Outdated Show resolved Hide resolved

viirya approved these changes Jan 10, 2025

View reviewed changes

Update native/core/src/execution/shuffle/codec.rs

7c93806

Co-authored-by: Liang-Chi Hsieh <[email protected]>

tustvold reviewed Jan 11, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

andygrove commented Dec 20, 2024 •

edited

Loading

parthchandra Dec 20, 2024

andygrove Dec 20, 2024

andygrove Dec 21, 2024

comphead Dec 20, 2024

codecov-commenter commented Dec 30, 2024 •

edited

Loading

comphead commented Jan 10, 2025

andygrove commented Jan 10, 2025

alamb commented Jan 10, 2025

kazuyukitanimura left a comment

kazuyukitanimura Jan 10, 2025

viirya Jan 10, 2025

tustvold commented Jan 11, 2025 •

edited

Loading

tustvold Jan 11, 2025

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

Are you sure you want to change the base?

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

Conversation

andygrove commented Dec 20, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

TPC-H

Microbenchmarks

Benchmark Results

What changes are included in this PR?

How are these changes tested?

parthchandra Dec 20, 2024

Choose a reason for hiding this comment

andygrove Dec 20, 2024

Choose a reason for hiding this comment

andygrove Dec 21, 2024

Choose a reason for hiding this comment

comphead Dec 20, 2024

Choose a reason for hiding this comment

codecov-commenter commented Dec 30, 2024 • edited Loading

Codecov Report

comphead commented Jan 10, 2025

andygrove commented Jan 10, 2025

alamb commented Jan 10, 2025

kazuyukitanimura left a comment

Choose a reason for hiding this comment

kazuyukitanimura Jan 10, 2025

Choose a reason for hiding this comment

viirya Jan 10, 2025

Choose a reason for hiding this comment

tustvold commented Jan 11, 2025 • edited Loading

tustvold Jan 11, 2025

Choose a reason for hiding this comment

andygrove commented Dec 20, 2024 •

edited

Loading

codecov-commenter commented Dec 30, 2024 •

edited

Loading

tustvold commented Jan 11, 2025 •

edited

Loading