perf: store shuffle offsets as lengths with zstd

docs/source/contributor-guide/jvm_shuffle.md

@@ -215,9 +215,10 @@ writes the same Arrow IPC block format as native shuffle, so direct read applies
 
 ## Configuration
 
-| Config                                   | Description                                       |
-| ---------------------------------------- | ------------------------------------------------- |
-| `spark.comet.shuffle.jvm.batchSize`      | Rows per Arrow batch                              |
-| `spark.comet.shuffle.jvm.spillThreshold` | Row count threshold for spill                     |
-| `spark.comet.shuffle.compression.codec`  | Compression codec, `lz4` by default               |
-| `spark.comet.shuffle.directRead.enabled` | Decode blocks natively on read, `true` by default |
+| Config                                                     | Description                                                                                   |
+| ---------------------------------------------------------- | --------------------------------------------------------------------------------------------- |
+| `spark.comet.shuffle.jvm.batchSize`                        | Rows per Arrow batch                                                                          |
+| `spark.comet.shuffle.jvm.spillThreshold`                   | Row count threshold for spill                                                                 |
+| `spark.comet.shuffle.compression.codec`                    | Compression codec, `lz4` by default                                                           |
+| `spark.comet.shuffle.compression.offsetsAsLengths.enabled` | Store offsets as lengths in zstd blocks without dictionary-encoded columns, `true` by default |
+| `spark.comet.shuffle.directRead.enabled`                   | Decode blocks natively on read, `true` by default                                             |

docs/source/contributor-guide/native_shuffle.md

@@ -158,6 +158,7 @@ The native shuffle implementation is its own workspace crate, `datafusion-comet-
 | `writers/buf_batch_writer.rs`      | `BufBatchWriter`, which coalesces sub-`batch_size` batches through Arrow's `BatchCoalescer` before serializing.                  |
 | `writers/local/spill.rs`           | `PartitionedSpill`, the one spill file shared by every output partition, and its per-partition byte ranges.                      |
 | `ipc.rs`                           | `read_ipc_compressed` and `read_ipc_compressed_validated`, the decode side used by the shuffle reader.                           |
+| `offsets_as_lengths.rs`            | Which IPC body buffers hold offsets, the rewrite of offsets as lengths in byte planes, and the `LENGTHS_TAG_FLAG` codec-tag bit. |
 
 ## Data Flow
 
@@ -191,6 +192,10 @@ The native shuffle implementation is its own workspace crate, `datafusion-comet-
    - Writes compression type header
    - Writes field count header
    - Writes compressed IPC stream
+   - With zstd and `spark.comet.shuffle.compression.offsetsAsLengths.enabled`, first rewrites
+     the IPC body's offsets buffers as lengths in byte planes, and marks the block by setting the
+     high bit of the codec tag's first byte (`native/shuffle/src/offsets_as_lengths.rs`). The
+     reader rebuilds the offsets before decoding.
 
 6. **Output**: One data file holds the concatenated partition data. The writer records the byte
    offset where each partition begins, plus the total length, and keeps them in memory.
@@ -270,8 +275,10 @@ leave that walk consistent with it.
 
 `CometShuffleBlockIterator` reads a 16-byte header per block: an 8-byte compressed length, which
 includes the field count but not itself, followed by an 8-byte field count that is discarded because
-the schema comes from the `ShuffleScan` protobuf fields. The compressed body is read into a reused
-`DirectByteBuffer`, **valid only until the next `hasNext()` call**. Native must fully consume it
+the schema comes from the `ShuffleScan` protobuf fields. The 4-byte codec tag after it is read only by
+native code; the high bit of its first byte marks a block whose offsets buffers hold lengths. The
+compressed body is read into a reused `DirectByteBuffer`, **valid only until the next `hasNext()`
+call**. Native must fully consume it
 before pulling the next block, which `read_ipc_compressed` satisfies because it allocates fresh
 native memory for the decoded data.
 
