comphead commented on code in PR #7258: URL: https://github.com/apache/datafusion-comet/pull/7258#discussion_r4402681924
########## spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorSuite.scala: ########## @@ -0,0 +1,122 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet.execution.shuffle + +import java.io.{ByteArrayInputStream, EOFException} + +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.execution.metric.SQLMetrics + +import org.apache.comet.Native +import org.apache.comet.vector.NativeUtil + +class NativeBatchDecoderIteratorSuite extends CometTestBase { + + import NativeShuffleTestFrames.intFrame + + private final class TrackingInputStream(bytes: Array[Byte]) Review Comment: This is a copy of `TrackingInputStream` in `CometCelebornShuffleReaderSuite`, which still has its own. Since `NativeShuffleTestFrames` was added to share fixtures between these suites, it could hold this class too. ########## spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala: ########## @@ -121,6 +124,45 @@ class CometBlockStoreShuffleReader[K, C]( }) .map(b => (0, b)) + withReadMetrics(recordIter) + } + + /** One decoder groups frames across fetched blocks for a single broadcast producer task. */ + private[shuffle] def readWithCoalescing(maxRows: Int): Iterator[Product2[K, C]] = { + val input = readAsRawStream() + val nativeUtil = + try new NativeUtil() Review Comment: The two nested `try`/`catch` blocks guard constructors that do not fail in practice. `NativeUtil` only takes allocator references and creates a `CDataDictionaryProvider`. `NativeBatchDecoderIterator` allocates an 8-byte buffer and runs `require` checks that pass for a local stream. `NonFatal` also excludes `OutOfMemoryError`. If construction did fail, `ShuffleBlockFetcherIterator` registers its own task-completion cleanup in `initialize()`. Building the decoder directly and then registering the listener, as `read()` above does, would cut this to a few lines. ########## native/core/src/execution/jni_api.rs: ########## @@ -2105,6 +2106,150 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlock( }) } +/// Decode a local frame into owned arrays. The direct buffer can be reused after this call. +fn append_shuffle_batch( + coalescer: &mut DecodedBatchCoalescer, + block: &[u8], + num_columns: usize, +) -> CometResult<()> { + let batch = read_ipc_compressed(block)?; + if batch.num_columns() != num_columns { + return Err(CometError::Internal( + "Shuffle frame column count mismatch".into(), + )); + } + coalescer.push(batch)?; + Ok(()) +} + +/// The reader owns this handle and serializes native calls with cleanup. +/// # Safety +/// A nonzero handle must refer to a live coalescer exclusively borrowed for this call. +unsafe fn shuffle_batch_coalescer<'a>(handle: jlong) -> CometResult<&'a mut DecodedBatchCoalescer> { + unsafe { (handle as *mut DecodedBatchCoalescer).as_mut() } + .ok_or_else(|| CometError::Internal("Closed shuffle batch coalescer".into())) +} + +fn shuffle_batch_ready_columns(coalescer: &DecodedBatchCoalescer) -> jint { + coalescer + .peek_ready() + .map_or(-1, |batch| batch.num_columns() as jint) +} + +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_Native_createShuffleBatchCoalescer( + e: EnvUnowned, + _class: JClass, + max_rows: jint, +) -> jlong { + try_unwrap_or_throw(&e, |_| { + if max_rows <= 0 { + return Err(CometError::Internal( + "Decoded shuffle grouping requires a positive row target".into(), + )); + } + let coalescer = DecodedBatchCoalescer::try_new(max_rows as usize)?; + Ok(Box::into_raw(Box::new(coalescer)) as jlong) + }) +} + +#[no_mangle] +/// Decode and retain one local frame, returning ready columns or -1 for more input. +/// # Safety +/// The handle must be live and exclusively owned by this reader. On error, close the reader. +pub unsafe extern "system" fn Java_org_apache_comet_Native_appendShuffleBatch( + e: EnvUnowned, + _class: JClass, + handle: jlong, + byte_buffer: JByteBuffer, + length: jint, + num_columns: jint, + tracing_enabled: jboolean, +) -> jint { + try_unwrap_or_throw(&e, |env| { + with_trace("appendShuffleBatch", tracing_enabled != JNI_FALSE, || { + let capacity = env.get_direct_buffer_capacity(&byte_buffer)?; + if length < 4 || length as usize > capacity || num_columns < 0 { Review Comment: The direct-buffer slicing here repeats `decode_shuffle_block_jni` and adds a bounds check that the existing decode path does not have. A shared helper that returns `&[u8]` for `(byte_buffer, length)` and does the capacity check would give `decodeShuffleBlock` the same protection and keep the two paths in step. In `takeShuffleBatch` below, the `array_addrs`/`schema_addrs` length check repeats the column-count check in `export_batch`. `NativeUtil.getNextBatch` allocates exactly the column count that `peekShuffleBatch` returned, so the arrays always match. ########## native/shuffle/src/decoded_batch_coalescer.rs: ########## @@ -0,0 +1,579 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Group small decoded shuffle batches before exporting them to the JVM. +//! +//! Groups are bounded by rows, retained Arrow bytes, and input batch count. An +//! oversized input passes through unchanged instead of being split or copied. +//! +//! Keep inputs until a group is ready, then concatenate once at its actual size. + +use std::collections::VecDeque; +use std::ptr::NonNull; + +use arrow::array::RecordBatch; +use arrow::buffer::Buffer; +use arrow::compute::concat_batches; +use arrow::datatypes::DataType; +use datafusion::common::{DataFusionError, Result}; + +const MAX_ROWS: usize = 65_536; +const MAX_GROUP_BATCHES: usize = 1024; +const MAX_GROUP_BYTES: usize = 8 * 1024 * 1024; + +/// Charge the shared IPC body once, but count each column's output copy separately. +/// Other shared allocations may be conservatively overcounted. +#[derive(Default)] +struct BatchBytes { + first_owner: Option<NonNull<u8>>, + retained: usize, + copied: usize, +} + +impl BatchBytes { + fn add_buffer(&mut self, buffer: &Buffer) { + let owner = buffer.data_ptr(); + if self.first_owner == Some(owner) { + return; + } + let capacity = buffer.capacity().max(buffer.len()); + if capacity != 0 { + if self.first_owner.is_none() { + self.first_owner = Some(owner); + } + self.retained = self.retained.saturating_add(capacity); + } + } +} + +/// Ordered grouping of batches returned by the native shuffle IPC decoder. +/// +/// `push` and `finish` enqueue output. Inspect it with `peek_ready`, and drain +/// it with `take_ready` before pushing another input. Call `finish` at EOF. +/// A singleton, an oversized input, or an input at least half the row target is +/// never copied. The byte budget bounds retained group inputs, not the whole +/// reader: concatenation also owns its output and may have one decoded lookahead +/// input. Callers must drop the helper on any error. +#[derive(Debug)] +pub struct DecodedBatchCoalescer { + max_rows: usize, + max_bytes: usize, + batches: Vec<RecordBatch>, + rows: usize, + bytes: usize, + ready: VecDeque<RecordBatch>, +} + +impl DecodedBatchCoalescer { + /// Creates an empty, reader-owned grouper with the production byte budget. + /// Caps positive row targets at 65536. + pub fn try_new(max_rows: usize) -> Result<Self> { + Self::try_new_with_limits(max_rows, MAX_GROUP_BYTES) + } + + /// Creates an empty grouper with explicit limits, allowing small test budgets. + /// Invalid limits return a configuration error. Keeping the byte budget at + /// most 8 MiB also keeps every concatenated STRING/BINARY offset within i32. + fn try_new_with_limits(max_rows: usize, max_bytes: usize) -> Result<Self> { + if max_rows == 0 { + return Err(DataFusionError::Configuration( + "Decoded shuffle grouping requires a positive row target".into(), + )); + } + if !(1..=MAX_GROUP_BYTES).contains(&max_bytes) { + return Err(DataFusionError::Configuration( + "Decoded shuffle grouping requires 1..8388608 bytes".into(), + )); + } + Ok(Self { + max_rows: max_rows.min(MAX_ROWS), + max_bytes, + batches: Vec::new(), + rows: 0, + bytes: 0, + ready: VecDeque::new(), + }) + } + + /// Transfers the oldest ready output to the caller without copying buffers. + /// Returns None when empty and leaves any pending group unchanged. + pub fn take_ready(&mut self) -> Option<RecordBatch> { + self.ready.pop_front() + } + + /// Borrows the oldest ready output until the next mutation, without draining it. + pub fn peek_ready(&self) -> Option<&RecordBatch> { + self.ready.front() + } + + /// Takes ownership of one decoded input. A flush followed by a passthrough can make two + /// outputs ready; drain both before reading another input. Drop this helper on error. + pub fn push(&mut self, batch: RecordBatch) -> Result<()> { + if !self.ready.is_empty() { + return Err(DataFusionError::Execution( + "Drain shuffle grouping output before reading another input".into(), + )); + } + let rows = batch.num_rows(); + let bytes = self.eligible_bytes(&batch); + + if !self.batches.is_empty() + && (bytes.is_none_or(|bytes| bytes > self.max_bytes - self.bytes) + || self.batches[0].schema() != batch.schema() + || rows > self.max_rows - self.rows) + { + self.flush()?; + } + if let Some(bytes) = bytes { + self.rows += rows; + self.bytes += bytes; + self.batches.push(batch); + if self.rows == self.max_rows + || self.bytes == self.max_bytes + || self.batches.len() == MAX_GROUP_BATCHES + { + self.flush()?; + } + } else { + self.ready.push_back(batch); + } + Ok(()) + } + + /// Flushes the final group. Repeated calls do not copy or emit it again. + pub fn finish(&mut self) -> Result<()> { + self.flush() + } + + /// Bounds both retained backing allocations and output copies for flat columns. + /// None means pass through without copying. + fn eligible_bytes(&self, batch: &RecordBatch) -> Option<usize> { + if batch.num_rows() == 0 + || batch.num_rows() >= self.max_rows.div_ceil(2) + || batch.num_columns() == 0 + { + return None; + } + let rows = batch.num_rows(); + // ArrayData reports logical bytes; bitmap builders round up to 64-byte blocks. + let bitmap_bytes = rows.div_ceil(8); + let padded_bitmap_bytes = bitmap_bytes.next_multiple_of(64); + let mut bytes = BatchBytes::default(); + for column in batch.columns() { + if !matches!( + column.data_type(), + DataType::Int32 + | DataType::Int64 + | DataType::Boolean + | DataType::Utf8 + | DataType::Binary + ) { + return None; + } + let data = column.to_data(); + let metadata = std::mem::size_of_val(column.as_ref()); + // to_data may omit an all-valid bitmap; account for the original nulls below. + let copied = + data.get_slice_memory_size().ok()? - data.nulls().map_or(0, |_| bitmap_bytes); + bytes.copied = bytes.copied.saturating_add(copied).saturating_add(metadata); + bytes.retained = bytes.retained.saturating_add(metadata); + for buffer in data.buffers() { + bytes.add_buffer(buffer); + } + if let Some(nulls) = column.nulls() { + bytes.add_buffer(nulls.buffer()); + bytes.copied = bytes.copied.saturating_add(padded_bitmap_bytes); + } + if data.data_type() == &DataType::Boolean { + bytes.copied = bytes + .copied + .saturating_add(padded_bitmap_bytes - bitmap_bytes); + } + if bytes.retained.max(bytes.copied) > self.max_bytes { + return None; + } + } + Some(bytes.retained.max(bytes.copied)) + } + + fn flush(&mut self) -> Result<()> { + if self.batches.len() < 2 { + self.ready.extend(self.batches.drain(..)); + } else { + self.ready + .push_back(concat_batches(&self.batches[0].schema(), &self.batches)?); + self.batches.clear(); + } + self.rows = 0; + self.bytes = 0; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + use arrow::array::{ + ArrayRef, AsArray, BinaryArray, BooleanArray, Int32Array, Int64Array, StringArray, + }; + use arrow::datatypes::{Field, Int32Type, Schema}; + use arrow::ipc::writer::StreamWriter; + + fn one_column(array: ArrayRef) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new( + "x", + array.data_type().clone(), + true, + )])); + RecordBatch::try_new(schema, vec![array]).unwrap() + } + + fn ints(values: impl IntoIterator<Item = i32>) -> RecordBatch { + one_column(Arc::new(Int32Array::from_iter_values(values))) + } + + fn encode_frame(batch: &RecordBatch) -> Vec<u8> { + let mut frame = b"NONE".to_vec(); + let mut writer = StreamWriter::try_new(&mut frame, batch.schema().as_ref()).unwrap(); + writer.write(batch).unwrap(); + writer.finish().unwrap(); + frame + } + + fn collect(max_rows: usize, batches: Vec<RecordBatch>) -> Vec<RecordBatch> { + collect_with_limits(max_rows, MAX_GROUP_BYTES, batches) + } + + fn collect_with_limits( + max_rows: usize, + max_bytes: usize, + batches: Vec<RecordBatch>, + ) -> Vec<RecordBatch> { + let mut coalescer = + DecodedBatchCoalescer::try_new_with_limits(max_rows, max_bytes).unwrap(); + let mut output = Vec::new(); + for batch in batches { + coalescer.push(batch).unwrap(); + while let Some(batch) = coalescer.take_ready() { + output.push(batch); + } + } + coalescer.finish().unwrap(); + while let Some(batch) = coalescer.take_ready() { + output.push(batch); + } + output + } + + #[test] + fn finishing_repeatedly_emits_the_pending_group_once() { + let mut coalescer = DecodedBatchCoalescer::try_new(8).unwrap(); + coalescer.push(ints(0..1)).unwrap(); + coalescer.push(ints(1..2)).unwrap(); + coalescer.finish().unwrap(); + coalescer.finish().unwrap(); + assert_eq!(coalescer.take_ready(), Some(ints(0..2))); + coalescer.finish().unwrap(); + assert!(coalescer.take_ready().is_none()); + } + + #[test] + fn singleton_large_and_near_full_inputs_are_not_copied() { + for batches in [ + vec![ints(0..1)], + vec![ints(0..1), ints(1..7)], + vec![ints(0..6), ints(6..7)], + vec![ints(0..1), ints(1..10)], + ] { + let originals: Vec<_> = batches + .iter() + .map(|batch| batch.column(0).clone()) + .collect(); + let output = collect(8, batches); + assert_eq!(output.len(), originals.len()); + for (batch, original) in output.iter().zip(originals) { + assert!(Arc::ptr_eq(batch.column(0), &original)); + } + } + } + + #[test] + fn row_and_input_count_limits_preserve_order() { + let batches = (0..4).map(|i| ints(i * 2..i * 2 + 2)).collect(); + assert_eq!(collect(8, batches), vec![ints(0..8)]); + assert_eq!( + collect(8, vec![ints(0..3), ints(3..6), ints(6..9)]), + vec![ints(0..6), ints(6..9)] + ); + let batches = (0..1025).map(|i| ints(i..i + 1)).collect(); + assert_eq!( + collect(MAX_ROWS, batches), + vec![ints(0..1024), ints(1024..1025)] + ); + } + + #[test] + fn byte_limit_flushes_at_and_before_crossing_the_boundary() { + let inputs: Vec<_> = (0..3).map(|i| ints(i..i + 1)).collect(); + let two_batches = inputs[0].get_array_memory_size() + inputs[1].get_array_memory_size(); + assert_eq!( + collect_with_limits(32, two_batches, inputs.clone()), + vec![ints(0..2), ints(2..3)] + ); + let output = collect_with_limits(32, two_batches - 1, inputs.clone()); + assert_eq!(output, inputs); + for (output, input) in output.iter().zip(&inputs) { + assert!(Arc::ptr_eq(output.column(0), input.column(0))); + } + + let mut coalescer = DecodedBatchCoalescer::try_new_with_limits(32, two_batches).unwrap(); + coalescer.push(inputs[0].clone()).unwrap(); + assert!(coalescer.peek_ready().is_none()); + coalescer.push(inputs[1].clone()).unwrap(); + assert_eq!(coalescer.take_ready(), Some(ints(0..2))); + coalescer.push(inputs[2].clone()).unwrap(); + assert!(coalescer.take_ready().is_none()); + coalescer.finish().unwrap(); + assert_eq!(coalescer.take_ready(), Some(ints(2..3))); + } + + #[test] + fn wide_json_values_pass_through_without_copying() { Review Comment: Several tests here cover the same paths: - This test checks that an over-budget input passes through uncopied and that grouping resumes afterwards. The second block of `byte_limit_flushes_at_and_before_crossing_the_boundary` already checks the uncopied passthrough, and `sliced_backing_buffers_count_toward_the_byte_budget` checks a passthrough between two small inputs. The one extra check, that grouping resumes after a passthrough, could be a final input in the sliced test. Nothing in it is JSON-specific. - The third block of `byte_limit_flushes_at_and_before_crossing_the_boundary`, which steps through push, peek and take, repeats the first block's `[0..2], [2..3]` result. - `ready_lookahead_and_pending_inputs_have_owned_lifetimes` covers flush-plus-passthrough, which the `[ints(0..1), ints(1..10)]` case of `singleton_large_and_near_full_inputs_are_not_copied` and the new `[3, 3, 8, 2]` JVM case also cover. Its `Weak` checks mostly confirm that Rust drops owned values. ########## spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastProducerBenchmark.scala: ########## @@ -0,0 +1,229 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.benchmark + +import java.io.{ByteArrayInputStream, ByteArrayOutputStream} +import java.nio.{ByteBuffer, ByteOrder} +import java.nio.channels.Channels +import java.nio.charset.StandardCharsets + +import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.{BigIntVector, IntVector, VarCharVector, VectorSchemaRoot} +import org.apache.arrow.vector.ipc.ArrowStreamWriter +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.comet.execution.shuffle.NativeBatchDecoderIterator +import org.apache.spark.sql.comet.util.Utils +import org.apache.spark.sql.execution.metric.SQLMetric + +import org.apache.comet.Native +import org.apache.comet.vector.NativeUtil + +/** + * Measures native shuffle decoding, Arrow export and broadcast serialization together. Input + * frames are prepared before timing; output validation and driver decoding are also untimed. Both + * arms use the same artifact and differ only in producer batching. This is a producer + * microbenchmark, not a complete Spark query or a shuffle transport benchmark. + * + * Pass --on-first to reverse the case order, --decode-only to include the decoder/FFI cost + * without serialization, and an optional scenario-name substring to filter. + * + * To run: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make \ + * benchmark-org.apache.spark.sql.benchmark.CometBroadcastProducerBenchmark BENCH_HEAP=4g + * }}} + */ +object CometBroadcastProducerBenchmark extends CometBenchmarkBase { + + private val targetRows = 8192 + private val wideValue = "x" * (6 * 1024) + + private case class Scenario(name: String, rows: Int, rowsPerFrame: Int, wide: Boolean = false) + private case class Result(rows: Long, streams: Long, bytes: Long) + + private val scenarios = Seq( + Scenario("tiny: 1 row per frame", 8192, 1), + Scenario("tiny: 8 rows per frame", 8192, 8), + Scenario("full: 8192 rows per frame", 8192, 8192), + Scenario("byte cap: 64 rows of 6 KiB strings", 8192, 64, wide = true)) + + private def frames(scenario: Scenario): Array[Byte] = { + val allocator = new RootAllocator(Long.MaxValue) + val ids = new BigIntVector("id", allocator) + val values = new IntVector("value", allocator) + val strings = new VarCharVector("payload", allocator) + val root = + if (scenario.wide) VectorSchemaRoot.of(ids, values, strings) + else VectorSchemaRoot.of(ids, values) + val output = new ByteArrayOutputStream() + val payload = wideValue.getBytes(StandardCharsets.UTF_8) + try { + ids.allocateNew(scenario.rowsPerFrame) + values.allocateNew(scenario.rowsPerFrame) + if (scenario.wide) strings.allocateNew() + var start = 0 + while (start < scenario.rows) { + val count = math.min(scenario.rowsPerFrame, scenario.rows - start) + var row = 0 + while (row < count) { + val id = start + row + ids.setSafe(row, id.toLong) + if (id % 7 == 0) values.setNull(row) else values.setSafe(row, id * 3) + if (scenario.wide) strings.setSafe(row, payload) + row += 1 + } + root.setRowCount(count) + val ipc = new ByteArrayOutputStream() + val writer = new ArrowStreamWriter(root, null, Channels.newChannel(ipc)) + try { + writer.start() + writer.writeBatch() + writer.end() + } finally writer.close() + val bytes = ipc.toByteArray + val header = ByteBuffer Review Comment: This hand-builds the same frame as `NativeShuffleTestFrames.arrowFrame`. Widening that object to `private[sql]` would let `frames` call `arrowFrame(root)` for each slice and drop the header code and the `ArrowStreamWriter` block. `produce` and `decodeOnly` also build the same decoder, and the `--on-first` mode list is the same in both benchmarks. Both scaladocs also restate the PR's measurement caveats ("not a complete Spark query or a shuffle transport benchmark", "neither TPC-H nor AQE's SMJ-to-BHJ conversion"). That context fits the PR description better than the source. A sentence on what each benchmark runs would be enough. ########## spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala: ########## @@ -42,6 +42,33 @@ import org.apache.comet.vector.CometVector class UtilsSuite extends CometTestBase { + test("serializeBatches round trips binary values across output chunks") { Review Comment: A 3 MiB payload gives a 1 MiB chunk size, the same as before `070de7f1bf`, so I expect this test to pass on `main` too. I have not run it. The new code is the small-chunk floor, so a case that serializes to more than its buffer estimate would cover it. One example is a one-row batch with a few hundred columns, where the IPC schema message alone is larger than 4 KiB. ########## native/shuffle/src/decoded_batch_coalescer.rs: ########## @@ -0,0 +1,579 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Group small decoded shuffle batches before exporting them to the JVM. +//! +//! Groups are bounded by rows, retained Arrow bytes, and input batch count. An +//! oversized input passes through unchanged instead of being split or copied. +//! +//! Keep inputs until a group is ready, then concatenate once at its actual size. + +use std::collections::VecDeque; +use std::ptr::NonNull; + +use arrow::array::RecordBatch; +use arrow::buffer::Buffer; +use arrow::compute::concat_batches; +use arrow::datatypes::DataType; +use datafusion::common::{DataFusionError, Result}; + +const MAX_ROWS: usize = 65_536; +const MAX_GROUP_BATCHES: usize = 1024; +const MAX_GROUP_BYTES: usize = 8 * 1024 * 1024; + +/// Charge the shared IPC body once, but count each column's output copy separately. +/// Other shared allocations may be conservatively overcounted. +#[derive(Default)] +struct BatchBytes { + first_owner: Option<NonNull<u8>>, + retained: usize, + copied: usize, +} + +impl BatchBytes { + fn add_buffer(&mut self, buffer: &Buffer) { + let owner = buffer.data_ptr(); + if self.first_owner == Some(owner) { + return; + } + let capacity = buffer.capacity().max(buffer.len()); + if capacity != 0 { + if self.first_owner.is_none() { + self.first_owner = Some(owner); + } + self.retained = self.retained.saturating_add(capacity); + } + } +} + +/// Ordered grouping of batches returned by the native shuffle IPC decoder. +/// +/// `push` and `finish` enqueue output. Inspect it with `peek_ready`, and drain +/// it with `take_ready` before pushing another input. Call `finish` at EOF. +/// A singleton, an oversized input, or an input at least half the row target is +/// never copied. The byte budget bounds retained group inputs, not the whole +/// reader: concatenation also owns its output and may have one decoded lookahead +/// input. Callers must drop the helper on any error. +#[derive(Debug)] +pub struct DecodedBatchCoalescer { + max_rows: usize, + max_bytes: usize, + batches: Vec<RecordBatch>, + rows: usize, + bytes: usize, + ready: VecDeque<RecordBatch>, +} + +impl DecodedBatchCoalescer { + /// Creates an empty, reader-owned grouper with the production byte budget. + /// Caps positive row targets at 65536. + pub fn try_new(max_rows: usize) -> Result<Self> { + Self::try_new_with_limits(max_rows, MAX_GROUP_BYTES) + } + + /// Creates an empty grouper with explicit limits, allowing small test budgets. + /// Invalid limits return a configuration error. Keeping the byte budget at + /// most 8 MiB also keeps every concatenated STRING/BINARY offset within i32. + fn try_new_with_limits(max_rows: usize, max_bytes: usize) -> Result<Self> { + if max_rows == 0 { + return Err(DataFusionError::Configuration( + "Decoded shuffle grouping requires a positive row target".into(), + )); + } + if !(1..=MAX_GROUP_BYTES).contains(&max_bytes) { + return Err(DataFusionError::Configuration( + "Decoded shuffle grouping requires 1..8388608 bytes".into(), + )); + } + Ok(Self { + max_rows: max_rows.min(MAX_ROWS), + max_bytes, + batches: Vec::new(), + rows: 0, + bytes: 0, + ready: VecDeque::new(), + }) + } + + /// Transfers the oldest ready output to the caller without copying buffers. + /// Returns None when empty and leaves any pending group unchanged. + pub fn take_ready(&mut self) -> Option<RecordBatch> { + self.ready.pop_front() + } + + /// Borrows the oldest ready output until the next mutation, without draining it. + pub fn peek_ready(&self) -> Option<&RecordBatch> { + self.ready.front() + } + + /// Takes ownership of one decoded input. A flush followed by a passthrough can make two + /// outputs ready; drain both before reading another input. Drop this helper on error. + pub fn push(&mut self, batch: RecordBatch) -> Result<()> { + if !self.ready.is_empty() { + return Err(DataFusionError::Execution( + "Drain shuffle grouping output before reading another input".into(), + )); + } + let rows = batch.num_rows(); + let bytes = self.eligible_bytes(&batch); + + if !self.batches.is_empty() + && (bytes.is_none_or(|bytes| bytes > self.max_bytes - self.bytes) + || self.batches[0].schema() != batch.schema() + || rows > self.max_rows - self.rows) + { + self.flush()?; + } + if let Some(bytes) = bytes { + self.rows += rows; + self.bytes += bytes; + self.batches.push(batch); + if self.rows == self.max_rows + || self.bytes == self.max_bytes + || self.batches.len() == MAX_GROUP_BATCHES + { + self.flush()?; + } + } else { + self.ready.push_back(batch); + } + Ok(()) + } + + /// Flushes the final group. Repeated calls do not copy or emit it again. + pub fn finish(&mut self) -> Result<()> { + self.flush() + } + + /// Bounds both retained backing allocations and output copies for flat columns. + /// None means pass through without copying. + fn eligible_bytes(&self, batch: &RecordBatch) -> Option<usize> { + if batch.num_rows() == 0 + || batch.num_rows() >= self.max_rows.div_ceil(2) + || batch.num_columns() == 0 + { + return None; + } + let rows = batch.num_rows(); + // ArrayData reports logical bytes; bitmap builders round up to 64-byte blocks. + let bitmap_bytes = rows.div_ceil(8); Review Comment: Is the bitmap padding accounting needed? `concat_batches` builds one null buffer per output column, so the 64-byte padding is allocated once per group, not once per input. Charging `padded_bitmap_bytes`, plus the extra Boolean padding, for every input overcounts the copy. The doc comment allows overcounting, and the 8 MiB budget is far below the i32 offset limit. Using `get_slice_memory_size()` plus metadata as the copy estimate would remove `bitmap_bytes`, `padded_bitmap_bytes`, the subtraction and the Boolean branch. `BatchBytes` could then become three locals in `eligible_bytes`. ########## native/core/src/execution/jni_api.rs: ########## @@ -2105,6 +2106,150 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlock( }) } +/// Decode a local frame into owned arrays. The direct buffer can be reused after this call. +fn append_shuffle_batch( + coalescer: &mut DecodedBatchCoalescer, + block: &[u8], + num_columns: usize, +) -> CometResult<()> { + let batch = read_ipc_compressed(block)?; + if batch.num_columns() != num_columns { + return Err(CometError::Internal( + "Shuffle frame column count mismatch".into(), + )); + } + coalescer.push(batch)?; + Ok(()) +} + +/// The reader owns this handle and serializes native calls with cleanup. +/// # Safety +/// A nonzero handle must refer to a live coalescer exclusively borrowed for this call. +unsafe fn shuffle_batch_coalescer<'a>(handle: jlong) -> CometResult<&'a mut DecodedBatchCoalescer> { + unsafe { (handle as *mut DecodedBatchCoalescer).as_mut() } + .ok_or_else(|| CometError::Internal("Closed shuffle batch coalescer".into())) +} + +fn shuffle_batch_ready_columns(coalescer: &DecodedBatchCoalescer) -> jint { + coalescer + .peek_ready() + .map_or(-1, |batch| batch.num_columns() as jint) +} + +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_Native_createShuffleBatchCoalescer( + e: EnvUnowned, + _class: JClass, + max_rows: jint, +) -> jlong { + try_unwrap_or_throw(&e, |_| { + if max_rows <= 0 { Review Comment: This repeats the zero check in `try_new_with_limits` with the same message. `DecodedBatchCoalescer::try_new(max_rows.max(0) as usize)` would map negative values to zero, so `try_new` reports every invalid value and the message lives in one place. ########## spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala: ########## @@ -73,6 +74,88 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper import testImplicits._ + test("broadcast producer batching combines local native shuffle blocks with and without AQE") { + withSQLConf( + CometConf.COMET_BATCH_SIZE.key -> "4096", + SQLConf.SHUFFLE_PARTITIONS.key -> "4", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_SIZE.key -> "1") { + for (adaptive <- Seq(false, true)) { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + val streams = scala.collection.mutable.ArrayBuffer.empty[Long] + Seq(false, true).foreach { enabled => + withSQLConf( + CometConf.COMET_BROADCAST_PRODUCER_COALESCE_ENABLED.key -> enabled.toString) { + val build = spark + .range(0, 512, 1, 16) + .selectExpr( + "CAST(id AS INT) AS k", + "CASE WHEN id % 3 = 0 THEN NULL ELSE concat('v', id) END AS v", + "CASE WHEN id % 5 = 0 THEN NULL ELSE id + 4294967296 END AS l", + "CASE WHEN id % 7 = 0 THEN NULL ELSE id % 2 = 0 END AS b", + "CASE WHEN id % 11 = 0 THEN NULL ELSE CAST(concat('b', id) AS BINARY) END AS bytes") + .repartition(col("k")) + val probe = spark.range(0, 512, 1, 2).selectExpr("CAST(id AS INT) AS k") + val df = probe.join(broadcast(build), "k") + checkAnswer( + df, + (0 until 512).map { i => + Row( + i, + if (i % 3 == 0) null else s"v$i", + if (i % 5 == 0) null else i.toLong + 4294967296L, + if (i % 7 == 0) null else i % 2 == 0, + if (i % 11 == 0) null else s"b$i".getBytes(UTF_8)) + }) + val plan = df.queryExecution.executedPlan + val exchanges = collect(plan) { case exchange: CometBroadcastExchangeExec => + exchange + } + assert(exchanges.size == 1, plan.treeString) + val exchange = exchanges.head + assert(exchange.metrics("numOutputRows").value == 512L) + assert( + collect(plan) { + case shuffle: CometShuffleExchangeExec + if shuffle.shuffleType == CometNativeShuffle => + shuffle + }.nonEmpty, + plan.treeString) + streams += exchange.metrics("numCoalescedBatches").value + () + } + } + assert(streams(1) > 0L && streams(1) < streams(0), s"AQE=$adaptive, batches=$streams") + } + } + } + } + + test("broadcast producer batching leaves unsupported schemas on the native shuffle reader") { Review Comment: I think this test passes with or without the JVM type gate. Without the gate, native `eligible_bytes` passes the DECIMAL batches through unchanged, so the answer, the row count and the native shuffle in the plan all stay the same. As written it checks a DECIMAL broadcast with the flag on, not the routing. If the gate should be pinned, extracting the predicate in `getByteArrayRdd` and asserting on it directly would do that. Otherwise this test could be dropped. -- 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]
