sunchao commented on code in PR #7258:
URL: https://github.com/apache/datafusion-comet/pull/7258#discussion_r4394915385
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala:
##########
@@ -53,31 +54,44 @@ case class NativeBatchDecoderIterator(
private var batch: Option[ColumnarBatch] = None
private val validateRemoteFrames =
in.isInstanceOf[CometShuffleReadFailureHandler]
private var remoteDecoderHandle = 0L
+ private var coalescerHandle = 0L
+ private var endOfInput = false
require(
!validateRemoteFrames || expectedSchema.exists(_ != null),
"Remote shuffle decoding requires the expected Spark schema")
+ require(!validateRemoteFrames || coalescing.isEmpty, "Batch coalescing is
local-shuffle only")
import NativeBatchDecoderIterator._
override def hasNext: Boolean = {
- synchronized {
- if (in == null || isClosed) {
- return false
- }
- if (batch.isDefined) {
- return true
- }
+ try {
+ synchronized {
+ if (in == null || isClosed) {
+ return false
+ }
+ if (batch.isDefined) {
+ return true
+ }
- // Clear ownership before cleanup so a failed close cannot release this
batch twice.
- if (currentBatch != null) {
- val previous = currentBatch
- currentBatch = null
- previous.close()
+ // Clear ownership before cleanup so a failed close cannot release
this batch twice.
+ if (currentBatch != null) {
+ val previous = currentBatch
+ currentBatch = null
+ previous.close()
+ }
}
- }
- fetchNext()
+ if (coalescing.isDefined) fetchCoalesced() else fetchNext()
+ } catch {
+ case NonFatal(failure) if coalescing.isDefined =>
Review Comment:
There is a deliberate scope difference: the coalescer can retain pending and
ready native batches after an error, so its new path drops that state
immediately. I added this rationale in `b97daacb45` and kept the ordinary
decoder's error/close behavior. Existing remote lifecycle tests distinguish the
reported decode failure from a later explicit stream-close failure; an
unconditional catch would change that established behavior. Task completion
remains the cleanup backstop for both paths.
##########
native/core/src/execution/jni_api.rs:
##########
@@ -2105,6 +2106,145 @@ 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, |_| {
+ 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("decodeShuffleBlock", tracing_enabled != JNI_FALSE, || {
Review Comment:
Renamed the span to `appendShuffleBatch` in `b97daacb45`, so traces
distinguish decode-plus-grouping from the ordinary decode path.
##########
native/shuffle/src/decoded_batch_coalescer.rs:
##########
@@ -0,0 +1,582 @@
+// 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::{Array, AsArray, RecordBatch};
+use arrow::buffer::Buffer;
+use arrow::compute::concat_batches;
+use arrow::datatypes::{DataType, Int32Type, Int64Type};
+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;
+
+/// Allocation-free accounting for one decoded batch, discarded after
admission.
+/// IPC columns normally share their message-body allocation. Remember that
first
+/// owner so its slices are charged once; other owners may be overcounted
safely.
+/// Logical copy bytes are tracked separately because columns may alias values
+/// that concatenation will materialize independently.
+#[derive(Default)]
+struct BatchBytes {
+ first_owner: Option<NonNull<u8>>,
+ retained: usize,
+ copied: usize,
+}
+
+impl BatchBytes {
+ /// Charges a borrowed buffer's allocation and the bytes its output will
copy.
+ /// Uses the allocation base, not the slice address, to deduplicate the
common
+ /// IPC owner before inspecting its capacity again. Saturating counters
+ /// conservatively reject overflow. No ownership is retained and no memory
+ /// is allocated or read from the buffer contents.
+ fn add_buffer(&mut self, buffer: &Buffer, copied: usize) {
+ self.copied = self.copied.saturating_add(copied);
+ 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);
+ }
+ }
+
+ /// Charges a borrowed concrete array's null buffer and object metadata
after
+ /// its values. The caller supplies the rounded output bitmap size. Static
+ /// dispatch avoids repeating trait-object lookups after the array
downcast;
+ /// ownership, allocation-free accounting, and saturation match
`add_buffer`.
+ #[inline]
+ fn add_nulls_and_metadata<T: Array>(&mut self, array: &T,
copied_null_bytes: usize) {
+ if let Some(nulls) = array.nulls() {
+ self.add_buffer(nulls.buffer(), copied_null_bytes);
+ }
+ let array_metadata = std::mem::size_of::<T>();
+ self.retained = self.retained.saturating_add(array_metadata);
+ self.copied = self.copied.saturating_add(array_metadata);
+ }
+}
+
+/// 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.
+ /// Rejects row targets outside 1..=65536 without retaining any buffers.
+ 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 !(1..=MAX_ROWS).contains(&max_rows) {
+ return Err(DataFusionError::Configuration(
+ "Decoded shuffle grouping requires 1..65536 rows".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_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()
+ }
+
+ /// Returns the larger of retained and logical copy bytes for a groupable
input.
+ /// Examines flat-column buffer metadata and string/binary endpoint
offsets,
+ /// never rows. The common shared IPC allocation is charged once. Other
shared
+ /// owners can be overcounted safely, and slices retain their full backing
+ /// allocation charge. RecordBatch guarantees matching field and array
types,
+ /// so dispatch borrows schema fields without cloning or inspecting array
+ /// trait objects twice. 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();
+ // Arrow bitmap builders round their allocation to 64-byte blocks.
+ let bitmap_bytes = rows.div_ceil(8).next_multiple_of(64);
+ let mut bytes = BatchBytes::default();
+ for (field, column) in
batch.schema_ref().fields().iter().zip(batch.columns()) {
+ match field.data_type() {
+ DataType::Int32 => {
+ let array = column.as_primitive::<Int32Type>();
+ bytes.add_buffer(array.values().inner(), rows * 4);
+ bytes.add_nulls_and_metadata(array, bitmap_bytes);
+ }
+ DataType::Int64 => {
+ let array = column.as_primitive::<Int64Type>();
+ bytes.add_buffer(array.values().inner(), rows * 8);
+ bytes.add_nulls_and_metadata(array, bitmap_bytes);
+ }
+ DataType::Boolean => {
+ let array = column.as_boolean();
+ bytes.add_buffer(array.values().inner(), bitmap_bytes);
+ bytes.add_nulls_and_metadata(array, bitmap_bytes);
+ }
+ DataType::Utf8 => {
+ let array = column.as_string::<i32>();
+ let offsets = array.value_offsets();
+ bytes.add_buffer(array.offsets().inner().inner(), (rows +
1) * 4);
+ bytes.add_buffer(array.values(), (offsets[rows] -
offsets[0]) as usize);
+ bytes.add_nulls_and_metadata(array, bitmap_bytes);
+ }
+ DataType::Binary => {
+ let array = column.as_binary::<i32>();
+ let offsets = array.value_offsets();
+ bytes.add_buffer(array.offsets().inner().inner(), (rows +
1) * 4);
+ bytes.add_buffer(array.values(), (offsets[rows] -
offsets[0]) as usize);
+ bytes.add_nulls_and_metadata(array, bitmap_bytes);
+ }
+ _ => return None,
+ }
+ if bytes.retained.max(bytes.copied) > self.max_bytes {
+ return None;
+ }
+ }
+ Some(bytes.retained.max(bytes.copied))
+ }
+
+ fn flush(&mut self) -> Result<()> {
+ let batches = std::mem::take(&mut self.batches);
+ self.rows = 0;
+ self.bytes = 0;
+ if batches.len() < 2 {
+ self.ready.extend(batches);
+ } else {
+ self.ready
+ .push_back(concat_batches(&batches[0].schema(), &batches)?);
+ }
+ Ok(())
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use std::sync::Arc;
+
+ use arrow::array::{ArrayRef, BinaryArray, BooleanArray, Int32Array,
Int64Array, StringArray};
+ use arrow::datatypes::{Field, 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();
+ coalescer.finish().unwrap();
+ while let Some(batch) = coalescer.take_ready() {
+ output.push(batch);
+ }
+ coalescer.finish().unwrap();
+ assert!(coalescer.peek_ready().is_none());
+ output
+ }
+
+ #[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)));
+ assert_eq!(coalescer.bytes, 0);
+ }
+
+ #[test]
+ fn wide_json_values_pass_through_without_copying() {
+ let small_json = format!("{{\"payload\":\"{}\"}}", "x".repeat(32));
+ let large_json = format!("{{\"payload\":\"{}\"}}", "x".repeat(2048));
+ let small =
one_column(Arc::new(StringArray::from(vec![small_json.as_str()])));
+ let large =
one_column(Arc::new(StringArray::from(vec![large_json.as_str()])));
+ let original = large.column(0).clone();
+ let max_bytes = 2 * small.get_array_memory_size();
+ assert!(large.get_array_memory_size() > max_bytes);
+ let output = collect_with_limits(
+ 32,
+ max_bytes,
+ vec![small.clone(), large.clone(), small.clone(), small, large],
+ );
+ let expected = one_column(Arc::new(StringArray::from(vec![
+ small_json.as_str(),
+ small_json.as_str(),
+ ])));
+ assert_eq!(output.len(), 4);
+ assert_eq!(output[0].num_rows(), 1);
+ assert!(Arc::ptr_eq(output[1].column(0), &original));
+ assert_eq!(output[2], expected);
+ assert!(Arc::ptr_eq(output[3].column(0), &original));
+ }
+
+ #[test]
+ fn sliced_backing_buffers_count_toward_the_byte_budget() {
+ let wide = "x".repeat(2048);
+ let backing = one_column(Arc::new(StringArray::from(vec!["a",
wide.as_str()])));
+ let slice = backing.slice(0, 1);
+ let small = one_column(Arc::new(StringArray::from(vec!["b"])));
+ let max_bytes = 2 * small.get_array_memory_size();
+ assert!(slice.get_array_memory_size() > max_bytes);
+ let original = slice.column(0).clone();
+ let output = collect_with_limits(32, max_bytes, vec![small.clone(),
slice, small]);
+ assert_eq!(output.len(), 3);
+ assert!(Arc::ptr_eq(output[1].column(0), &original));
+ }
+
+ #[test]
+ fn shared_ipc_body_is_charged_once_for_wide_batches() {
+ let schema = Arc::new(Schema::new(
+ (0..128)
+ .map(|column| Field::new(format!("c{column}"),
DataType::Int32, true))
+ .collect::<Vec<_>>(),
+ ));
+ let arrays: Vec<ArrayRef> = (0..128)
+ .map(|column| {
+ Arc::new(Int32Array::from_iter((0..64).map(|row| {
+ ((row + column + 1) % 17 != 0).then_some(row * 31 + column)
+ }))) as ArrayRef
+ })
+ .collect();
+ let input = RecordBatch::try_new(schema, arrays).unwrap();
+ let frame = encode_frame(&input);
+ let decoded = crate::read_ipc_compressed(&frame).unwrap();
+ let mut coalescer = DecodedBatchCoalescer::try_new(8192).unwrap();
+ let charged = coalescer.eligible_bytes(&decoded).unwrap();
+ let first_buffer = decoded
+ .column(0)
+ .as_primitive::<Int32Type>()
+ .values()
+ .inner();
+ let last_buffer = decoded
+ .column(127)
+ .as_primitive::<Int32Type>()
+ .values()
+ .inner();
+ assert_eq!(first_buffer.data_ptr(), last_buffer.data_ptr());
+ assert!(charged < 2 * first_buffer.capacity());
+ assert!(charged * 128 <= MAX_GROUP_BYTES);
+ assert!(decoded.get_array_memory_size() > MAX_GROUP_BYTES);
+
+ for _ in 0..128 {
+ coalescer
+ .push(crate::read_ipc_compressed(&frame).unwrap())
+ .unwrap();
+ }
+ let output = coalescer.take_ready().unwrap();
+ assert_eq!(output.num_rows(), 8192);
+ assert!(coalescer.take_ready().is_none());
+ for offset in (0..8192).step_by(64) {
+ assert_eq!(output.slice(offset, 64), input);
+ }
+ }
+
+ #[test]
+ fn aliased_columns_are_charged_for_each_output_copy() {
+ let array = ints(0..160).column(0).clone();
+ let schema = Arc::new(Schema::new(
+ (0..16)
+ .map(|column| Field::new(format!("c{column}"),
DataType::Int32, false))
+ .collect::<Vec<_>>(),
+ ));
+ let input = RecordBatch::try_new(schema, vec![array.clone();
16]).unwrap();
+ let metadata = 16 * std::mem::size_of_val(array.as_ref());
+ let copied = metadata + 16 * 160 * 4;
+ let retained = metadata
+ + array
+ .as_primitive::<Int32Type>()
+ .values()
+ .inner()
+ .capacity();
+ assert!(copied > retained);
+ let coalescer = DecodedBatchCoalescer::try_new(8192).unwrap();
+ assert_eq!(coalescer.eligible_bytes(&input), Some(copied));
+ let output = collect_with_limits(8192, copied * 2 - 1,
vec![input.clone(), input.clone()]);
+ assert_eq!(output, vec![input.clone(), input]);
+ for batch in output {
+ for column in batch.columns() {
+ assert!(Arc::ptr_eq(column, &array));
+ }
+ }
+ }
+
+ #[test]
+ fn invalid_byte_budgets_are_configuration_errors() {
Review Comment:
Cleaned this up in `b97daacb45`: the collection helper finishes once;
repeated finish calls have one dedicated idempotence test; invalid row/byte
limits are consolidated; and the byte-reset assertion is replaced with another
observable input/flush cycle. The backing-buffer tests now assert
grouping/passthrough and buffer identity rather than the exact internal
accounting value.
##########
native/shuffle/src/decoded_batch_coalescer.rs:
##########
@@ -0,0 +1,582 @@
+// 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::{Array, AsArray, RecordBatch};
+use arrow::buffer::Buffer;
+use arrow::compute::concat_batches;
+use arrow::datatypes::{DataType, Int32Type, Int64Type};
+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;
+
+/// Allocation-free accounting for one decoded batch, discarded after
admission.
+/// IPC columns normally share their message-body allocation. Remember that
first
+/// owner so its slices are charged once; other owners may be overcounted
safely.
+/// Logical copy bytes are tracked separately because columns may alias values
+/// that concatenation will materialize independently.
+#[derive(Default)]
+struct BatchBytes {
+ first_owner: Option<NonNull<u8>>,
+ retained: usize,
+ copied: usize,
+}
+
+impl BatchBytes {
+ /// Charges a borrowed buffer's allocation and the bytes its output will
copy.
+ /// Uses the allocation base, not the slice address, to deduplicate the
common
+ /// IPC owner before inspecting its capacity again. Saturating counters
+ /// conservatively reject overflow. No ownership is retained and no memory
+ /// is allocated or read from the buffer contents.
+ fn add_buffer(&mut self, buffer: &Buffer, copied: usize) {
+ self.copied = self.copied.saturating_add(copied);
+ 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);
+ }
+ }
+
+ /// Charges a borrowed concrete array's null buffer and object metadata
after
+ /// its values. The caller supplies the rounded output bitmap size. Static
+ /// dispatch avoids repeating trait-object lookups after the array
downcast;
+ /// ownership, allocation-free accounting, and saturation match
`add_buffer`.
+ #[inline]
+ fn add_nulls_and_metadata<T: Array>(&mut self, array: &T,
copied_null_bytes: usize) {
+ if let Some(nulls) = array.nulls() {
+ self.add_buffer(nulls.buffer(), copied_null_bytes);
+ }
+ let array_metadata = std::mem::size_of::<T>();
+ self.retained = self.retained.saturating_add(array_metadata);
+ self.copied = self.copied.saturating_add(array_metadata);
+ }
+}
+
+/// 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 {
Review Comment:
I checked the pinned Arrow 59.3.1 implementation and ran focused probes. It
eagerly copies primitive arrays, but Boolean/Utf8/Binary use
`GenericInProgressArray`, which retains source slices and concatenates at
finish. A one-row Utf8 slice retained its large backing allocation, and 1,025
one-row inputs remained buffered without an output, so the byte/input bounds
and their tests are still needed.
The half-threshold adjustment is correct, but the two-output state also
remains possible: with target 8 and bypass limit 3, inputs [2,2,4] produce two
ready outputs [4,4]. Inputs [2,4] become [6], changing this PR's large-input
passthrough behavior; primitive singletons also acquire a copy. I kept the
custom helper to preserve those requirements, while simplifying its accounting
and retaining Vec capacity in `b97daacb45`. A wrapper around Arrow would still
need those policies and would offer less simplification than expected.
##########
native/shuffle/src/decoded_batch_coalescer.rs:
##########
@@ -0,0 +1,582 @@
+// 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::{Array, AsArray, RecordBatch};
+use arrow::buffer::Buffer;
+use arrow::compute::concat_batches;
+use arrow::datatypes::{DataType, Int32Type, Int64Type};
+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;
+
+/// Allocation-free accounting for one decoded batch, discarded after
admission.
+/// IPC columns normally share their message-body allocation. Remember that
first
+/// owner so its slices are charged once; other owners may be overcounted
safely.
+/// Logical copy bytes are tracked separately because columns may alias values
+/// that concatenation will materialize independently.
+#[derive(Default)]
+struct BatchBytes {
+ first_owner: Option<NonNull<u8>>,
+ retained: usize,
+ copied: usize,
+}
+
+impl BatchBytes {
+ /// Charges a borrowed buffer's allocation and the bytes its output will
copy.
+ /// Uses the allocation base, not the slice address, to deduplicate the
common
+ /// IPC owner before inspecting its capacity again. Saturating counters
+ /// conservatively reject overflow. No ownership is retained and no memory
+ /// is allocated or read from the buffer contents.
+ fn add_buffer(&mut self, buffer: &Buffer, copied: usize) {
+ self.copied = self.copied.saturating_add(copied);
+ 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);
+ }
+ }
+
+ /// Charges a borrowed concrete array's null buffer and object metadata
after
+ /// its values. The caller supplies the rounded output bitmap size. Static
+ /// dispatch avoids repeating trait-object lookups after the array
downcast;
+ /// ownership, allocation-free accounting, and saturation match
`add_buffer`.
+ #[inline]
+ fn add_nulls_and_metadata<T: Array>(&mut self, array: &T,
copied_null_bytes: usize) {
+ if let Some(nulls) = array.nulls() {
+ self.add_buffer(nulls.buffer(), copied_null_bytes);
+ }
+ let array_metadata = std::mem::size_of::<T>();
+ self.retained = self.retained.saturating_add(array_metadata);
+ self.copied = self.copied.saturating_add(array_metadata);
+ }
+}
+
+/// 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.
+ /// Rejects row targets outside 1..=65536 without retaining any buffers.
+ 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 !(1..=MAX_ROWS).contains(&max_rows) {
+ return Err(DataFusionError::Configuration(
+ "Decoded shuffle grouping requires 1..65536 rows".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_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()
+ }
+
+ /// Returns the larger of retained and logical copy bytes for a groupable
input.
+ /// Examines flat-column buffer metadata and string/binary endpoint
offsets,
+ /// never rows. The common shared IPC allocation is charged once. Other
shared
+ /// owners can be overcounted safely, and slices retain their full backing
+ /// allocation charge. RecordBatch guarantees matching field and array
types,
+ /// so dispatch borrows schema fields without cloning or inspecting array
+ /// trait objects twice. 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();
+ // Arrow bitmap builders round their allocation to 64-byte blocks.
+ let bitmap_bytes = rows.div_ceil(8).next_multiple_of(64);
+ let mut bytes = BatchBytes::default();
+ for (field, column) in
batch.schema_ref().fields().iter().zip(batch.columns()) {
+ match field.data_type() {
+ DataType::Int32 => {
+ let array = column.as_primitive::<Int32Type>();
+ bytes.add_buffer(array.values().inner(), rows * 4);
+ bytes.add_nulls_and_metadata(array, bitmap_bytes);
+ }
+ DataType::Int64 => {
+ let array = column.as_primitive::<Int64Type>();
+ bytes.add_buffer(array.values().inner(), rows * 8);
+ bytes.add_nulls_and_metadata(array, bitmap_bytes);
+ }
+ DataType::Boolean => {
+ let array = column.as_boolean();
+ bytes.add_buffer(array.values().inner(), bitmap_bytes);
+ bytes.add_nulls_and_metadata(array, bitmap_bytes);
+ }
+ DataType::Utf8 => {
+ let array = column.as_string::<i32>();
+ let offsets = array.value_offsets();
+ bytes.add_buffer(array.offsets().inner().inner(), (rows +
1) * 4);
+ bytes.add_buffer(array.values(), (offsets[rows] -
offsets[0]) as usize);
+ bytes.add_nulls_and_metadata(array, bitmap_bytes);
+ }
+ DataType::Binary => {
+ let array = column.as_binary::<i32>();
+ let offsets = array.value_offsets();
+ bytes.add_buffer(array.offsets().inner().inner(), (rows +
1) * 4);
+ bytes.add_buffer(array.values(), (offsets[rows] -
offsets[0]) as usize);
+ bytes.add_nulls_and_metadata(array, bitmap_bytes);
+ }
+ _ => return None,
+ }
+ if bytes.retained.max(bytes.copied) > self.max_bytes {
+ return None;
+ }
+ }
+ Some(bytes.retained.max(bytes.copied))
+ }
+
+ fn flush(&mut self) -> Result<()> {
+ let batches = std::mem::take(&mut self.batches);
Review Comment:
Changed `flush` in `b97daacb45` to concatenate from the borrowed batch
vector and then clear it; singleton passthrough uses `drain(..)`. The reader
now retains the vector's allocation between groups.
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala:
##########
@@ -110,9 +115,24 @@ case class CometBroadcastExchangeExec(
private lazy val maxBroadcastRows = 512000000L
private def getByteArrayRdd(plan: SparkPlan): RDD[(Long, ChunkedByteBuffer)]
= {
- plan.executeColumnar().mapPartitionsInternal { iter =>
- Utils.serializeBatches(iter)
+ // Capture the read mode on the driver. A new RDD keeps it local to this
broadcast consumer,
+ // including when AQE or exchange reuse shares the original shuffle with
another operator.
+ val input = plan.executeColumnar()
+ val supported = plan.output.nonEmpty && plan.output.forall { attribute =>
+ attribute.dataType match {
+ case IntegerType | LongType | BooleanType | StringType | BinaryType =>
true
+ case _ => false
+ }
+ }
+ val readInput = input match {
+ case shuffle: CometShuffledBatchRDD
+ if CometConf.COMET_BROADCAST_PRODUCER_COALESCE_ENABLED.get(conf) &&
+ supported && shuffle.supportsNativeBatchCoalescing =>
+ longMetric("nativeProducerBatching") += 1L
+
shuffle.withNativeBatchCoalescing(math.min(CometConf.COMET_BATCH_SIZE.get(conf),
65536))
Review Comment:
Moved the upper clamp to native in `b97daacb45`. Scala passes
`spark.comet.batchSize` directly; native caps positive targets at `MAX_ROWS`,
and JNI rejects nonpositive values before converting to usize. The row-limit
test now checks observable passthrough at the clamped boundary.
--
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]
