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 b73f81e8968..4e7248f2ba3 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. 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 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..629ce8f02c3 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -214,6 +214,14 @@ 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 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: 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..4c86a2dcb89 100644 --- a/docs/source/user-guide/latest/tuning/shuffle.md +++ b/docs/source/user-guide/latest/tuning/shuffle.md @@ -53,6 +53,20 @@ 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. + +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 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..50df7a786b0 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,137 @@ 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.max(0) as usize)?; + Ok(Box::into_raw(Box::new(coalescer)) as jlong) + }) +} + +#[no_mangle] +/// Decode and retain one local frame, returning ready columns or -1 for more input. +/// # Safety +/// The handle must be live and exclusively owned by this reader. On error, close the reader. +pub unsafe extern "system" fn Java_org_apache_comet_Native_appendShuffleBatch( + e: EnvUnowned, + _class: JClass, + handle: jlong, + byte_buffer: JByteBuffer, + length: jint, + num_columns: jint, + tracing_enabled: jboolean, +) -> jint { + try_unwrap_or_throw(&e, |env| { + with_trace("appendShuffleBatch", tracing_enabled != JNI_FALSE, || { + if num_columns < 0 { + return Err(CometError::Internal("Invalid shuffle frame bounds".into())); + } + 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)) + }) + }) +} + +#[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_peekShuffleBatch( + 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. +/// 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, + 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()))?; + 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. /// @@ -2251,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( @@ -2261,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) } @@ -3398,6 +3544,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..91d902a4892 --- /dev/null +++ b/native/shuffle/src/decoded_batch_coalescer.rs @@ -0,0 +1,526 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Group small decoded shuffle batches before exporting them to the JVM. +//! +//! Groups are bounded by rows, retained Arrow bytes, and input batch count. An +//! oversized input passes through unchanged instead of being split or copied. +//! +//! Keep inputs until a group is ready, then concatenate once at its actual size. + +use std::collections::VecDeque; +use std::ptr::NonNull; + +use arrow::array::RecordBatch; +use arrow::buffer::Buffer; +use arrow::compute::concat_batches; +use arrow::datatypes::DataType; +use datafusion::common::{DataFusionError, Result}; + +const MAX_ROWS: usize = 65_536; +const MAX_GROUP_BATCHES: usize = 1024; +const MAX_GROUP_BYTES: usize = 8 * 1024 * 1024; + +/// Charge the shared IPC body once, but count each column's output copy separately. +/// Other shared allocations may be conservatively overcounted. +#[derive(Default)] +struct BatchBytes { + first_owner: Option>, + retained: usize, + copied: usize, +} + +impl BatchBytes { + fn add_buffer(&mut self, buffer: &Buffer) { + let owner = buffer.data_ptr(); + if self.first_owner == Some(owner) { + return; + } + let capacity = buffer.capacity().max(buffer.len()); + if capacity != 0 { + if self.first_owner.is_none() { + self.first_owner = Some(owner); + } + self.retained = self.retained.saturating_add(capacity); + } + } +} + +/// Ordered grouping of batches returned by the native shuffle IPC decoder. +/// +/// `push` and `finish` enqueue output. Inspect it with `peek_ready`, and drain +/// it with `take_ready` before pushing another input. Call `finish` at EOF. +/// A singleton, an oversized input, or an input at least half the row target is +/// never copied. The byte budget bounds retained group inputs, not the whole +/// reader: concatenation also owns its output and may have one decoded lookahead +/// input. Callers must drop the helper on any error. +#[derive(Debug)] +pub struct DecodedBatchCoalescer { + max_rows: usize, + max_bytes: usize, + batches: Vec, + rows: usize, + bytes: usize, + ready: VecDeque, +} + +impl DecodedBatchCoalescer { + /// Creates an empty, reader-owned grouper with the production byte budget. + /// Caps positive row targets at 65536. + pub fn try_new(max_rows: usize) -> Result { + Self::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 max_rows == 0 { + return Err(DataFusionError::Configuration( + "Decoded shuffle grouping requires a positive row target".into(), + )); + } + if !(1..=MAX_GROUP_BYTES).contains(&max_bytes) { + return Err(DataFusionError::Configuration( + "Decoded shuffle grouping requires 1..8388608 bytes".into(), + )); + } + Ok(Self { + max_rows: max_rows.min(MAX_ROWS), + max_bytes, + batches: Vec::new(), + rows: 0, + bytes: 0, + ready: VecDeque::new(), + }) + } + + /// Transfers the oldest ready output to the caller without copying buffers. + /// Returns None when empty and leaves any pending group unchanged. + pub fn take_ready(&mut self) -> Option { + self.ready.pop_front() + } + + /// Borrows the oldest ready output until the next mutation, without draining it. + pub fn peek_ready(&self) -> Option<&RecordBatch> { + self.ready.front() + } + + /// Takes ownership of one decoded input. A flush followed by a passthrough can make two + /// outputs ready; drain both before reading another input. Drop this helper on error. + pub fn push(&mut self, batch: RecordBatch) -> Result<()> { + if !self.ready.is_empty() { + return Err(DataFusionError::Execution( + "Drain shuffle grouping output before reading another input".into(), + )); + } + let rows = batch.num_rows(); + let bytes = self.eligible_bytes(&batch); + + if !self.batches.is_empty() + && (bytes.is_none_or(|bytes| bytes > self.max_bytes - self.bytes) + || self.batches[0].schema() != batch.schema() + || rows > self.max_rows - self.rows) + { + self.flush()?; + } + if let Some(bytes) = bytes { + self.rows += rows; + self.bytes += bytes; + self.batches.push(batch); + if self.rows == self.max_rows + || self.bytes == self.max_bytes + || self.batches.len() == MAX_GROUP_BATCHES + { + self.flush()?; + } + } else { + self.ready.push_back(batch); + } + Ok(()) + } + + /// Flushes the final group. Repeated calls do not copy or emit it again. + pub fn finish(&mut self) -> Result<()> { + self.flush() + } + + /// Bounds both retained backing allocations and output copies for flat columns. + /// None means pass through without copying. + fn eligible_bytes(&self, batch: &RecordBatch) -> Option { + 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(); + // 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(); + for column in batch.columns() { + if !matches!( + column.data_type(), + DataType::Int32 + | DataType::Int64 + | DataType::Boolean + | DataType::Utf8 + | DataType::Binary + ) { + return None; + } + let data = column.to_data(); + let metadata = std::mem::size_of_val(column.as_ref()); + // to_data may omit an all-valid bitmap; account for the original nulls below. + let copied = + data.get_slice_memory_size().ok()? - data.nulls().map_or(0, |_| bitmap_bytes); + bytes.copied = bytes.copied.saturating_add(copied).saturating_add(metadata); + bytes.retained = bytes.retained.saturating_add(metadata); + for buffer in data.buffers() { + bytes.add_buffer(buffer); + } + if let Some(nulls) = column.nulls() { + bytes.add_buffer(nulls.buffer()); + bytes.copied = bytes.copied.saturating_add(padded_bitmap_bytes); + } + if data.data_type() == &DataType::Boolean { + bytes.copied = bytes + .copied + .saturating_add(padded_bitmap_bytes - bitmap_bytes); + } + if bytes.retained.max(bytes.copied) > self.max_bytes { + return None; + } + } + Some(bytes.retained.max(bytes.copied)) + } + + fn flush(&mut self) -> Result<()> { + if self.batches.len() < 2 { + self.ready.extend(self.batches.drain(..)); + } else { + self.ready + .push_back(concat_batches(&self.batches[0].schema(), &self.batches)?); + self.batches.clear(); + } + self.rows = 0; + self.bytes = 0; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + use arrow::array::{ + ArrayRef, AsArray, BinaryArray, BooleanArray, Int32Array, Int64Array, StringArray, + }; + use arrow::datatypes::{Field, Int32Type, Schema}; + use arrow::ipc::writer::StreamWriter; + + fn one_column(array: ArrayRef) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new( + "x", + array.data_type().clone(), + true, + )])); + RecordBatch::try_new(schema, vec![array]).unwrap() + } + + fn ints(values: impl IntoIterator) -> 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(); + while let Some(batch) = coalescer.take_ready() { + output.push(batch); + } + output + } + + #[test] + fn finishing_repeatedly_emits_the_pending_group_once() { + let mut coalescer = DecodedBatchCoalescer::try_new(8).unwrap(); + coalescer.push(ints(0..1)).unwrap(); + coalescer.push(ints(1..2)).unwrap(); + coalescer.finish().unwrap(); + coalescer.finish().unwrap(); + assert_eq!(coalescer.take_ready(), Some(ints(0..2))); + coalescer.finish().unwrap(); + assert!(coalescer.take_ready().is_none()); + } + + #[test] + fn singleton_large_and_near_full_inputs_are_not_copied() { + for batches in [ + vec![ints(0..1)], + vec![ints(0..1), ints(1..7)], + vec![ints(0..6), ints(6..7)], + vec![ints(0..1), ints(1..10)], + ] { + let originals: Vec<_> = batches + .iter() + .map(|batch| batch.column(0).clone()) + .collect(); + let output = collect(8, batches); + assert_eq!(output.len(), originals.len()); + for (batch, original) in output.iter().zip(originals) { + assert!(Arc::ptr_eq(batch.column(0), &original)); + } + } + } + + #[test] + fn row_and_input_count_limits_preserve_order() { + let batches = (0..4).map(|i| ints(i * 2..i * 2 + 2)).collect(); + assert_eq!(collect(8, batches), vec![ints(0..8)]); + assert_eq!( + collect(8, vec![ints(0..3), ints(3..6), ints(6..9)]), + vec![ints(0..6), ints(6..9)] + ); + let batches = (0..1025).map(|i| ints(i..i + 1)).collect(); + assert_eq!( + collect(MAX_ROWS, batches), + vec![ints(0..1024), ints(1024..1025)] + ); + } + + #[test] + fn byte_limit_flushes_at_and_before_crossing_the_boundary() { + let inputs: Vec<_> = (0..3).map(|i| ints(i..i + 1)).collect(); + let two_batches = inputs[0].get_array_memory_size() + inputs[1].get_array_memory_size(); + assert_eq!( + collect_with_limits(32, two_batches, inputs.clone()), + vec![ints(0..2), ints(2..3)] + ); + let output = collect_with_limits(32, two_batches - 1, inputs.clone()); + assert_eq!(output, inputs); + for (output, input) in output.iter().zip(&inputs) { + assert!(Arc::ptr_eq(output.column(0), input.column(0))); + } + } + + #[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.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] + 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( + (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 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!(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 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_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![ + ( + 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 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)]); + } +} 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..a89cdfb5b2b 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -646,6 +646,17 @@ 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. 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) + 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..c24bd762c3e 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 peekShuffleBatch(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..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 @@ -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 @@ -110,9 +112,25 @@ 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() + // 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 + case _ => false + } + } + val readInput = input match { + case shuffle: CometShuffledBatchRDD + if CometConf.COMET_BROADCAST_PRODUCER_COALESCE_ENABLED.get(conf) && + supported && shuffle.supportsNativeBatchCoalescing => + shuffle.withNativeBatchCoalescing(CometConf.COMET_BATCH_SIZE.get(conf)) + 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..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 @@ -34,6 +34,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 +122,33 @@ 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 = 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] { _ => + 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..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 @@ -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,46 @@ 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 { + // 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 { + case NonFatal(cleanupFailure) => + if (cleanupFailure ne failure) failure.addSuppressed(cleanupFailure) + } + throw failure + } } def next(): ColumnarBatch = { @@ -163,6 +179,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.peekShuffleBatch(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.peekShuffleBatch(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 +299,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 +315,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..f6414620c15 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 @@ -41,7 +42,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 +74,64 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper import testImplicits._ + test("broadcast producer batching combines local native shuffle blocks with and without AQE") { + withSQLConf( + CometConf.COMET_BATCH_SIZE.key -> "4096", + SQLConf.SHUFFLE_PARTITIONS.key -> "4", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_SIZE.key -> "1") { + for (adaptive <- Seq(false, true)) { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + val streams = scala.collection.mutable.ArrayBuffer.empty[Long] + Seq(false, true).foreach { enabled => + withSQLConf( + CometConf.COMET_BROADCAST_PRODUCER_COALESCE_ENABLED.key -> enabled.toString) { + val build = spark + .range(0, 512, 1, 16) + .selectExpr( + "CAST(id AS INT) AS k", + "CASE WHEN id % 3 = 0 THEN NULL ELSE concat('v', id) END AS v", + "CASE WHEN id % 5 = 0 THEN NULL ELSE id + 4294967296 END AS l", + "CASE WHEN id % 7 = 0 THEN NULL ELSE id % 2 = 0 END AS b", + "CASE WHEN id % 11 = 0 THEN NULL ELSE CAST(concat('b', id) AS BINARY) END AS bytes") + .repartition(col("k")) + val probe = spark.range(0, 512, 1, 2).selectExpr("CAST(id AS INT) AS k") + val df = probe.join(broadcast(build), "k") + checkAnswer( + df, + (0 until 512).map { i => + Row( + i, + if (i % 3 == 0) null else s"v$i", + if (i % 5 == 0) null else i.toLong + 4294967296L, + if (i % 7 == 0) null else i % 2 == 0, + if (i % 11 == 0) null else s"b$i".getBytes(UTF_8)) + }) + val plan = df.queryExecution.executedPlan + val exchanges = collect(plan) { case exchange: CometBroadcastExchangeExec => + exchange + } + assert(exchanges.size == 1, plan.treeString) + val exchange = exchanges.head + assert(exchange.metrics("numOutputRows").value == 512L) + assert( + collect(plan) { + case shuffle: CometShuffleExchangeExec + if shuffle.shuffleType == CometNativeShuffle => + shuffle + }.nonEmpty, + plan.treeString) + streams += exchange.metrics("numCoalescedBatches").value + () + } + } + assert(streams(1) > 0L && streams(1) < streams(0), s"AQE=$adaptive, batches=$streams") + } + } + } + } + /** * 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 new file mode 100644 index 00000000000..295e76e69d6 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastLocalShuffleBenchmark.scala @@ -0,0 +1,116 @@ +/* + * 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 + +/** + * 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: + * {{{ + * 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 = CometBenchmarkBase.onOffModes(mainArgs.contains("--on-first")) + 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 new file mode 100644 index 00000000000..4e50045dab7 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastProducerBenchmark.scala @@ -0,0 +1,204 @@ +/* + * 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.charset.StandardCharsets + +import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.{BigIntVector, IntVector, VarCharVector, VectorSchemaRoot} +import org.apache.spark.benchmark.Benchmark +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 + +import org.apache.comet.Native +import org.apache.comet.vector.NativeUtil + +/** + * 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. + * + * 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) + output.write(NativeShuffleTestFrames.arrowFrame(root)) + start += count + } + output.toByteArray + } finally { + root.close() + if (!scenario.wide) strings.close() + allocator.close() + } + } + + 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 = newDecoder(input, enabled, util) + 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() + } + } + + private def decodeOnly(input: Array[Byte], scenario: Scenario, enabled: Boolean): Unit = { + val util = new NativeUtil + val decoder = newDecoder(input, enabled, util) + 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(Set("--on-first", "--decode-only")) + 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) { + 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)) + } + } + 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 0e618b96088..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 @@ -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,27 +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 - - override def close(): Unit = { - closeCalls += 1 - super.close() - } - } - private final class RecordingReaderBackend(client: RecordingCelebornRawClient) extends ShuffleManager { var readRange: Option[(Int, Int, Int, Int)] = None @@ -1227,42 +1209,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() - } - } - 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..37c5ceceb1c --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorSuite.scala @@ -0,0 +1,112 @@ +/* + * 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.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, TrackingInputStream} + + 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..8e6ff3fd603 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeShuffleTestFrames.scala @@ -0,0 +1,85 @@ +/* + * 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, 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[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) + .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() + } + } +}