From 953b9b91ca49257d246ce4f7bdd7f92055d0b7eb Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 27 Oct 2026 16:35:52 +0000 Subject: [PATCH 1/3] perf: coalesce native shuffle batches before broadcast serialization --- .../contributor-guide/memory_management.md | 6 + .../contributor-guide/native_shuffle.md | 9 + .../user-guide/latest/tuning/shuffle.md | 17 + native/core/src/execution/jni_api.rs | 170 +++++ native/shuffle/src/decoded_batch_coalescer.rs | 582 ++++++++++++++++++ native/shuffle/src/lib.rs | 1 + .../scala/org/apache/comet/CometConf.scala | 14 + .../main/scala/org/apache/comet/Native.scala | 23 + .../comet/CometBroadcastExchangeExec.scala | 24 +- .../CometBlockStoreShuffleReader.scala | 42 ++ .../shuffle/CometShuffledRowRDD.scala | 21 +- .../shuffle/NativeBatchDecoderIterator.scala | 101 ++- .../comet/exec/CometNativeShuffleSuite.scala | 74 ++- .../CometBroadcastProducerBenchmark.scala | 203 ++++++ .../CometCelebornShuffleReaderSuite.scala | 66 ++ .../CometNativeShuffleInputRDDSuite.scala | 24 + 16 files changed, 1356 insertions(+), 21 deletions(-) create mode 100644 native/shuffle/src/decoded_batch_coalescer.rs create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastProducerBenchmark.scala diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index b73f81e8968..dbd77afc239 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -120,6 +120,12 @@ the tasks of an executor that read it at the same time share one such copy. The usage log reports the allocator's total, and the part of it charged to the import allocator described below, so that the overhead can be sized for it. +The opt-in broadcast producer batching path also retains native decoded shuffle arrays outside +the execution pool. It bounds each group's retained/copy estimates to 8 MiB and 1,024 input +batches, then concatenates once before Arrow export. Inputs and output coexist during that +copy, and one decoded lookahead can be retained. The grouping limit therefore does not cap +total reader RSS, and oversized input batches pass through unchanged. + Sharing lowers the bytes that are resident, not the bytes that are reserved. Each task's native join reserves its whole build side, because DataFusion skips only the buffers that the same plan has already counted, so with N tasks sharing one copy the pool holds N times the bytes that are diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md index a9d47f0f572..76a5d1f2c6f 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -214,6 +214,15 @@ read time. See [Direct Read](#direct-read-shufflescan) below for how the choice 3. The native `ShuffleScanExec` decodes each block with `read_ipc_compressed` and feeds the resulting `RecordBatch` straight into the plan. No Arrow FFI export or import happens. +For a direct local native shuffle feeding a broadcast exchange, the opt-in +`spark.comet.exec.broadcastExchange.producerCoalesce.enabled` mode uses one decoder across the task's fetched +blocks. `DecodedBatchCoalescer` groups eligible small decoded batches before FFI export and +`Utils.serializeBatches`. The read RDD is copied with the same dependency and AQE partition +specifications, so other consumers of a reused shuffle keep their original reader. Grouping is +bounded by rows, 1,024 inputs, and 8 MiB of retained/copy buffer estimates. It does not reserve +Spark execution memory; concatenation also temporarily owns an output and a lookahead input. +Task completion releases pending native batches, exported Arrow batches, and the raw stream. + **JVM decode**, used when direct read does not apply: 1. `CometBlockStoreShuffleReader` fetches shuffle blocks via `ShuffleBlockFetcherIterator`. diff --git a/docs/source/user-guide/latest/tuning/shuffle.md b/docs/source/user-guide/latest/tuning/shuffle.md index edcf3b49706..d0751119599 100644 --- a/docs/source/user-guide/latest/tuning/shuffle.md +++ b/docs/source/user-guide/latest/tuning/shuffle.md @@ -53,6 +53,23 @@ scalar types. Hash partitioning keys must be scalar types unless and later) map keys. That setting is disabled by default until the performance of the nested hashing paths has been measured. Columns that are not partitioning keys may contain complex types like maps, structs, and arrays. +### Batching shuffle inputs to broadcast + +The experimental `spark.comet.exec.broadcastExchange.producerCoalesce.enabled` setting combines small +local native shuffle batches before broadcast serialization. It is disabled by default. It +can reduce Arrow export, serialization, and driver assembly work when a broadcast build reads +many small shuffle blocks. + +The broadcast must read directly from a local Comet native shuffle, including an AQE shuffle +read, and have only INT, BIGINT, BOOLEAN, STRING, or BINARY columns. Grouping preserves the +selected shuffle partitions and row order. Large batches pass through without a merge copy. + +Each group is limited by `spark.comet.batchSize` (at most 65,536 rows), 1,024 input batches, +and an 8 MiB estimate of retained buffers and bytes to copy. This is not a total memory limit: +concatenation temporarily holds inputs and output together, and may retain one lookahead batch. +These reader buffers are outside Spark's execution memory pool. Measure the setting on your +workload before enabling it, particularly for wide rows. + ### Columnar (JVM) Shuffle Comet Columnar shuffle is JVM-based and supports `HashPartitioning`, `RoundRobinPartitioning`, `RangePartitioning`, and diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 494e7acdff9..c875bd85b38 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -62,6 +62,7 @@ use datafusion::{ use datafusion_comet_common::decode_string_arrays; use datafusion_comet_proto::spark_expression::agg_expr::ExprStruct as AggExprStruct; use datafusion_comet_proto::spark_operator::{AggregateMode, Operator, ShuffleScan}; +use datafusion_comet_shuffle::decoded_batch_coalescer::DecodedBatchCoalescer; use datafusion_comet_spark_expr::url_funcs::{CometParseUrl, CometTryParseUrl}; use datafusion_spark::function::array::array_contains::SparkArrayContains; use datafusion_spark::function::array::repeat::SparkArrayRepeat; @@ -2165,6 +2166,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, || { + let capacity = env.get_direct_buffer_capacity(&byte_buffer)?; + if length < 4 || length as usize > capacity || num_columns < 0 { + return Err(CometError::Internal("Invalid shuffle frame bounds".into())); + } + let pointer = env.get_direct_buffer_address(&byte_buffer)?; + let block = unsafe { std::slice::from_raw_parts(pointer, length as usize) }; + let coalescer = unsafe { shuffle_batch_coalescer(handle)? }; + append_shuffle_batch(coalescer, block, num_columns as usize)?; + Ok(shuffle_batch_ready_columns(coalescer)) + }) + }) +} + +#[no_mangle] +/// Report ready output, flushing the pending group at EOF. +/// # Safety +/// The handle must be live and exclusively owned by this reader. +pub unsafe extern "system" fn Java_org_apache_comet_Native_drainShuffleBatch( + e: EnvUnowned, + _class: JClass, + handle: jlong, + end_of_input: jboolean, + tracing_enabled: jboolean, +) -> jint { + try_unwrap_or_throw(&e, |_| { + let coalescer = unsafe { shuffle_batch_coalescer(handle)? }; + if end_of_input != JNI_FALSE { + with_trace( + "coalesceShuffleBatches", + tracing_enabled != JNI_FALSE, + || coalescer.finish(), + )?; + } + Ok(shuffle_batch_ready_columns(coalescer)) + }) +} + +#[no_mangle] +/// Transfer one ready output using the ordinary export path, including offset repair. +/// # Safety +/// The handle and Arrow FFI output addresses must be live and exclusively owned by this reader. +pub unsafe extern "system" fn Java_org_apache_comet_Native_takeShuffleBatch( + e: EnvUnowned, + _class: JClass, + handle: jlong, + array_addrs: JLongArray, + schema_addrs: JLongArray, + tracing_enabled: jboolean, +) -> jlong { + try_unwrap_or_throw(&e, |env| { + with_trace("exportShuffleBatch", tracing_enabled != JNI_FALSE, || { + let batch = unsafe { shuffle_batch_coalescer(handle)? } + .take_ready() + .ok_or_else(|| CometError::Internal("No decoded shuffle batch is ready".into()))?; + if array_addrs.len(env)? != batch.num_columns() + || schema_addrs.len(env)? != batch.num_columns() + { + return Err(CometError::Internal( + "Shuffle output column count mismatch".into(), + )); + } + prepare_output(env, array_addrs, schema_addrs, batch, false) + }) + }) +} + +#[no_mangle] +/// # Safety +/// A nonzero handle must be live, must not be in use, and must be released exactly once. +pub unsafe extern "system" fn Java_org_apache_comet_Native_closeShuffleBatchCoalescer( + e: EnvUnowned, + _class: JClass, + handle: jlong, +) { + try_unwrap_or_throw(&e, |_| { + if handle != 0 { + drop(unsafe { Box::from_raw(handle as *mut DecodedBatchCoalescer) }); + } + Ok(()) + }) +} + #[no_mangle] /// Parse the expected schema once for a remote shuffle iterator. /// @@ -3398,6 +3538,36 @@ mod tests { Ok((rows, arrow::array::make_array(data))) } + #[test] + fn append_shuffle_batches_owns_decoded_input_and_preserves_values() { + let input = int_batch(); + let mut block = shuffle_block(&input); + let mut coalescer = DecodedBatchCoalescer::try_new(16).unwrap(); + append_shuffle_batch(&mut coalescer, &block, 1).unwrap(); + assert_eq!(shuffle_batch_ready_columns(&coalescer), -1); + // The JVM reuses its direct input buffer before the pending group is exported. + block.fill(0); + append_shuffle_batch(&mut coalescer, &shuffle_block(&input), 1).unwrap(); + coalescer.finish().unwrap(); + assert_eq!(shuffle_batch_ready_columns(&coalescer), 1); + let output = coalescer.take_ready().unwrap(); + assert_eq!(output.num_rows(), 6); + assert_eq!(output.slice(0, 3), input); + assert_eq!(output.slice(3, 3), input); + } + + #[test] + fn append_shuffle_batch_rejects_frame_count_mismatch() { + let mut coalescer = DecodedBatchCoalescer::try_new(16).unwrap(); + let error = + append_shuffle_batch(&mut coalescer, &shuffle_block(&int_batch()), 2).unwrap_err(); + assert!(error + .to_string() + .contains("Shuffle frame column count mismatch")); + coalescer.finish().unwrap(); + assert!(coalescer.take_ready().is_none()); + } + #[test] fn decode_shuffle_block_exports_columns() { let batch = int_batch(); diff --git a/native/shuffle/src/decoded_batch_coalescer.rs b/native/shuffle/src/decoded_batch_coalescer.rs new file mode 100644 index 00000000000..a629df74813 --- /dev/null +++ b/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>, + 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(&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::(); + 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, + rows: usize, + bytes: usize, + ready: VecDeque, +} + +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::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 { + 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 { + 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 { + 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::(); + bytes.add_buffer(array.values().inner(), rows * 4); + bytes.add_nulls_and_metadata(array, bitmap_bytes); + } + DataType::Int64 => { + let array = column.as_primitive::(); + 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::(); + 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::(); + 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) -> RecordBatch { + one_column(Arc::new(Int32Array::from_iter_values(values))) + } + + fn encode_frame(batch: &RecordBatch) -> Vec { + 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) -> Vec { + collect_with_limits(max_rows, MAX_GROUP_BYTES, batches) + } + + fn collect_with_limits( + max_rows: usize, + max_bytes: usize, + batches: Vec, + ) -> Vec { + 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::>(), + )); + let arrays: Vec = (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::() + .values() + .inner(); + let last_buffer = decoded + .column(127) + .as_primitive::() + .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::>(), + )); + 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::() + .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() { + assert!(DecodedBatchCoalescer::try_new_with_limits(8, 0).is_err()); + assert!(DecodedBatchCoalescer::try_new_with_limits(8, MAX_GROUP_BYTES + 1).is_err()); + } + + #[test] + fn nulls_and_supported_variable_width_values_keep_order() { + let families: Vec<(ArrayRef, ArrayRef, ArrayRef)> = vec![ + ( + Arc::new(Int64Array::from(vec![Some(1), None])), + Arc::new(Int64Array::from(vec![Some(2)])), + Arc::new(Int64Array::from(vec![Some(1), None, Some(2)])), + ), + ( + Arc::new(BooleanArray::from(vec![Some(true), None])), + Arc::new(BooleanArray::from(vec![Some(false)])), + Arc::new(BooleanArray::from(vec![Some(true), None, Some(false)])), + ), + ( + Arc::new(StringArray::from(vec![Some(""), None])), + Arc::new(StringArray::from(vec![Some("json")])), + Arc::new(StringArray::from(vec![Some(""), None, Some("json")])), + ), + ( + Arc::new(BinaryArray::from(vec![Some(b"".as_slice()), None])), + Arc::new(BinaryArray::from(vec![Some(b"bytes".as_slice())])), + Arc::new(BinaryArray::from(vec![ + Some(b"".as_slice()), + None, + Some(b"bytes".as_slice()), + ])), + ), + ]; + for (first, second, expected) in families { + assert_eq!( + collect(8, vec![one_column(first), one_column(second)]), + vec![one_column(expected)] + ); + } + } + + #[test] + fn schema_changes_and_unsupported_types_pass_in_order() { + let first = ints(0..1); + let renamed = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "other", + DataType::Int32, + true, + )])), + first.columns().to_vec(), + ) + .unwrap(); + let unsupported = one_column(Arc::new(arrow::array::UInt32Array::from(vec![1]))); + let batches = vec![first, renamed, unsupported, ints(1..2)]; + assert_eq!(collect(8, batches.clone()), batches); + } + + #[test] + fn ready_lookahead_and_pending_inputs_have_owned_lifetimes() { + let first = ints(0..1); + let large = ints(1..9); + let first_weak = Arc::downgrade(first.column(0)); + let large_weak = Arc::downgrade(large.column(0)); + let mut coalescer = DecodedBatchCoalescer::try_new(8).unwrap(); + coalescer.push(first).unwrap(); + assert!(coalescer.peek_ready().is_none()); + coalescer.push(large).unwrap(); + assert_eq!(coalescer.peek_ready().unwrap().num_rows(), 1); + let output = coalescer.take_ready().unwrap(); + assert!(first_weak.upgrade().is_some()); + drop(output); + assert!(first_weak.upgrade().is_none()); + assert!(large_weak.upgrade().is_some()); + let lookahead = coalescer.take_ready().unwrap(); + assert_eq!(lookahead, ints(1..9)); + drop(lookahead); + assert!(large_weak.upgrade().is_none()); + let pending = ints(10..11); + let pending_weak = Arc::downgrade(pending.column(0)); + coalescer.push(pending).unwrap(); + assert!(coalescer.peek_ready().is_none()); + drop(coalescer); + assert!(pending_weak.upgrade().is_none()); + } + + #[test] + fn empty_inputs_and_invalid_limits() { + let empty = ints(0..0); + assert_eq!(collect(8, vec![empty.clone()]), vec![empty]); + assert_eq!(collect(1, vec![ints(0..1)]), vec![ints(0..1)]); + assert!(DecodedBatchCoalescer::try_new(0).is_err()); + assert!(DecodedBatchCoalescer::try_new(MAX_ROWS + 1).is_err()); + } +} diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index b66bfcb74ae..4447ad8e1f4 100644 --- a/native/shuffle/src/lib.rs +++ b/native/shuffle/src/lib.rs @@ -20,6 +20,7 @@ pub mod bench_support; mod codec_context; pub(crate) mod comet_partitioning; +pub mod decoded_batch_coalescer; pub mod ipc; pub(crate) mod metrics; mod offsets_as_lengths; diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6003703a630..cf7dd902ae0 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -646,6 +646,20 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(true) + val COMET_BROADCAST_PRODUCER_COALESCE_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.exec.broadcastExchange.producerCoalesce.enabled") + .category(CATEGORY_EXEC) + .doc( + "Experimental batching for direct local native shuffle reads feeding a broadcast " + + "exchange. Combines small INT, BIGINT, BOOLEAN, STRING and BINARY batches before " + + "Arrow export and serialization, with at most min(spark.comet.batchSize, 65536) " + + "rows, 1024 input batches and an 8 MiB decoded-buffer budget per group. Large " + + "batches pass through without copying. The budget bounds grouped inputs, not " + + "total memory: concatenation also owns its output and one lookahead input. " + + "Other input paths and driver broadcast assembly are unchanged.") + .booleanConf + .createWithDefault(false) + val COMET_FORCE_SHJ: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.forceShuffledHashJoin") .withAlternative(s"$COMET_EXEC_CONFIG_PREFIX.replaceSortMergeJoin") diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index f960785647b..2ac9d7a7222 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -303,6 +303,29 @@ class Native extends NativeBase { schemaAddrs: Array[Long], tracingEnabled: Boolean): Long + /** Own a bounded group of decoded local shuffle batches until export or reader cleanup. */ + @native def createShuffleBatchCoalescer(maxRows: Int): Long + + /** Decode one frame and return the next output's column count, or -1 for more input. */ + @native def appendShuffleBatch( + handle: Long, + shuffleBlock: ByteBuffer, + length: Int, + numColumns: Int, + tracingEnabled: Boolean): Int + + /** Return ready columns, flushing any pending group at EOF. */ + @native def drainShuffleBatch(handle: Long, endOfInput: Boolean, tracingEnabled: Boolean): Int + + /** Transfer one ready batch through Arrow FFI, returning its row count. */ + @native def takeShuffleBatch( + handle: Long, + arrayAddrs: Array[Long], + schemaAddrs: Array[Long], + tracingEnabled: Boolean): Long + + @native def closeShuffleBatchCoalescer(handle: Long): Unit + /** * Create a remote shuffle decoder that retains the expected Spark types for one iterator. The * expected schema is serialized as a ShuffleScan protobuf. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala index e106fd5c67f..5c1d86e2e74 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala @@ -32,12 +32,14 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.plans.logical.Statistics import org.apache.spark.sql.catalyst.plans.physical.{BroadcastMode, BroadcastPartitioning, Partitioning} +import org.apache.spark.sql.comet.execution.shuffle.CometShuffledBatchRDD import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, SQLExecution} import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} +import org.apache.spark.sql.types.{BinaryType, BooleanType, IntegerType, LongType, StringType} import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.spark.util.{SparkFatalException, ThreadUtils} import org.apache.spark.util.io.ChunkedByteBuffer @@ -77,6 +79,9 @@ case class CometBroadcastExchangeExec( "collectTime" -> SQLMetrics.createTimingMetric(sparkContext, "time to collect"), "buildTime" -> SQLMetrics.createTimingMetric(sparkContext, "time to build"), "broadcastTime" -> SQLMetrics.createTimingMetric(sparkContext, "time to broadcast"), + "nativeProducerBatching" -> SQLMetrics.createMetric( + sparkContext, + "broadcast exchanges using native producer batching"), "numCoalescedBatches" -> SQLMetrics.createMetric( sparkContext, "number of coalesced batches for broadcast"), @@ -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)) + case _ => input } + readInput.mapPartitionsInternal(iter => Utils.serializeBatches(iter)) } def getNumPartitions(): Int = { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala index c5064029b1a..e34d7875b2d 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala @@ -22,6 +22,8 @@ package org.apache.spark.sql.comet.execution.shuffle import java.io.{FilterInputStream, InputStream} import java.util.concurrent.atomic.AtomicBoolean +import scala.util.control.NonFatal + import org.apache.spark.{InterruptibleIterator, MapOutputTracker, SparkEnv, TaskContext} import org.apache.spark.internal.{config, Logging} import org.apache.spark.io.CompressionCodec @@ -34,6 +36,7 @@ import org.apache.spark.storage.BlockManager import org.apache.spark.storage.BlockManagerId import org.apache.spark.storage.ShuffleBlockFetcherIterator import org.apache.spark.util.CompletionIterator +import org.apache.spark.util.Utils import org.apache.comet.{CometConf, Native} import org.apache.comet.vector.NativeUtil @@ -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() + catch { + case NonFatal(failure) => Utils.tryWithSafeFinally(throw failure)(input.close()) + } + val decoder = + try { + NativeBatchDecoderIterator( + input, + dep.decodeTime, + new Native(), + nativeUtil, + CometConf.COMET_TRACING_ENABLED.get(), + readBufferSize = CometConf.COMET_SHUFFLE_READ_BUFFER_SIZE.get().toInt, + coalescing = Some(maxRows)) + } catch { + case NonFatal(failure) => + Utils.tryWithSafeFinally(throw failure) { + Utils.tryWithSafeFinally(input.close())(nativeUtil.close()) + } + } + // Registration invokes the callback immediately if completion raced construction. The + // decoder is lazy, so a closed task cannot acquire native batching state after cleanup. + context.addTaskCompletionListener[Unit] { _ => + Utils.tryWithSafeFinally { + Utils.tryWithSafeFinally(decoder.close())(nativeUtil.close()) + }(context.taskMetrics().mergeShuffleReadMetrics()) + } + withReadMetrics(decoder.map(batch => (0, batch))) + } + + private def withReadMetrics( + recordIter: Iterator[(Int, ColumnarBatch)]): Iterator[Product2[K, C]] = { // Update the context task metrics for each record read. val metricIter = CompletionIterator[(Any, Any), Iterator[(Any, Any)]]( recordIter.map { record => diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala index dddb9767761..fcea4da68da 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala @@ -37,7 +37,8 @@ import org.apache.comet.{CometConf, CometShuffleBlockIterator} class CometShuffledBatchRDD( var dependency: ShuffleDependency[Int, _, _], val metrics: Map[String, SQLMetric], - partitionSpecs: Array[ShufflePartitionSpec]) + partitionSpecs: Array[ShufflePartitionSpec], + private[comet] val nativeBatchCoalescing: Option[Int] = None) extends RDD[ColumnarBatch](dependency.rdd.context, Nil) { def this(dependency: ShuffleDependency[Int, _, _], metrics: Map[String, SQLMetric]) = { @@ -51,6 +52,16 @@ class CometShuffledBatchRDD( SortShuffleManager.FETCH_SHUFFLE_BLOCKS_IN_BATCH_ENABLED_KEY, SQLConf.get.fetchShuffleBlocksInBatch.toString) + /** Only the local native shuffle decoder supports producer batching. */ + private[comet] def supportsNativeBatchCoalescing: Boolean = dependency match { + case native: CometShuffleDependency[_, _, _] if native.shuffleType == CometNativeShuffle => + native.shuffleHandle.isInstanceOf[CometNativeShuffleHandle[_, _]] + case _ => false + } + + private[comet] def withNativeBatchCoalescing(maxRows: Int): CometShuffledBatchRDD = + new CometShuffledBatchRDD(dependency, metrics, partitionSpecs.clone(), Some(maxRows)) + // Start without waiting so constructing a join or union can start every independent input. // Spark resolves dependencies on the submitting thread before queueing a downstream job; // freeze the selected destination there, before Spark caches this RDD's dependency graph. @@ -182,8 +193,12 @@ class CometShuffledBatchRDD( override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { val reader: CometShuffleReader[_, _] = createReader(split, context)._1 - // TODO: Reads IPC by native code - reader.read().asInstanceOf[Iterator[Product2[Int, ColumnarBatch]]].map(_._2) + val batches = (reader, nativeBatchCoalescing) match { + case (local: CometBlockStoreShuffleReader[_, _], Some(maxRows)) => + local.readWithCoalescing(maxRows) + case _ => reader.read() + } + batches.asInstanceOf[Iterator[Product2[Int, ColumnarBatch]]].map(_._2) } override def clearDependencies(): Unit = { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala index b74732a8c11..b8a86de0782 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala @@ -42,7 +42,8 @@ case class NativeBatchDecoderIterator( nativeUtil: NativeUtil, tracingEnabled: Boolean, expectedSchema: Option[Array[Byte]] = None, - readBufferSize: Int = CometShuffleBlockIterator.DEFAULT_READ_BUFFER_SIZE) + readBufferSize: Int = CometShuffleBlockIterator.DEFAULT_READ_BUFFER_SIZE, + coalescing: Option[Int] = None) extends Iterator[ColumnarBatch] { // One consumer reads this iterator, while task completion may close it from another thread. @@ -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 => + try close() + catch { + case NonFatal(cleanupFailure) => + if (cleanupFailure ne failure) failure.addSuppressed(cleanupFailure) + } + throw failure + } } def next(): ColumnarBatch = { @@ -163,6 +177,60 @@ case class NativeBatchDecoderIterator( } } + /** Keep decoded arrays native until a bounded group is ready for a single FFI export. */ + private def fetchCoalesced(): Boolean = { + var columns = synchronized { + if (isClosed) return false + if (coalescerHandle == 0L) { + coalescerHandle = nativeLib.createShuffleBatchCoalescer(coalescing.get) + } + nativeLib.drainShuffleBatch(coalescerHandle, endOfInput, tracingEnabled) + } + while (columns < 0 && !endOfInput) { + // As on the ordinary path, transport reads must not prevent task completion from closing + // the input. Native state is touched only under the same monitor as close(). + val block = readNextBlock() + synchronized { + if (isClosed) return false + val startTime = System.nanoTime() + try { + columns = block match { + case Some((fieldCount, dataBuf, bytesToRead)) => + nativeLib.appendShuffleBatch( + coalescerHandle, + dataBuf, + bytesToRead, + fieldCount, + tracingEnabled) + case None => + endOfInput = true + nativeLib.drainShuffleBatch(coalescerHandle, true, tracingEnabled) + } + } finally decodeTime.add(System.nanoTime() - startTime) + } + } + synchronized { + if (isClosed) return false + if (columns < 0) { + close() + false + } else { + val startTime = System.nanoTime() + try { + batch = nativeUtil.getNextBatch( + columns, + (arrayAddrs, schemaAddrs) => + nativeLib.takeShuffleBatch( + coalescerHandle, + arrayAddrs, + schemaAddrs, + tracingEnabled)) + batch.isDefined + } finally decodeTime.add(System.nanoTime() - startTime) + } + } + } + private def readNextBlock(): Option[(Int, ByteBuffer, Int)] = { // read compressed batch size from header longBuf.clear() @@ -229,6 +297,8 @@ case class NativeBatchDecoderIterator( batch = None val decoderHandle = remoteDecoderHandle remoteDecoderHandle = 0L + val groupingHandle = coalescerHandle + coalescerHandle = 0L var failure: Throwable = null def release(resource: => Unit): Unit = { @@ -243,6 +313,7 @@ case class NativeBatchDecoderIterator( if (previous != null) release(previous.close()) prefetched.filterNot(_ eq previous).foreach(pending => release(pending.close())) if (decoderHandle != 0L) release(nativeLib.releaseRemoteShuffleDecoder(decoderHandle)) + if (groupingHandle != 0L) release(nativeLib.closeShuffleBatchCoalescer(groupingHandle)) if (in != null) release(in.close()) release(resetDataBuf()) if (failure != null) throw failure diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index 06660cef2a6..b2711082dea 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -41,7 +41,7 @@ import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, import org.apache.spark.sql.catalyst.expressions.aggregate.Final import org.apache.spark.sql.catalyst.plans.logical.LocalRelation import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RangePartitioning} -import org.apache.spark.sql.comet.{CometExec, CometHashAggregateExec, CometLocalTableScanExec, CometMetricNode, CometNativeExec, CometNativeScanExec, CometScanWrapper, CometSparkToColumnarExec, CometTakeOrderedAndProjectExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometExec, CometHashAggregateExec, CometLocalTableScanExec, CometMetricNode, CometNativeExec, CometNativeScanExec, CometScanWrapper, CometSparkToColumnarExec, CometTakeOrderedAndProjectExec} import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, CometNativeArrowSource} import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.execution.{FileSourceScanExec, LocalTableScanExec, SparkPlan} @@ -73,6 +73,78 @@ 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") + .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"))) + 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("nativeProducerBatching").value == (if (enabled) 1L else 0L)) + 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") { + withSQLConf( + CometConf.COMET_BROADCAST_PRODUCER_COALESCE_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val build = spark + .range(0, 32, 1, 4) + .selectExpr("CAST(id AS INT) AS k", "CAST(id AS DECIMAL(10, 2)) AS v") + .repartition(2, col("k")) + val probe = spark.range(0, 32, 1, 2).selectExpr("CAST(id AS INT) AS k") + val df = probe.join(broadcast(build), "k") + checkAnswer(df, (0 until 32).map(i => Row(i, BigDecimal(i).setScale(2)))) + val plan = df.queryExecution.executedPlan + val exchanges = collect(plan) { case exchange: CometBroadcastExchangeExec => exchange } + assert(exchanges.size == 1, plan.treeString) + assert(exchanges.head.metrics("nativeProducerBatching").value == 0L) + assert( + collect(plan) { + case shuffle: CometShuffleExchangeExec if shuffle.shuffleType == CometNativeShuffle => + shuffle + }.nonEmpty, + plan.treeString) + } + } + /** * Runs `f` with nested hash partitioning keys enabled. The config is disabled by default, so * every test that expects a nested key to reach native shuffle has to opt in. diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastProducerBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastProducerBenchmark.scala new file mode 100644 index 00000000000..b9eea51c363 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastProducerBenchmark.scala @@ -0,0 +1,203 @@ +/* + * 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 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 + .allocate(20) + .order(ByteOrder.LITTLE_ENDIAN) + .putLong(12L + bytes.length.toLong) + .putLong(root.getFieldVectors.size().toLong) + .put(Array[Byte](78, 79, 78, 69)) // NONE codec, followed by Arrow IPC. + output.write(header.array()) + output.write(bytes) + start += count + } + output.toByteArray + } finally { + root.close() + if (!scenario.wide) strings.close() + allocator.close() + } + } + + private def produce( + input: Array[Byte], + scenario: Scenario, + enabled: Boolean, + verify: Boolean): Result = { + val util = new NativeUtil + val decoder = NativeBatchDecoderIterator( + new ByteArrayInputStream(input), + new SQLMetric("nsTiming", 0L), + new Native, + util, + tracingEnabled = false, + coalescing = if (enabled) Some(targetRows) else None) + var rows = 0L + var streams = 0L + var bytes = 0L + var checkedRows = 0L + try { + Utils.serializeBatches(decoder).foreach { case (count, buffer) => + try { + rows += count + streams += 1L + bytes += buffer.size + if (verify) { + val batches = Utils.decodeBatches(buffer, "broadcast-producer-benchmark") + while (batches.hasNext) { + val batch = batches.next() + var row = 0 + while (row < batch.numRows()) { + assert(batch.column(0).getLong(row) == checkedRows) + assert(batch.column(1).isNullAt(row) == (checkedRows % 7 == 0)) + if (checkedRows % 7 != 0) { + assert(batch.column(1).getInt(row) == checkedRows.toInt * 3) + } + if (scenario.wide) { + assert(batch.column(2).getUTF8String(row).toString == wideValue) + } + checkedRows += 1L + row += 1 + } + } + } + } finally buffer.dispose() + } + assert(rows == scenario.rows.toLong) + if (verify) assert(checkedRows == rows) + Result(rows, streams, bytes) + } finally { + try decoder.close() + finally util.close() + } + } + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + // Initialize SparkEnv, whose compression codec serializeBatches uses. + spark.sparkContext.setLogLevel("WARN") + val scenarioFilters = mainArgs.filterNot(_ == "--on-first") + val modes = if (mainArgs.contains("--on-first")) { + Seq("ON" -> true, "OFF" -> false, "ON (repeat)" -> true) + } else { + Seq("OFF" -> false, "ON" -> true, "OFF (repeat)" -> false) + } + scenarios.filter(s => scenarioFilters.forall(s.name.contains)).foreach { scenario => + val input = frames(scenario) + runBenchmark(scenario.name) { + val benchmark = + new Benchmark(scenario.name, scenario.rows.toLong, minNumIters = 5, output = output) + val off = produce(input, scenario, enabled = false, verify = true) + val on = produce(input, scenario, enabled = true, verify = true) + benchmark.out.println(s"OFF: $off; ON: $on") + if (scenario.rowsPerFrame == targetRows) { + assert(on.streams == off.streams) + } else { + assert(on.streams < off.streams) + } + modes.foreach { case (name, enabled) => + benchmark.addCase(name) { _ => + val result = produce(input, scenario, enabled, verify = false) + assert(result == (if (enabled) on else off)) + } + } + benchmark.run() + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala index 0e618b96088..f6c38b4f0d3 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala @@ -1263,6 +1263,72 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { } } + test("local decoder coalesces across frames and flushes the final partial batch") { + val frames = Array.fill(8)(intFrame()).flatten + for (coalescing <- Seq(None, Some(16))) { + val nativeUtil = new org.apache.comet.vector.NativeUtil() + val input = new TrackingInputStream(frames) + val decoder = NativeBatchDecoderIterator( + input, + SQLMetrics.createNanoTimingMetric(spark.sparkContext, "decode time"), + new Native(), + nativeUtil, + tracingEnabled = false, + coalescing = coalescing) + try { + val sizes = scala.collection.mutable.ArrayBuffer.empty[Int] + val values = scala.collection.mutable.ArrayBuffer.empty[Int] + while (decoder.hasNext) { + val batch = decoder.next() + sizes += batch.numRows() + values ++= (0 until batch.numRows()).map(batch.column(0).getInt) + } + assert(values.toSeq == Seq.fill(8)(Seq(-1, 0, 1)).flatten) + assert(sizes.toSeq == (if (coalescing.isDefined) Seq(15, 9) else Seq.fill(8)(3))) + assert(input.closeCalls == 1) + } finally { + decoder.close() + nativeUtil.close() + } + } + } + + test("coalescing decoder releases native state on early close and truncated input") { + for (truncated <- Seq(false, true)) { + var releases = 0 + val native = new Native() { + override def closeShuffleBatchCoalescer(handle: Long): Unit = { + releases += 1 + super.closeShuffleBatchCoalescer(handle) + } + } + val complete = Array.fill(8)(intFrame()).flatten + val input = new TrackingInputStream( + if (truncated) complete.take(complete.length - 1) + else complete) + val nativeUtil = new org.apache.comet.vector.NativeUtil() + val decoder = NativeBatchDecoderIterator( + input, + SQLMetrics.createNanoTimingMetric(spark.sparkContext, "decode time"), + native, + nativeUtil, + tracingEnabled = false, + coalescing = Some(if (truncated) 128 else 16)) + try { + if (truncated) intercept[EOFException](decoder.hasNext) + else assert(decoder.hasNext) // Close with both exported and pending native batches. + decoder.close() + decoder.close() + assert(releases == 1) + assert(input.closeCalls == 1) + assert(!decoder.hasNext) + } finally { + decoder.close() + nativeUtil.close() + } + } + } + private def dictionaryFrame(): Array[Byte] = { val allocator = new RootAllocator(Long.MaxValue) val encoding = new DictionaryEncoding(0L, false, new ArrowType.Int(32, true)) diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala index 871255b1ecc..d959786c487 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala @@ -24,6 +24,7 @@ import org.apache.spark.rdd.{DeterministicLevel, RDD} import org.apache.spark.serializer.JavaSerializer import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.comet.{CometMetricNode, NativeExecContext} +import org.apache.spark.sql.execution.{CoalescedMapperPartitionSpec, CoalescedPartitionSpec, PartialMapperPartitionSpec, PartialReducerPartitionSpec, ShufflePartitionSpec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch @@ -45,6 +46,29 @@ import org.apache.comet.serde.OperatorOuterClass.Operator */ class CometNativeShuffleInputRDDSuite extends CometTestBase { + test("broadcast producer batching preserves shared shuffle dependencies and AQE ranges") { + val input = spark.sparkContext.parallelize(Seq.empty[(Int, ColumnarBatch)], 4) + val dependency = new CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]( + input, + new HashPartitioner(4), + decodeTime = SQLMetrics.createMetric(spark.sparkContext, "decode time")) + val specs = Array[ShufflePartitionSpec]( + CoalescedPartitionSpec(1, 3), + PartialReducerPartitionSpec(2, 1, 3, 100L), + PartialMapperPartitionSpec(2, 1, 3), + CoalescedMapperPartitionSpec(1, 3, 4)) + val original = new CometShuffledBatchRDD(dependency, Map.empty, specs) + val batched = original.withNativeBatchCoalescing(4096) + assert(original.supportsNativeBatchCoalescing) + assert(batched.dependency eq original.dependency) + assert(batched.dependencies == original.dependencies) + assert(original.nativeBatchCoalescing.isEmpty) + assert(batched.nativeBatchCoalescing.contains(4096)) + assert( + batched.partitions.map(_.asInstanceOf[ShuffledRowRDDPartition].spec).sameElements(specs)) + assert(batched.partitioner == original.partitioner) + } + test("native shuffle input preserves its parents' determinism") { Seq( DeterministicLevel.DETERMINATE, From de679cdaaf2304a52fc7f14d7a194abb6ba59454 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 27 Oct 2026 22:29:59 +0000 Subject: [PATCH 2/3] refactor: simplify broadcast producer batching and cover decoder ordering --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../contributor-guide/memory_management.md | 7 +- .../contributor-guide/native_shuffle.md | 5 +- .../user-guide/latest/tuning/shuffle.md | 8 +- native/core/src/execution/jni_api.rs | 9 +- native/shuffle/src/decoded_batch_coalescer.rs | 185 +++++++++--------- .../scala/org/apache/comet/CometConf.scala | 13 +- .../main/scala/org/apache/comet/Native.scala | 2 +- .../comet/CometBroadcastExchangeExec.scala | 8 +- .../shuffle/NativeBatchDecoderIterator.scala | 6 +- .../comet/exec/CometNativeShuffleSuite.scala | 21 +- .../CometBroadcastLocalShuffleBenchmark.scala | 126 ++++++++++++ .../CometBroadcastProducerBenchmark.scala | 30 ++- .../CometCelebornShuffleReaderSuite.scala | 114 +---------- .../NativeBatchDecoderIteratorSuite.scala | 122 ++++++++++++ .../shuffle/NativeShuffleTestFrames.scala | 76 +++++++ 17 files changed, 492 insertions(+), 242 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastLocalShuffleBenchmark.scala create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorSuite.scala create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeShuffleTestFrames.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 52e86f59ffd..381df39aa95 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -543,6 +543,7 @@ jobs: org.apache.spark.sql.comet.execution.shuffle.CometNativePositionalRoundRobinSuite org.apache.spark.sql.comet.execution.shuffle.CometDiskBlockWriterSuite org.apache.spark.sql.comet.execution.shuffle.CometBlockStoreShuffleReaderSuite + org.apache.spark.sql.comet.execution.shuffle.NativeBatchDecoderIteratorSuite org.apache.comet.exec.CometShuffleEncryptionSuite org.apache.comet.exec.CometShuffleManagerSuite org.apache.comet.exec.DisableAQECometShuffleSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index b1d89ec1240..373e530dbfa 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -209,6 +209,7 @@ jobs: org.apache.spark.sql.comet.execution.shuffle.CometNativePositionalRoundRobinSuite org.apache.spark.sql.comet.execution.shuffle.CometDiskBlockWriterSuite org.apache.spark.sql.comet.execution.shuffle.CometBlockStoreShuffleReaderSuite + org.apache.spark.sql.comet.execution.shuffle.NativeBatchDecoderIteratorSuite org.apache.comet.exec.CometShuffleEncryptionSuite org.apache.comet.exec.CometShuffleManagerSuite org.apache.comet.exec.DisableAQECometShuffleSuite diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index dbd77afc239..e023fd535d4 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -121,9 +121,10 @@ usage log reports the allocator's total, and the part of it charged to the impor described below, so that the overhead can be sized for it. The opt-in broadcast producer batching path also retains native decoded shuffle arrays outside -the execution pool. It bounds each group's retained/copy estimates to 8 MiB and 1,024 input -batches, then concatenates once before Arrow export. Inputs and output coexist during that -copy, and one decoded lookahead can be retained. The grouping limit therefore does not cap +the execution pool. Each group is bounded by `min(spark.comet.batchSize, 65536)` rows, 8 MiB of +retained/copy estimates, and 1,024 input batches, then concatenated once before Arrow export. +Inputs and output coexist during that copy, and one decoded lookahead can be retained. +The grouping limit therefore does not cap total reader RSS, and oversized input batches pass through unchanged. Sharing lowers the bytes that are resident, not the bytes that are reserved. Each task's native diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md index 76a5d1f2c6f..629ce8f02c3 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -219,9 +219,8 @@ For a direct local native shuffle feeding a broadcast exchange, the opt-in blocks. `DecodedBatchCoalescer` groups eligible small decoded batches before FFI export and `Utils.serializeBatches`. The read RDD is copied with the same dependency and AQE partition specifications, so other consumers of a reused shuffle keep their original reader. Grouping is -bounded by rows, 1,024 inputs, and 8 MiB of retained/copy buffer estimates. It does not reserve -Spark execution memory; concatenation also temporarily owns an output and a lookahead input. -Task completion releases pending native batches, exported Arrow batches, and the raw stream. +bounded as described in [Memory Management](memory_management.md). Task completion releases +pending native batches, exported Arrow batches, and the raw stream. **JVM decode**, used when direct read does not apply: diff --git a/docs/source/user-guide/latest/tuning/shuffle.md b/docs/source/user-guide/latest/tuning/shuffle.md index d0751119599..0bd63b285c7 100644 --- a/docs/source/user-guide/latest/tuning/shuffle.md +++ b/docs/source/user-guide/latest/tuning/shuffle.md @@ -64,11 +64,9 @@ The broadcast must read directly from a local Comet native shuffle, including an read, and have only INT, BIGINT, BOOLEAN, STRING, or BINARY columns. Grouping preserves the selected shuffle partitions and row order. Large batches pass through without a merge copy. -Each group is limited by `spark.comet.batchSize` (at most 65,536 rows), 1,024 input batches, -and an 8 MiB estimate of retained buffers and bytes to copy. This is not a total memory limit: -concatenation temporarily holds inputs and output together, and may retain one lookahead batch. -These reader buffers are outside Spark's execution memory pool. Measure the setting on your -workload before enabling it, particularly for wide rows. +Reader buffers are outside Spark's execution memory pool. Measure the setting on your workload +before enabling it, particularly for wide rows. The contributor guide describes the +[grouping limits and memory lifetime](../../../contributor-guide/memory_management.md). ### Columnar (JVM) Shuffle diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index c875bd85b38..90c1b666a24 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -2203,6 +2203,11 @@ pub extern "system" fn Java_org_apache_comet_Native_createShuffleBatchCoalescer( 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) }) @@ -2222,7 +2227,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_appendShuffleBatch( tracing_enabled: jboolean, ) -> jint { try_unwrap_or_throw(&e, |env| { - with_trace("decodeShuffleBlock", tracing_enabled != JNI_FALSE, || { + 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 { return Err(CometError::Internal("Invalid shuffle frame bounds".into())); @@ -2240,7 +2245,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_appendShuffleBatch( /// Report ready output, flushing the pending group at EOF. /// # Safety /// The handle must be live and exclusively owned by this reader. -pub unsafe extern "system" fn Java_org_apache_comet_Native_drainShuffleBatch( +pub unsafe extern "system" fn Java_org_apache_comet_Native_peekShuffleBatch( e: EnvUnowned, _class: JClass, handle: jlong, diff --git a/native/shuffle/src/decoded_batch_coalescer.rs b/native/shuffle/src/decoded_batch_coalescer.rs index a629df74813..a535afe6f1d 100644 --- a/native/shuffle/src/decoded_batch_coalescer.rs +++ b/native/shuffle/src/decoded_batch_coalescer.rs @@ -25,21 +25,18 @@ use std::collections::VecDeque; use std::ptr::NonNull; -use arrow::array::{Array, AsArray, RecordBatch}; +use arrow::array::RecordBatch; use arrow::buffer::Buffer; use arrow::compute::concat_batches; -use arrow::datatypes::{DataType, Int32Type, Int64Type}; +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; -/// 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. +/// 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>, @@ -48,13 +45,7 @@ struct BatchBytes { } 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); + fn add_buffer(&mut self, buffer: &Buffer) { let owner = buffer.data_ptr(); if self.first_owner == Some(owner) { return; @@ -67,20 +58,6 @@ impl BatchBytes { 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(&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::(); - 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. @@ -103,7 +80,7 @@ pub struct DecodedBatchCoalescer { impl DecodedBatchCoalescer { /// Creates an empty, reader-owned grouper with the production byte budget. - /// Rejects row targets outside 1..=65536 without retaining any buffers. + /// Caps positive row targets at 65536. pub fn try_new(max_rows: usize) -> Result { Self::try_new_with_limits(max_rows, MAX_GROUP_BYTES) } @@ -112,9 +89,9 @@ impl DecodedBatchCoalescer { /// 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 { - if !(1..=MAX_ROWS).contains(&max_rows) { + if max_rows == 0 { return Err(DataFusionError::Configuration( - "Decoded shuffle grouping requires 1..65536 rows".into(), + "Decoded shuffle grouping requires a positive row target".into(), )); } if !(1..=MAX_GROUP_BYTES).contains(&max_bytes) { @@ -123,7 +100,7 @@ impl DecodedBatchCoalescer { )); } Ok(Self { - max_rows, + max_rows: max_rows.min(MAX_ROWS), max_bytes, batches: Vec::new(), rows: 0, @@ -182,13 +159,8 @@ impl DecodedBatchCoalescer { 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. + /// Bounds both retained backing allocations and output copies for flat columns. + /// None means pass through without copying. fn eligible_bytes(&self, batch: &RecordBatch) -> Option { if batch.num_rows() == 0 || batch.num_rows() >= self.max_rows.div_ceil(2) @@ -197,41 +169,39 @@ impl DecodedBatchCoalescer { 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); + // 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 (field, column) in batch.schema_ref().fields().iter().zip(batch.columns()) { - match field.data_type() { - DataType::Int32 => { - let array = column.as_primitive::(); - bytes.add_buffer(array.values().inner(), rows * 4); - bytes.add_nulls_and_metadata(array, bitmap_bytes); - } - DataType::Int64 => { - let array = column.as_primitive::(); - 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::(); - 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::(); - 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, + 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; @@ -241,15 +211,15 @@ impl DecodedBatchCoalescer { } 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); + if self.batches.len() < 2 { + self.ready.extend(self.batches.drain(..)); } else { self.ready - .push_back(concat_batches(&batches[0].schema(), &batches)?); + .push_back(concat_batches(&self.batches[0].schema(), &self.batches)?); + self.batches.clear(); } + self.rows = 0; + self.bytes = 0; Ok(()) } } @@ -259,8 +229,10 @@ mod tests { use super::*; use std::sync::Arc; - use arrow::array::{ArrayRef, BinaryArray, BooleanArray, Int32Array, Int64Array, StringArray}; - use arrow::datatypes::{Field, Schema}; + 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 { @@ -303,15 +275,24 @@ mod tests { } } 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 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 [ @@ -366,7 +347,10 @@ mod tests { 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); + 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] @@ -408,6 +392,16 @@ mod tests { assert!(Arc::ptr_eq(output[1].column(0), &original)); } + #[test] + fn all_valid_null_backing_counts_toward_the_byte_budget() { + let nulls = arrow::buffer::NullBuffer::new_valid(16 * 1024).slice(0, 1); + let input = one_column(Arc::new(Int32Array::new(vec![1].into(), Some(nulls)))); + assert!(input.column(0).to_data().nulls().is_none()); + let output = collect_with_limits(8, 512, vec![ints(0..1), input.clone(), ints(2..3)]); + assert_eq!(output, vec![ints(0..1), input.clone(), ints(2..3)]); + assert!(Arc::ptr_eq(output[1].column(0), input.column(0))); + } + #[test] fn shared_ipc_body_is_charged_once_for_wide_batches() { let schema = Arc::new(Schema::new( @@ -426,7 +420,6 @@ mod tests { 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::() @@ -438,8 +431,6 @@ mod tests { .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 { @@ -473,8 +464,6 @@ mod tests { .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 { @@ -485,11 +474,21 @@ mod tests { } #[test] - fn invalid_byte_budgets_are_configuration_errors() { + fn invalid_limits_are_configuration_errors() { + assert!(DecodedBatchCoalescer::try_new(0).is_err()); assert!(DecodedBatchCoalescer::try_new_with_limits(8, 0).is_err()); assert!(DecodedBatchCoalescer::try_new_with_limits(8, MAX_GROUP_BYTES + 1).is_err()); } + #[test] + fn large_row_targets_are_clamped() { + let batch = ints(0..(MAX_ROWS / 2) as i32); + let mut coalescer = DecodedBatchCoalescer::try_new(usize::MAX).unwrap(); + coalescer.push(batch.clone()).unwrap(); + let output = coalescer.take_ready().unwrap(); + assert!(Arc::ptr_eq(output.column(0), batch.column(0))); + } + #[test] fn nulls_and_supported_variable_width_values_keep_order() { let families: Vec<(ArrayRef, ArrayRef, ArrayRef)> = vec![ @@ -572,11 +571,9 @@ mod tests { } #[test] - fn empty_inputs_and_invalid_limits() { + fn empty_inputs_and_unit_row_target_pass_through() { let empty = ints(0..0); assert_eq!(collect(8, vec![empty.clone()]), vec![empty]); assert_eq!(collect(1, vec![ints(0..1)]), vec![ints(0..1)]); - assert!(DecodedBatchCoalescer::try_new(0).is_err()); - assert!(DecodedBatchCoalescer::try_new(MAX_ROWS + 1).is_err()); } } diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index cf7dd902ae0..a89cdfb5b2b 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -649,14 +649,11 @@ object CometConf extends ShimCometConf { val COMET_BROADCAST_PRODUCER_COALESCE_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.exec.broadcastExchange.producerCoalesce.enabled") .category(CATEGORY_EXEC) - .doc( - "Experimental batching for direct local native shuffle reads feeding a broadcast " + - "exchange. Combines small INT, BIGINT, BOOLEAN, STRING and BINARY batches before " + - "Arrow export and serialization, with at most min(spark.comet.batchSize, 65536) " + - "rows, 1024 input batches and an 8 MiB decoded-buffer budget per group. Large " + - "batches pass through without copying. The budget bounds grouped inputs, not " + - "total memory: concatenation also owns its output and one lookahead input. " + - "Other input paths and driver broadcast assembly are unchanged.") + .doc("Experimental batching for direct local native shuffle reads feeding a broadcast " + + "exchange. Combines small INT, BIGINT, BOOLEAN, STRING and BINARY batches before " + + "Arrow export and serialization. Large batches pass through without copying. " + + "Measure the setting on your workload before enabling it, particularly for wide rows. " + + s"$TUNING_GUIDE.") .booleanConf .createWithDefault(false) diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index 2ac9d7a7222..c24bd762c3e 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -315,7 +315,7 @@ class Native extends NativeBase { tracingEnabled: Boolean): Int /** Return ready columns, flushing any pending group at EOF. */ - @native def drainShuffleBatch(handle: Long, endOfInput: Boolean, tracingEnabled: Boolean): Int + @native def peekShuffleBatch(handle: Long, endOfInput: Boolean, tracingEnabled: Boolean): Int /** Transfer one ready batch through Arrow FFI, returning its row count. */ @native def takeShuffleBatch( diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala index 5c1d86e2e74..93c0a65b5eb 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala @@ -79,9 +79,6 @@ case class CometBroadcastExchangeExec( "collectTime" -> SQLMetrics.createTimingMetric(sparkContext, "time to collect"), "buildTime" -> SQLMetrics.createTimingMetric(sparkContext, "time to build"), "broadcastTime" -> SQLMetrics.createTimingMetric(sparkContext, "time to broadcast"), - "nativeProducerBatching" -> SQLMetrics.createMetric( - sparkContext, - "broadcast exchanges using native producer batching"), "numCoalescedBatches" -> SQLMetrics.createMetric( sparkContext, "number of coalesced batches for broadcast"), @@ -118,6 +115,8 @@ case class CometBroadcastExchangeExec( // 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() + // Skip known unsupported schemas without paying for coalescer JNI calls. Native checks + // still handle physical representations such as dictionary-encoded strings. val supported = plan.output.nonEmpty && plan.output.forall { attribute => attribute.dataType match { case IntegerType | LongType | BooleanType | StringType | BinaryType => true @@ -128,8 +127,7 @@ case class CometBroadcastExchangeExec( 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)) + shuffle.withNativeBatchCoalescing(CometConf.COMET_BATCH_SIZE.get(conf)) case _ => input } readInput.mapPartitionsInternal(iter => Utils.serializeBatches(iter)) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala index b8a86de0782..16ec3e87ea6 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala @@ -84,6 +84,8 @@ case class NativeBatchDecoderIterator( if (coalescing.isDefined) fetchCoalesced() else fetchNext() } catch { + // A failed append/import may leave pending native batches. Release them immediately, + // while preserving the ordinary reader's existing failure and stream-close behavior. case NonFatal(failure) if coalescing.isDefined => try close() catch { @@ -184,7 +186,7 @@ case class NativeBatchDecoderIterator( if (coalescerHandle == 0L) { coalescerHandle = nativeLib.createShuffleBatchCoalescer(coalescing.get) } - nativeLib.drainShuffleBatch(coalescerHandle, endOfInput, tracingEnabled) + nativeLib.peekShuffleBatch(coalescerHandle, endOfInput, tracingEnabled) } while (columns < 0 && !endOfInput) { // As on the ordinary path, transport reads must not prevent task completion from closing @@ -204,7 +206,7 @@ case class NativeBatchDecoderIterator( tracingEnabled) case None => endOfInput = true - nativeLib.drainShuffleBatch(coalescerHandle, true, tracingEnabled) + nativeLib.peekShuffleBatch(coalescerHandle, true, tracingEnabled) } } finally decodeTime.add(System.nanoTime() - startTime) } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index b2711082dea..72af2960583 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -21,6 +21,7 @@ package org.apache.comet.exec import java.io.{ByteArrayInputStream, IOException} import java.nio.{ByteBuffer, ByteOrder} +import java.nio.charset.StandardCharsets.UTF_8 import scala.collection.mutable import scala.collection.mutable.ListBuffer @@ -90,19 +91,29 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper .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 % 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"))) + 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("nativeProducerBatching").value == (if (enabled) 1L else 0L)) assert(exchange.metrics("numOutputRows").value == 512L) assert( collect(plan) { @@ -135,7 +146,7 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper val plan = df.queryExecution.executedPlan val exchanges = collect(plan) { case exchange: CometBroadcastExchangeExec => exchange } assert(exchanges.size == 1, plan.treeString) - assert(exchanges.head.metrics("nativeProducerBatching").value == 0L) + assert(exchanges.head.metrics("numOutputRows").value == 32L) assert( collect(plan) { case shuffle: CometShuffleExchangeExec if shuffle.shuffleType == CometNativeShuffle => diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastLocalShuffleBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastLocalShuffleBenchmark.scala new file mode 100644 index 00000000000..796819b3267 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastLocalShuffleBenchmark.scala @@ -0,0 +1,126 @@ +/* + * 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 org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.{DataFrame, Row} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec} +import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.functions.{broadcast, col, count, sum} +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf + +/** + * A complete local query with a native shuffle on the broadcast build side. Explicit repartition + * and broadcast keep the plan stable; this measures neither TPC-H nor AQE's SMJ-to-BHJ + * conversion. Every iteration generates both inputs, writes and reads the shuffle, builds and + * distributes the broadcast, joins, and aggregates the result. No stage or input is cached + * between iterations. + * + * Before timing, check results, native join/shuffle execution, producer batching activation, and + * the number of Arrow streams reaching driver assembly. Pass --on-first to reverse case order. + * + * To run with four local workers: + * {{{ + * COMET_BENCHMARK_MASTER='local[4]' SPARK_GENERATE_BENCHMARK_FILES=1 make \ + * benchmark-org.apache.spark.sql.benchmark.CometBroadcastLocalShuffleBenchmark BENCH_HEAP=4g + * }}} + */ +object CometBroadcastLocalShuffleBenchmark extends CometBenchmarkBase { + + private val buildRows = 65536L + private val probeRows = buildRows * 16 + + private def query(): DataFrame = { + val build = spark + .range(0, buildRows, 1, 128) + .selectExpr( + "CAST(id AS INT) AS k", + "CASE WHEN id % 7 = 0 THEN NULL ELSE CAST(id * 3 AS INT) END AS v") + .repartition(64, col("k")) + val probe = spark + .range(0, probeRows, 1, 8) + .selectExpr(s"CAST(id % $buildRows AS INT) AS k") + probe.join(broadcast(build), "k").agg(count("*"), sum("v")) + } + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + spark.sparkContext.setLogLevel("WARN") + withSQLConf( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_BATCH_SIZE.key -> "8192", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val expectedSum = (0L until buildRows).filter(_ % 7 != 0).map(_ * 3).sum * 16 + val expected = Seq(Row(probeRows, expectedSum)) + val streams = Seq(false, true).map { enabled => + var count = 0L + withSQLConf(CometConf.COMET_BROADCAST_PRODUCER_COALESCE_ENABLED.key -> enabled.toString) { + val df = query() + assert(df.collect().toSeq == expected) + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case join: CometBroadcastHashJoinExec => join }.nonEmpty, plan) + assert( + collect(plan) { + case shuffle: CometShuffleExchangeExec + if shuffle.shuffleType == CometNativeShuffle => + shuffle + }.nonEmpty, + plan) + val exchanges = collect(plan) { case exchange: CometBroadcastExchangeExec => exchange } + assert(exchanges.size == 1, plan) + val exchange = exchanges.head + assert(exchange.metrics("numOutputRows").value == buildRows) + count = exchange.metrics("numCoalescedBatches").value + println(s"enabled=$enabled, broadcast input streams=$count, result=$expected") + println(plan.treeString) + } + count + } + assert(streams(1) > 0 && streams(1) < streams(0), streams) + + val modes = if (mainArgs.contains("--on-first")) { + Seq("ON" -> true, "OFF" -> false, "ON (repeat)" -> true) + } else { + Seq("OFF" -> false, "ON" -> true, "OFF (repeat)" -> false) + } + runBenchmark("local native shuffle to broadcast join") { + val benchmark = new Benchmark( + "local native shuffle to broadcast join", + probeRows, + minNumIters = 5, + output = output) + modes.foreach { case (name, enabled) => + benchmark.addCase(name) { _ => + withSQLConf( + CometConf.COMET_BROADCAST_PRODUCER_COALESCE_ENABLED.key -> enabled.toString) { + assert(query().collect().toSeq == expected) + } + } + } + benchmark.run() + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastProducerBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastProducerBenchmark.scala index b9eea51c363..5b36f093d07 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastProducerBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastProducerBenchmark.scala @@ -41,7 +41,8 @@ import org.apache.comet.vector.NativeUtil * 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 and an optional scenario-name substring to filter. + * 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: * {{{ @@ -168,10 +169,29 @@ object CometBroadcastProducerBenchmark extends CometBenchmarkBase { } } + private def decodeOnly(input: Array[Byte], scenario: Scenario, enabled: Boolean): Unit = { + val util = new NativeUtil + val decoder = NativeBatchDecoderIterator( + new ByteArrayInputStream(input), + new SQLMetric("nsTiming", 0L), + new Native, + util, + tracingEnabled = false, + coalescing = if (enabled) Some(targetRows) else None) + var rows = 0L + try { + while (decoder.hasNext) rows += decoder.next().numRows() + assert(rows == scenario.rows.toLong) + } finally { + try decoder.close() + finally util.close() + } + } + override def runCometBenchmark(mainArgs: Array[String]): Unit = { // Initialize SparkEnv, whose compression codec serializeBatches uses. spark.sparkContext.setLogLevel("WARN") - val scenarioFilters = mainArgs.filterNot(_ == "--on-first") + val scenarioFilters = mainArgs.filterNot(Set("--on-first", "--decode-only")) val modes = if (mainArgs.contains("--on-first")) { Seq("ON" -> true, "OFF" -> false, "ON (repeat)" -> true) } else { @@ -196,6 +216,12 @@ object CometBroadcastProducerBenchmark extends CometBenchmarkBase { assert(result == (if (enabled) on else off)) } } + if (mainArgs.contains("--decode-only")) { + benchmark.addCase("decode only OFF") { _ => + decodeOnly(input, scenario, enabled = false) + } + benchmark.addCase("decode only ON") { _ => decodeOnly(input, scenario, enabled = true) } + } benchmark.run() } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala index f6c38b4f0d3..9ca690f3b69 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala @@ -35,7 +35,7 @@ import org.apache.arrow.flatbuf.{Int => IpcInt, Message, MessageHeader, Schema = import org.apache.arrow.memory.RootAllocator import org.apache.arrow.vector.{FieldVector, IntVector, VarCharVector, VectorSchemaRoot} import org.apache.arrow.vector.complex.{ListVector, MapVector} -import org.apache.arrow.vector.dictionary.{Dictionary, DictionaryProvider} +import org.apache.arrow.vector.dictionary.Dictionary import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider import org.apache.arrow.vector.ipc.ArrowStreamWriter import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, Field, FieldType} @@ -56,17 +56,9 @@ import org.apache.comet.serde.{OperatorOuterClass, QueryPlanSerde} class CometCelebornShuffleReaderSuite extends CometTestBase { + import NativeShuffleTestFrames._ import testImplicits._ - private def frameBytes(payload: Byte*): Array[Byte] = ByteBuffer - .allocate(20 + payload.size) - .order(ByteOrder.LITTLE_ENDIAN) - .putLong(12L + payload.size) - .putLong(0L) - .put(Array[Byte](78, 79, 78, 69)) // NONE codec; raw-reader tests need only a framed payload. - .put(payload.toArray) - .array() - private final class TrackingInputStream(bytes: Array[Byte]) extends ByteArrayInputStream(bytes) { var closeCalls = 0 @@ -1227,108 +1219,6 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { } } - private def arrowFrame( - root: VectorSchemaRoot, - dictionaries: DictionaryProvider = null): Array[Byte] = { - val output = new ByteArrayOutputStream() - val writer = new ArrowStreamWriter(root, dictionaries, Channels.newChannel(output)) - val payload = - try { - writer.start() - writer.writeBatch() - writer.end() - output.toByteArray - } finally writer.close() - val frame = frameBytes(payload: _*) - ByteBuffer - .wrap(frame) - .order(ByteOrder.LITTLE_ENDIAN) - .putLong(8, root.getFieldVectors.size().toLong) - frame - } - - private def intFrame(): Array[Byte] = { - val allocator = new RootAllocator(Long.MaxValue) - val values = new IntVector("value", allocator) - val root = VectorSchemaRoot.of(values) - try { - values.allocateNew() - Seq(-1, 0, 1).zipWithIndex.foreach { case (value, index) => values.setSafe(index, value) } - values.setValueCount(3) - root.setRowCount(3) - arrowFrame(root) - } finally { - root.close() - allocator.close() - } - } - - test("local decoder coalesces across frames and flushes the final partial batch") { - val frames = Array.fill(8)(intFrame()).flatten - for (coalescing <- Seq(None, Some(16))) { - val nativeUtil = new org.apache.comet.vector.NativeUtil() - val input = new TrackingInputStream(frames) - val decoder = NativeBatchDecoderIterator( - input, - SQLMetrics.createNanoTimingMetric(spark.sparkContext, "decode time"), - new Native(), - nativeUtil, - tracingEnabled = false, - coalescing = coalescing) - try { - val sizes = scala.collection.mutable.ArrayBuffer.empty[Int] - val values = scala.collection.mutable.ArrayBuffer.empty[Int] - while (decoder.hasNext) { - val batch = decoder.next() - sizes += batch.numRows() - values ++= (0 until batch.numRows()).map(batch.column(0).getInt) - } - assert(values.toSeq == Seq.fill(8)(Seq(-1, 0, 1)).flatten) - assert(sizes.toSeq == (if (coalescing.isDefined) Seq(15, 9) else Seq.fill(8)(3))) - assert(input.closeCalls == 1) - } finally { - decoder.close() - nativeUtil.close() - } - } - } - - test("coalescing decoder releases native state on early close and truncated input") { - for (truncated <- Seq(false, true)) { - var releases = 0 - val native = new Native() { - override def closeShuffleBatchCoalescer(handle: Long): Unit = { - releases += 1 - super.closeShuffleBatchCoalescer(handle) - } - } - val complete = Array.fill(8)(intFrame()).flatten - val input = new TrackingInputStream( - if (truncated) complete.take(complete.length - 1) - else complete) - val nativeUtil = new org.apache.comet.vector.NativeUtil() - val decoder = NativeBatchDecoderIterator( - input, - SQLMetrics.createNanoTimingMetric(spark.sparkContext, "decode time"), - native, - nativeUtil, - tracingEnabled = false, - coalescing = Some(if (truncated) 128 else 16)) - try { - if (truncated) intercept[EOFException](decoder.hasNext) - else assert(decoder.hasNext) // Close with both exported and pending native batches. - decoder.close() - decoder.close() - assert(releases == 1) - assert(input.closeCalls == 1) - assert(!decoder.hasNext) - } finally { - decoder.close() - nativeUtil.close() - } - } - } - private def dictionaryFrame(): Array[Byte] = { val allocator = new RootAllocator(Long.MaxValue) val encoding = new DictionaryEncoding(0L, false, new ArrowType.Int(32, true)) diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorSuite.scala new file mode 100644 index 00000000000..14c52c788da --- /dev/null +++ b/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]) + extends ByteArrayInputStream(bytes) { + var closeCalls = 0 + + override def close(): Unit = { + closeCalls += 1 + super.close() + } + } + + test("local decoder groups frames, drains passthrough output, and flushes the final batch") { + val scenarios = Seq( + (Seq.fill(8)(3), Seq(15, 9)), + // The 8-row frame flushes the pending 6 rows and queues a second, uncopied output. + (Seq(3, 3, 8, 2), Seq(6, 8, 2))) + for { + (rowCounts, coalescedSizes) <- scenarios + coalescing <- Seq(None, Some(16)) + } { + val offsets = rowCounts.scanLeft(0)(_ + _) + val frames = offsets + .zip(rowCounts) + .flatMap { case (offset, count) => + intFrame(offset until offset + count) + } + .toArray + val nativeUtil = new NativeUtil() + val input = new TrackingInputStream(frames) + val decoder = NativeBatchDecoderIterator( + input, + SQLMetrics.createNanoTimingMetric(spark.sparkContext, "decode time"), + new Native(), + nativeUtil, + tracingEnabled = false, + coalescing = coalescing) + try { + val sizes = scala.collection.mutable.ArrayBuffer.empty[Int] + val values = scala.collection.mutable.ArrayBuffer.empty[Int] + while (decoder.hasNext) { + val batch = decoder.next() + sizes += batch.numRows() + values ++= (0 until batch.numRows()).map(batch.column(0).getInt) + } + assert(values.toSeq == (0 until rowCounts.sum)) + assert(sizes.toSeq == (if (coalescing.isDefined) coalescedSizes else rowCounts)) + assert(input.closeCalls == 1) + } finally { + decoder.close() + nativeUtil.close() + } + } + } + + test("coalescing decoder releases native state on early close and truncated input") { + for (truncated <- Seq(false, true)) { + var releases = 0 + val native = new Native() { + override def closeShuffleBatchCoalescer(handle: Long): Unit = { + releases += 1 + super.closeShuffleBatchCoalescer(handle) + } + } + val complete = Array.fill(8)(intFrame()).flatten + val input = new TrackingInputStream( + if (truncated) complete.take(complete.length - 1) + else complete) + val nativeUtil = new NativeUtil() + val decoder = NativeBatchDecoderIterator( + input, + SQLMetrics.createNanoTimingMetric(spark.sparkContext, "decode time"), + native, + nativeUtil, + tracingEnabled = false, + coalescing = Some(if (truncated) 128 else 16)) + try { + if (truncated) intercept[EOFException](decoder.hasNext) + else assert(decoder.hasNext) // Close with both exported and pending native batches. + decoder.close() + decoder.close() + assert(releases == 1) + assert(input.closeCalls == 1) + assert(!decoder.hasNext) + } finally { + decoder.close() + nativeUtil.close() + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeShuffleTestFrames.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeShuffleTestFrames.scala new file mode 100644 index 00000000000..139f5e11609 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeShuffleTestFrames.scala @@ -0,0 +1,76 @@ +/* + * 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.ByteArrayOutputStream +import java.nio.{ByteBuffer, ByteOrder} +import java.nio.channels.Channels + +import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.{IntVector, VectorSchemaRoot} +import org.apache.arrow.vector.dictionary.DictionaryProvider +import org.apache.arrow.vector.ipc.ArrowStreamWriter + +/** Local and remote readers consume the same framed Arrow IPC format. */ +private[shuffle] object NativeShuffleTestFrames { + + def frameBytes(payload: Byte*): Array[Byte] = ByteBuffer + .allocate(20 + payload.size) + .order(ByteOrder.LITTLE_ENDIAN) + .putLong(12L + payload.size) + .putLong(0L) + .put(Array[Byte](78, 79, 78, 69)) // NONE codec + .put(payload.toArray) + .array() + + def arrowFrame(root: VectorSchemaRoot, dictionaries: DictionaryProvider = null): Array[Byte] = { + val output = new ByteArrayOutputStream() + val writer = new ArrowStreamWriter(root, dictionaries, Channels.newChannel(output)) + val payload = + try { + writer.start() + writer.writeBatch() + writer.end() + output.toByteArray + } finally writer.close() + val frame = frameBytes(payload: _*) + ByteBuffer + .wrap(frame) + .order(ByteOrder.LITTLE_ENDIAN) + .putLong(8, root.getFieldVectors.size().toLong) + frame + } + + def intFrame(input: Seq[Int] = Seq(-1, 0, 1)): Array[Byte] = { + val allocator = new RootAllocator(Long.MaxValue) + val values = new IntVector("value", allocator) + val root = VectorSchemaRoot.of(values) + try { + values.allocateNew() + input.zipWithIndex.foreach { case (value, index) => values.setSafe(index, value) } + values.setValueCount(input.size) + root.setRowCount(input.size) + arrowFrame(root) + } finally { + root.close() + allocator.close() + } + } +} From 48e1024bfc90aa146fd7bb809a655123e41408f4 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 28 Oct 2026 21:39:24 +0000 Subject: [PATCH 3/3] refactor: simplify broadcast batching readers and tests --- .../contributor-guide/memory_management.md | 5 +- .../user-guide/latest/tuning/shuffle.md | 3 +- native/core/src/execution/jni_api.rs | 41 +++++----- native/shuffle/src/decoded_batch_coalescer.rs | 77 +++---------------- .../CometBlockStoreShuffleReader.scala | 32 +++----- .../comet/exec/CometNativeShuffleSuite.scala | 24 ------ .../sql/benchmark/CometBenchmarkBase.scala | 8 ++ .../CometBroadcastLocalShuffleBenchmark.scala | 16 +--- .../CometBroadcastProducerBenchmark.scala | 61 +++++---------- .../CometCelebornShuffleReaderSuite.scala | 10 --- .../CometNativeShuffleInputRDDSuite.scala | 24 ------ .../NativeBatchDecoderIteratorSuite.scala | 14 +--- .../shuffle/NativeShuffleTestFrames.scala | 13 +++- 13 files changed, 87 insertions(+), 241 deletions(-) diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index e023fd535d4..4e7248f2ba3 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -123,9 +123,8 @@ described below, so that the overhead can be sized for it. The opt-in broadcast producer batching path also retains native decoded shuffle arrays outside the execution pool. Each group is bounded by `min(spark.comet.batchSize, 65536)` rows, 8 MiB of retained/copy estimates, and 1,024 input batches, then concatenated once before Arrow export. -Inputs and output coexist during that copy, and one decoded lookahead can be retained. -The grouping limit therefore does not cap -total reader RSS, and oversized input batches pass through unchanged. +Inputs and output coexist during that copy, and one decoded lookahead can be retained. The grouping +limit therefore does not cap total reader RSS, and oversized input batches pass through unchanged. Sharing lowers the bytes that are resident, not the bytes that are reserved. Each task's native join reserves its whole build side, because DataFusion skips only the buffers that the same plan diff --git a/docs/source/user-guide/latest/tuning/shuffle.md b/docs/source/user-guide/latest/tuning/shuffle.md index 0bd63b285c7..4c86a2dcb89 100644 --- a/docs/source/user-guide/latest/tuning/shuffle.md +++ b/docs/source/user-guide/latest/tuning/shuffle.md @@ -64,8 +64,7 @@ The broadcast must read directly from a local Comet native shuffle, including an read, and have only INT, BIGINT, BOOLEAN, STRING, or BINARY columns. Grouping preserves the selected shuffle partitions and row order. Large batches pass through without a merge copy. -Reader buffers are outside Spark's execution memory pool. Measure the setting on your workload -before enabling it, particularly for wide rows. The contributor guide describes the +Reader buffers are outside Spark's execution memory pool. The contributor guide describes the [grouping limits and memory lifetime](../../../contributor-guide/memory_management.md). ### Columnar (JVM) Shuffle diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 90c1b666a24..50df7a786b0 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -2203,12 +2203,7 @@ pub extern "system" fn Java_org_apache_comet_Native_createShuffleBatchCoalescer( 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)?; + let coalescer = DecodedBatchCoalescer::try_new(max_rows.max(0) as usize)?; Ok(Box::into_raw(Box::new(coalescer)) as jlong) }) } @@ -2228,12 +2223,10 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_appendShuffleBatch( ) -> 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 { + if num_columns < 0 { return Err(CometError::Internal("Invalid shuffle frame bounds".into())); } - let pointer = env.get_direct_buffer_address(&byte_buffer)?; - let block = unsafe { std::slice::from_raw_parts(pointer, length as usize) }; + let block = unsafe { shuffle_frame_bytes(env, &byte_buffer, length)? }; let coalescer = unsafe { shuffle_batch_coalescer(handle)? }; append_shuffle_batch(coalescer, block, num_columns as usize)?; Ok(shuffle_batch_ready_columns(coalescer)) @@ -2269,6 +2262,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_peekShuffleBatch( /// Transfer one ready output using the ordinary export path, including offset repair. /// # Safety /// The handle and Arrow FFI output addresses must be live and exclusively owned by this reader. +/// Both address arrays must have the column count reported by `peekShuffleBatch`. pub unsafe extern "system" fn Java_org_apache_comet_Native_takeShuffleBatch( e: EnvUnowned, _class: JClass, @@ -2282,13 +2276,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_takeShuffleBatch( let batch = unsafe { shuffle_batch_coalescer(handle)? } .take_ready() .ok_or_else(|| CometError::Internal("No decoded shuffle batch is ready".into()))?; - if array_addrs.len(env)? != batch.num_columns() - || schema_addrs.len(env)? != batch.num_columns() - { - return Err(CometError::Internal( - "Shuffle output column count mismatch".into(), - )); - } prepare_output(env, array_addrs, schema_addrs, batch, false) }) }) @@ -2396,6 +2383,22 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlockWit }) } +/// Borrows a codec header and IPC payload from a direct buffer. +/// # Safety +/// The buffer's backing memory must remain valid and unmodified for the borrow. +unsafe fn shuffle_frame_bytes<'a>( + env: &mut Env, + byte_buffer: &'a JByteBuffer<'_>, + length: jint, +) -> CometResult<&'a [u8]> { + let capacity = env.get_direct_buffer_capacity(byte_buffer)?; + if length < 4 || length as usize > capacity { + return Err(CometError::Internal("Invalid shuffle frame bounds".into())); + } + let pointer = env.get_direct_buffer_address(byte_buffer)?; + Ok(unsafe { std::slice::from_raw_parts(pointer, length as usize) }) +} + /// Converts the JNI arguments of the `decodeShuffleBlock` entry points for /// [`decode_shuffle_block`]. fn decode_shuffle_block_jni( @@ -2406,9 +2409,7 @@ fn decode_shuffle_block_jni( schema_addrs: JLongArray, expected_types: Option<&[ArrowDataType]>, ) -> CometResult { - let raw_pointer = env.get_direct_buffer_address(&byte_buffer)?; - let length = length as usize; - let block: &[u8] = unsafe { std::slice::from_raw_parts(raw_pointer, length) }; + let block = unsafe { shuffle_frame_bytes(env, &byte_buffer, length)? }; let array_addrs = unsafe { array_addrs.get_elements(env, ReleaseMode::NoCopyBack)? }; let schema_addrs = unsafe { schema_addrs.get_elements(env, ReleaseMode::NoCopyBack)? }; unsafe { decode_shuffle_block(block, &array_addrs, &schema_addrs, expected_types) } diff --git a/native/shuffle/src/decoded_batch_coalescer.rs b/native/shuffle/src/decoded_batch_coalescer.rs index a535afe6f1d..91d902a4892 100644 --- a/native/shuffle/src/decoded_batch_coalescer.rs +++ b/native/shuffle/src/decoded_batch_coalescer.rs @@ -169,7 +169,8 @@ impl DecodedBatchCoalescer { return None; } let rows = batch.num_rows(); - // ArrayData reports logical bytes; bitmap builders round up to 64-byte blocks. + // Bitmap builders round up to 64 bytes. Charging padding per input overcounts, + // but keeps the output-copy bound conservative even when columns share input buffers. let bitmap_bytes = rows.div_ceil(8); let padded_bitmap_bytes = bitmap_bytes.next_multiple_of(64); let mut bytes = BatchBytes::default(); @@ -341,41 +342,6 @@ mod tests { 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() { - 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] @@ -387,9 +353,18 @@ mod tests { 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]); + let output = collect_with_limits( + 32, + max_bytes, + vec![small.clone(), slice, small.clone(), small.clone()], + ); assert_eq!(output.len(), 3); + assert_eq!(output[0], small); assert!(Arc::ptr_eq(output[1].column(0), &original)); + assert_eq!( + output[2], + one_column(Arc::new(StringArray::from(vec!["b", "b"]))) + ); } #[test] @@ -542,34 +517,6 @@ mod tests { assert_eq!(collect(8, batches.clone()), batches); } - #[test] - fn ready_lookahead_and_pending_inputs_have_owned_lifetimes() { - let first = ints(0..1); - let large = ints(1..9); - let first_weak = Arc::downgrade(first.column(0)); - let large_weak = Arc::downgrade(large.column(0)); - let mut coalescer = DecodedBatchCoalescer::try_new(8).unwrap(); - coalescer.push(first).unwrap(); - assert!(coalescer.peek_ready().is_none()); - coalescer.push(large).unwrap(); - assert_eq!(coalescer.peek_ready().unwrap().num_rows(), 1); - let output = coalescer.take_ready().unwrap(); - assert!(first_weak.upgrade().is_some()); - drop(output); - assert!(first_weak.upgrade().is_none()); - assert!(large_weak.upgrade().is_some()); - let lookahead = coalescer.take_ready().unwrap(); - assert_eq!(lookahead, ints(1..9)); - drop(lookahead); - assert!(large_weak.upgrade().is_none()); - let pending = ints(10..11); - let pending_weak = Arc::downgrade(pending.column(0)); - coalescer.push(pending).unwrap(); - assert!(coalescer.peek_ready().is_none()); - drop(coalescer); - assert!(pending_weak.upgrade().is_none()); - } - #[test] fn empty_inputs_and_unit_row_target_pass_through() { let empty = ints(0..0); diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala index e34d7875b2d..40d2500a11a 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala @@ -22,8 +22,6 @@ package org.apache.spark.sql.comet.execution.shuffle import java.io.{FilterInputStream, InputStream} import java.util.concurrent.atomic.AtomicBoolean -import scala.util.control.NonFatal - import org.apache.spark.{InterruptibleIterator, MapOutputTracker, SparkEnv, TaskContext} import org.apache.spark.internal.{config, Logging} import org.apache.spark.io.CompressionCodec @@ -130,27 +128,15 @@ class CometBlockStoreShuffleReader[K, C]( /** 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() - catch { - case NonFatal(failure) => Utils.tryWithSafeFinally(throw failure)(input.close()) - } - val decoder = - try { - NativeBatchDecoderIterator( - input, - dep.decodeTime, - new Native(), - nativeUtil, - CometConf.COMET_TRACING_ENABLED.get(), - readBufferSize = CometConf.COMET_SHUFFLE_READ_BUFFER_SIZE.get().toInt, - coalescing = Some(maxRows)) - } catch { - case NonFatal(failure) => - Utils.tryWithSafeFinally(throw failure) { - Utils.tryWithSafeFinally(input.close())(nativeUtil.close()) - } - } + val nativeUtil = new NativeUtil() + val decoder = NativeBatchDecoderIterator( + input, + dep.decodeTime, + new Native(), + nativeUtil, + CometConf.COMET_TRACING_ENABLED.get(), + readBufferSize = CometConf.COMET_SHUFFLE_READ_BUFFER_SIZE.get().toInt, + coalescing = Some(maxRows)) // Registration invokes the callback immediately if completion raced construction. The // decoder is lazy, so a closed task cannot acquire native batching state after cleanup. context.addTaskCompletionListener[Unit] { _ => diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index 72af2960583..f6414620c15 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -132,30 +132,6 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper } } - test("broadcast producer batching leaves unsupported schemas on the native shuffle reader") { - withSQLConf( - CometConf.COMET_BROADCAST_PRODUCER_COALESCE_ENABLED.key -> "true", - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { - val build = spark - .range(0, 32, 1, 4) - .selectExpr("CAST(id AS INT) AS k", "CAST(id AS DECIMAL(10, 2)) AS v") - .repartition(2, col("k")) - val probe = spark.range(0, 32, 1, 2).selectExpr("CAST(id AS INT) AS k") - val df = probe.join(broadcast(build), "k") - checkAnswer(df, (0 until 32).map(i => Row(i, BigDecimal(i).setScale(2)))) - val plan = df.queryExecution.executedPlan - val exchanges = collect(plan) { case exchange: CometBroadcastExchangeExec => exchange } - assert(exchanges.size == 1, plan.treeString) - assert(exchanges.head.metrics("numOutputRows").value == 32L) - assert( - collect(plan) { - case shuffle: CometShuffleExchangeExec if shuffle.shuffleType == CometNativeShuffle => - shuffle - }.nonEmpty, - plan.treeString) - } - } - /** * Runs `f` with nested hash partitioning keys enabled. The config is disabled by default, so * every test that expects a nested key to reach native shuffle has to opt in. diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala index 5f72b9da4dc..0a8e186e2f6 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala @@ -454,6 +454,14 @@ trait CometBenchmarkBase object CometBenchmarkBase { + private[benchmark] def onOffModes(onFirst: Boolean): Seq[(String, Boolean)] = { + if (onFirst) { + Seq("ON" -> true, "OFF" -> false, "ON (repeat)" -> true) + } else { + Seq("OFF" -> false, "ON" -> true, "OFF (repeat)" -> false) + } + } + /** * SplitMix64 finalizer, used to turn a row id into a well distributed pseudo-random value. * diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastLocalShuffleBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastLocalShuffleBenchmark.scala index 796819b3267..295e76e69d6 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastLocalShuffleBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastLocalShuffleBenchmark.scala @@ -29,14 +29,8 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf /** - * A complete local query with a native shuffle on the broadcast build side. Explicit repartition - * and broadcast keep the plan stable; this measures neither TPC-H nor AQE's SMJ-to-BHJ - * conversion. Every iteration generates both inputs, writes and reads the shuffle, builds and - * distributes the broadcast, joins, and aggregates the result. No stage or input is cached - * between iterations. - * - * Before timing, check results, native join/shuffle execution, producer batching activation, and - * the number of Arrow streams reaching driver assembly. Pass --on-first to reverse case order. + * Measures a complete local native-shuffle to broadcast-join query with uncached inputs and a + * final aggregation. Pass --on-first to reverse case order. * * To run with four local workers: * {{{ @@ -100,11 +94,7 @@ object CometBroadcastLocalShuffleBenchmark extends CometBenchmarkBase { } assert(streams(1) > 0 && streams(1) < streams(0), streams) - val modes = if (mainArgs.contains("--on-first")) { - Seq("ON" -> true, "OFF" -> false, "ON (repeat)" -> true) - } else { - Seq("OFF" -> false, "ON" -> true, "OFF (repeat)" -> false) - } + val modes = CometBenchmarkBase.onOffModes(mainArgs.contains("--on-first")) runBenchmark("local native shuffle to broadcast join") { val benchmark = new Benchmark( "local native shuffle to broadcast join", diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastProducerBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastProducerBenchmark.scala index 5b36f093d07..4e50045dab7 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastProducerBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastProducerBenchmark.scala @@ -20,15 +20,12 @@ 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.execution.shuffle.{NativeBatchDecoderIterator, NativeShuffleTestFrames} import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.metric.SQLMetric @@ -36,10 +33,8 @@ 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. + * Measures native shuffle decoding, Arrow export and broadcast serialization, with input frame + * preparation and output validation outside timing. * * 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. @@ -90,22 +85,7 @@ object CometBroadcastProducerBenchmark extends CometBenchmarkBase { 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 - .allocate(20) - .order(ByteOrder.LITTLE_ENDIAN) - .putLong(12L + bytes.length.toLong) - .putLong(root.getFieldVectors.size().toLong) - .put(Array[Byte](78, 79, 78, 69)) // NONE codec, followed by Arrow IPC. - output.write(header.array()) - output.write(bytes) + output.write(NativeShuffleTestFrames.arrowFrame(root)) start += count } output.toByteArray @@ -116,19 +96,24 @@ object CometBroadcastProducerBenchmark extends CometBenchmarkBase { } } + private def newDecoder( + input: Array[Byte], + enabled: Boolean, + util: NativeUtil): NativeBatchDecoderIterator = NativeBatchDecoderIterator( + new ByteArrayInputStream(input), + new SQLMetric("nsTiming", 0L), + new Native, + util, + tracingEnabled = false, + coalescing = if (enabled) Some(targetRows) else None) + private def produce( input: Array[Byte], scenario: Scenario, enabled: Boolean, verify: Boolean): Result = { val util = new NativeUtil - val decoder = NativeBatchDecoderIterator( - new ByteArrayInputStream(input), - new SQLMetric("nsTiming", 0L), - new Native, - util, - tracingEnabled = false, - coalescing = if (enabled) Some(targetRows) else None) + val decoder = newDecoder(input, enabled, util) var rows = 0L var streams = 0L var bytes = 0L @@ -171,13 +156,7 @@ object CometBroadcastProducerBenchmark extends CometBenchmarkBase { private def decodeOnly(input: Array[Byte], scenario: Scenario, enabled: Boolean): Unit = { val util = new NativeUtil - val decoder = NativeBatchDecoderIterator( - new ByteArrayInputStream(input), - new SQLMetric("nsTiming", 0L), - new Native, - util, - tracingEnabled = false, - coalescing = if (enabled) Some(targetRows) else None) + val decoder = newDecoder(input, enabled, util) var rows = 0L try { while (decoder.hasNext) rows += decoder.next().numRows() @@ -192,11 +171,7 @@ object CometBroadcastProducerBenchmark extends CometBenchmarkBase { // Initialize SparkEnv, whose compression codec serializeBatches uses. spark.sparkContext.setLogLevel("WARN") val scenarioFilters = mainArgs.filterNot(Set("--on-first", "--decode-only")) - val modes = if (mainArgs.contains("--on-first")) { - Seq("ON" -> true, "OFF" -> false, "ON (repeat)" -> true) - } else { - Seq("OFF" -> false, "ON" -> true, "OFF (repeat)" -> false) - } + val modes = CometBenchmarkBase.onOffModes(mainArgs.contains("--on-first")) scenarios.filter(s => scenarioFilters.forall(s.name.contains)).foreach { scenario => val input = frames(scenario) runBenchmark(scenario.name) { diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala index 9ca690f3b69..d315bddfe02 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala @@ -59,16 +59,6 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { import NativeShuffleTestFrames._ import testImplicits._ - private final class TrackingInputStream(bytes: Array[Byte]) - extends ByteArrayInputStream(bytes) { - var closeCalls = 0 - - override def close(): Unit = { - closeCalls += 1 - super.close() - } - } - private final class RecordingReaderBackend(client: RecordingCelebornRawClient) extends ShuffleManager { var readRange: Option[(Int, Int, Int, Int)] = None diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala index d959786c487..871255b1ecc 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala @@ -24,7 +24,6 @@ import org.apache.spark.rdd.{DeterministicLevel, RDD} import org.apache.spark.serializer.JavaSerializer import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.comet.{CometMetricNode, NativeExecContext} -import org.apache.spark.sql.execution.{CoalescedMapperPartitionSpec, CoalescedPartitionSpec, PartialMapperPartitionSpec, PartialReducerPartitionSpec, ShufflePartitionSpec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch @@ -46,29 +45,6 @@ import org.apache.comet.serde.OperatorOuterClass.Operator */ class CometNativeShuffleInputRDDSuite extends CometTestBase { - test("broadcast producer batching preserves shared shuffle dependencies and AQE ranges") { - val input = spark.sparkContext.parallelize(Seq.empty[(Int, ColumnarBatch)], 4) - val dependency = new CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]( - input, - new HashPartitioner(4), - decodeTime = SQLMetrics.createMetric(spark.sparkContext, "decode time")) - val specs = Array[ShufflePartitionSpec]( - CoalescedPartitionSpec(1, 3), - PartialReducerPartitionSpec(2, 1, 3, 100L), - PartialMapperPartitionSpec(2, 1, 3), - CoalescedMapperPartitionSpec(1, 3, 4)) - val original = new CometShuffledBatchRDD(dependency, Map.empty, specs) - val batched = original.withNativeBatchCoalescing(4096) - assert(original.supportsNativeBatchCoalescing) - assert(batched.dependency eq original.dependency) - assert(batched.dependencies == original.dependencies) - assert(original.nativeBatchCoalescing.isEmpty) - assert(batched.nativeBatchCoalescing.contains(4096)) - assert( - batched.partitions.map(_.asInstanceOf[ShuffledRowRDDPartition].spec).sameElements(specs)) - assert(batched.partitioner == original.partitioner) - } - test("native shuffle input preserves its parents' determinism") { Seq( DeterministicLevel.DETERMINATE, diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorSuite.scala index 14c52c788da..37c5ceceb1c 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorSuite.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.comet.execution.shuffle -import java.io.{ByteArrayInputStream, EOFException} +import java.io.EOFException import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.execution.metric.SQLMetrics @@ -29,17 +29,7 @@ import org.apache.comet.vector.NativeUtil class NativeBatchDecoderIteratorSuite extends CometTestBase { - import NativeShuffleTestFrames.intFrame - - private final class TrackingInputStream(bytes: Array[Byte]) - extends ByteArrayInputStream(bytes) { - var closeCalls = 0 - - override def close(): Unit = { - closeCalls += 1 - super.close() - } - } + import NativeShuffleTestFrames.{intFrame, TrackingInputStream} test("local decoder groups frames, drains passthrough output, and flushes the final batch") { val scenarios = Seq( diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeShuffleTestFrames.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeShuffleTestFrames.scala index 139f5e11609..8e6ff3fd603 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeShuffleTestFrames.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeShuffleTestFrames.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.comet.execution.shuffle -import java.io.ByteArrayOutputStream +import java.io.{ByteArrayInputStream, ByteArrayOutputStream} import java.nio.{ByteBuffer, ByteOrder} import java.nio.channels.Channels @@ -29,7 +29,16 @@ import org.apache.arrow.vector.dictionary.DictionaryProvider import org.apache.arrow.vector.ipc.ArrowStreamWriter /** Local and remote readers consume the same framed Arrow IPC format. */ -private[shuffle] object NativeShuffleTestFrames { +private[sql] object NativeShuffleTestFrames { + + final class TrackingInputStream(bytes: Array[Byte]) extends ByteArrayInputStream(bytes) { + var closeCalls = 0 + + override def close(): Unit = { + closeCalls += 1 + super.close() + } + } def frameBytes(payload: Byte*): Array[Byte] = ByteBuffer .allocate(20 + payload.size)