@@ -510,6 +517,7 @@ independently compressed, allowing parallel decompression during reads.
 | `spark.comet.shuffle.directRead.enabled`                                  | `true`  | Decode shuffle blocks in native code, bypassing Arrow FFI                     |
 | `spark.comet.shuffle.compression.codec`                                   | `lz4`   | Compression codec                                                             |
 | `spark.comet.shuffle.compression.zstd.level`                              | `1`     | Zstd compression level                                                        |
+| `spark.comet.shuffle.compression.offsetsAsLengths.enabled`                | `true`  | Store offsets as lengths in zstd blocks                                       |
 | `spark.comet.shuffle.native.writeBufferSize`                              | `1MB`   | Write buffer size                                                             |
 | `spark.comet.shuffle.native.maxBufferBytes`                               | `0`     | Fixed spill threshold. `0` disables it, leaving pool pressure                 |
 | `spark.comet.shuffle.native.partitioning.hash.enabled`                    | `true`  | Allow `HashPartitioning` on the native path                                   |

docs/source/user-guide/latest/tuning/shuffle.md

@@ -115,6 +115,15 @@ both Comet's native shuffle and the JVM columnar shuffle writer. Supported value
 useful when shuffle I/O or network bandwidth is the bottleneck. When ZSTD is selected, the level is
 controlled by `spark.comet.shuffle.compression.zstd.level` (default `1`).
 
+With zstd, Comet also stores the offsets of string, binary, array and map columns as lengths before
+compressing each shuffle block. Lengths compress much better than offsets: in our measurements, shuffle
+files with such columns were 10% to 40% smaller, and zstd was no slower. This applies to native shuffle
+and to JVM shuffle blocks without dictionary-encoded columns. A JVM shuffle block in which
+`spark.comet.shuffle.jvm.preferDictionary.ratio` dictionary-encodes a column keeps the offsets of all
+its columns, and Celeborn shuffle blocks keep theirs too.
+`spark.comet.shuffle.compression.offsetsAsLengths.enabled` turns this off. LZ4 and snappy shuffle blocks
+are written as before.
+
 `spark.shuffle.compress=false` disables compression for Comet's native shuffle only. It has no effect on
 the JVM columnar shuffle writer, which always compresses spill files with the codec above. Disabling
 compression on the native path may result in faster shuffle times in certain environments, such as

native/core/src/execution/jni_api.rs

@@ -1773,6 +1773,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_writeSortedFileNative
     current_checksum: jlong,
     compression_codec: JString,
     compression_level: jint,
