andygrove opened a new pull request, #5006: URL: https://github.com/apache/datafusion-comet/pull/5006
## Which issue does this PR close? Addresses the "IPC schema is re-encoded for every block" item in #5002. That issue is an umbrella tracking several shuffle-writer optimizations, so this PR does not close it. ## Rationale for this change `ShuffleBlockWriter::write_batch` created a `StreamWriter` for every block, and `StreamWriter::try_new` re-serializes the schema flatbuffer into the stream each time. Every block in a shuffle shares one schema, so this repeats identical work per block. It matters for wide or deeply nested schemas (larger schema flatbuffer) and for shuffles with many small blocks (high partition counts), where per-block fixed costs dominate. ## What changes are included in this PR? - Pre-encode the schema IPC message once in `try_new` and write it verbatim at the start of every block; the record batch is written via `IpcDataGenerator::encode` plus the end-of-stream marker, instead of a per-block `StreamWriter`. - Schemas that contain dictionary types keep using `StreamWriter`. Dictionary encoding requires the schema and the record batch to share a dictionary tracker (the batch's dictionary ids come from the schema encoding), so those two cannot be separated; a `contains_dictionary` check selects the fallback. Output framing is unchanged: blocks remain standalone Arrow IPC streams in the same length-prefixed format, so the read side and wire format are unaffected. This is the schema-encoding half of #5002's compression items; zstd context reuse is a separate change. ### Benchmark A new criterion benchmark, `shuffle_block_schema_encoding`, encodes a single small batch (so schema serialization is a meaningful fraction) with the `None` codec, for a wide flat schema (50 columns) and a deeply nested schema (nested structs). Local results, main vs this PR: | Schema | main | this PR | | | --- | --- | --- | --- | | flat (50 cols) | ~11.0µs | ~7.75µs | ~30% faster | | nested (deep structs) | ~11.6µs | ~8.4µs | ~27% faster | ## How are these changes tested? Existing `datafusion-comet-shuffle` and `datafusion-comet` shuffle tests pass, including the roundtrip and dictionary-encoded shuffle-block tests. A new `roundtrip_ipc_dictionary` test covers a dictionary-typed column across all codecs, exercising the `StreamWriter` fallback and confirming the dictionary ids still resolve. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