+    offsets_as_lengths: jboolean,
     tracing_enabled: jboolean,
 ) -> jlongArray {
     try_unwrap_or_throw(&e, |env| unsafe {
@@ -1798,6 +1799,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_writeSortedFileNative
                     current_checksum,
                     &compression_codec,
                     compression_level,
+                    offsets_as_lengths != JNI_FALSE,
                 )?;
 
                 let long_array = env.new_long_array(results.len())?;
@@ -1829,6 +1831,7 @@ unsafe fn write_sorted_file(
     current_checksum: i64,
     compression_codec: &str,
     compression_level: i32,
+    offsets_as_lengths: bool,
 ) -> CometResult<[i64; 3]> {
     if row_sizes.len() < row_addresses.len() {
         return Err(CometError::Internal(format!(
@@ -1865,6 +1868,7 @@ unsafe fn write_sorted_file(
         checksum_algo,
         current_checksum,
         &compression_codec,
+        offsets_as_lengths,
     )?;
 
     let checksum = checksum.map_or(i64::MIN, |checksum| checksum as i64);
@@ -3376,6 +3380,7 @@ mod tests {
                 current_checksum,
                 codec,
                 1,
+                false,
             )
         }?;
         Ok((results, std::fs::read(path).unwrap()))

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

@@ -104,15 +104,18 @@ impl OperatorBuilder for ShuffleWriterBuilder {
         // only ever sees a real limit or none at all.
         let max_buffer_bytes =
             (writer.max_buffer_bytes > 0).then_some(writer.max_buffer_bytes as usize);
-        let shuffle_writer = Arc::new(ShuffleWriterExec::try_new_with_destination(
-            writer_input,
-            partitioning,
-            codec,
-            destination,
-            writer.tracing_enabled,
-            write_buffer_size,
-            max_buffer_bytes,
-        )?);
+        let shuffle_writer = Arc::new(
+            ShuffleWriterExec::try_new_with_destination(
+                writer_input,
+                partitioning,
+                codec,
+                destination,
+                writer.tracing_enabled,
+                write_buffer_size,
+                max_buffer_bytes,
+            )?
+            .with_offsets_as_lengths(writer.offsets_as_lengths),
+        );
 
         Ok((
             scans,

native/proto/src/proto/operator.proto

@@ -987,6 +987,9 @@ message ShuffleWriter {
   uint64 max_buffer_bytes = 10;
   // Explicit output destination. When absent, use the legacy output file fields.
   PartitionWriter partition_writer = 11;
+  // Store the offsets of variable-width columns as lengths in local shuffle blocks, which compress
+  // better. The native writer applies it only with zstd for now.
+  bool offsets_as_lengths = 12;
 }
 
 message ParquetWriter {

native/shuffle/benches/row_columnar.rs

@@ -260,6 +260,7 @@ fn run_benchmark(
                 0,
                 None,
                 &CompressionCodec::Zstd(1),
+                false,
             )
             .unwrap();
         });

native/shuffle/src/bin/shuffle_bench.rs

@@ -120,6 +120,11 @@ struct Args {
     /// Each task reads the same input and writes to its own output files.
     #[arg(long, default_value_t = 1)]
     concurrent_tasks: usize,
+
+    /// Store offsets as lengths in shuffle blocks (takes effect with zstd only), as
+    /// spark.comet.shuffle.compression.offsetsAsLengths.enabled does
+    #[arg(long, default_value_t = false)]
+    offsets_as_lengths: bool,
 }
 
 fn main() {
@@ -418,6 +423,7 @@ fn run_shuffle_write(
             args.max_buffer_bytes,
             args.limit,
             data_file.to_string(),
+            args.offsets_as_lengths,
         )
         .await
         .unwrap();
@@ -441,6 +447,7 @@ async fn execute_shuffle_write(
     max_buffer_bytes: Option<usize>,
     limit: usize,
     data_file: String,
+    offsets_as_lengths: bool,
 ) -> datafusion::common::Result<(MetricsSet, MetricsSet)> {
     let mut config = SessionConfig::new().with_batch_size(batch_size);
     // Comet never hands the shuffle writer view types: the serde maps Spark `String` to `Utf8`,
@@ -498,7 +505,8 @@ async fn execute_shuffle_write(
         write_buffer_size,
         max_buffer_bytes,
     )
-    .expect("Failed to create ShuffleWriterExec");
+    .expect("Failed to create ShuffleWriterExec")
+    .with_offsets_as_lengths(offsets_as_lengths);
 
     let task_ctx = ctx.task_ctx();
     let stream = exec.execute(0, task_ctx).unwrap();
@@ -561,6 +569,7 @@ fn run_concurrent_shuffle_writes(
             let write_buffer_size = args.write_buffer_size;
             let max_buffer_bytes = args.max_buffer_bytes;
             let limit = args.limit;
+            let offsets_as_lengths = args.offsets_as_lengths;
 
             handles.push(tokio::spawn(async move {
                 execute_shuffle_write(
@@ -573,6 +582,7 @@ fn run_concurrent_shuffle_writes(
                     max_buffer_bytes,
                     limit,
                     data_file,
+                    offsets_as_lengths,
                 )
                 .await
                 .unwrap()