sunchao commented on code in PR #7258: URL: https://github.com/apache/datafusion-comet/pull/7258#discussion_r4394913989
########## native/shuffle/src/decoded_batch_coalescer.rs: ########## @@ -0,0 +1,582 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Group small decoded shuffle batches before exporting them to the JVM. +//! +//! Groups are bounded by rows, retained Arrow bytes, and input batch count. An +//! oversized input passes through unchanged instead of being split or copied. +//! +//! Keep inputs until a group is ready, then concatenate once at its actual size. + +use std::collections::VecDeque; +use std::ptr::NonNull; + +use arrow::array::{Array, AsArray, RecordBatch}; +use arrow::buffer::Buffer; +use arrow::compute::concat_batches; +use arrow::datatypes::{DataType, Int32Type, Int64Type}; +use datafusion::common::{DataFusionError, Result}; + +const MAX_ROWS: usize = 65_536; +const MAX_GROUP_BATCHES: usize = 1024; +const MAX_GROUP_BYTES: usize = 8 * 1024 * 1024; + +/// Allocation-free accounting for one decoded batch, discarded after admission. +/// IPC columns normally share their message-body allocation. Remember that first +/// owner so its slices are charged once; other owners may be overcounted safely. +/// Logical copy bytes are tracked separately because columns may alias values +/// that concatenation will materialize independently. +#[derive(Default)] +struct BatchBytes { + first_owner: Option<NonNull<u8>>, + retained: usize, + copied: usize, +} + +impl BatchBytes { + /// Charges a borrowed buffer's allocation and the bytes its output will copy. + /// Uses the allocation base, not the slice address, to deduplicate the common + /// IPC owner before inspecting its capacity again. Saturating counters + /// conservatively reject overflow. No ownership is retained and no memory + /// is allocated or read from the buffer contents. + fn add_buffer(&mut self, buffer: &Buffer, copied: usize) { + self.copied = self.copied.saturating_add(copied); + let owner = buffer.data_ptr(); + if self.first_owner == Some(owner) { + return; + } + let capacity = buffer.capacity().max(buffer.len()); + if capacity != 0 { + if self.first_owner.is_none() { + self.first_owner = Some(owner); + } + self.retained = self.retained.saturating_add(capacity); + } + } + + /// Charges a borrowed concrete array's null buffer and object metadata after + /// its values. The caller supplies the rounded output bitmap size. Static + /// dispatch avoids repeating trait-object lookups after the array downcast; + /// ownership, allocation-free accounting, and saturation match `add_buffer`. + #[inline] + fn add_nulls_and_metadata<T: Array>(&mut self, array: &T, copied_null_bytes: usize) { + if let Some(nulls) = array.nulls() { + self.add_buffer(nulls.buffer(), copied_null_bytes); + } + let array_metadata = std::mem::size_of::<T>(); + self.retained = self.retained.saturating_add(array_metadata); + self.copied = self.copied.saturating_add(array_metadata); + } +} + +/// Ordered grouping of batches returned by the native shuffle IPC decoder. +/// +/// `push` and `finish` enqueue output. Inspect it with `peek_ready`, and drain +/// it with `take_ready` before pushing another input. Call `finish` at EOF. +/// A singleton, an oversized input, or an input at least half the row target is +/// never copied. The byte budget bounds retained group inputs, not the whole +/// reader: concatenation also owns its output and may have one decoded lookahead +/// input. Callers must drop the helper on any error. +#[derive(Debug)] +pub struct DecodedBatchCoalescer { + max_rows: usize, + max_bytes: usize, + batches: Vec<RecordBatch>, + rows: usize, + bytes: usize, + ready: VecDeque<RecordBatch>, +} + +impl DecodedBatchCoalescer { + /// Creates an empty, reader-owned grouper with the production byte budget. + /// Rejects row targets outside 1..=65536 without retaining any buffers. + pub fn try_new(max_rows: usize) -> Result<Self> { + Self::try_new_with_limits(max_rows, MAX_GROUP_BYTES) + } + + /// Creates an empty grouper with explicit limits, allowing small test budgets. + /// Invalid limits return a configuration error. Keeping the byte budget at + /// most 8 MiB also keeps every concatenated STRING/BINARY offset within i32. + fn try_new_with_limits(max_rows: usize, max_bytes: usize) -> Result<Self> { + if !(1..=MAX_ROWS).contains(&max_rows) { + return Err(DataFusionError::Configuration( + "Decoded shuffle grouping requires 1..65536 rows".into(), + )); + } + if !(1..=MAX_GROUP_BYTES).contains(&max_bytes) { + return Err(DataFusionError::Configuration( + "Decoded shuffle grouping requires 1..8388608 bytes".into(), + )); + } + Ok(Self { + max_rows, + max_bytes, + batches: Vec::new(), + rows: 0, + bytes: 0, + ready: VecDeque::new(), + }) + } + + /// Transfers the oldest ready output to the caller without copying buffers. + /// Returns None when empty and leaves any pending group unchanged. + pub fn take_ready(&mut self) -> Option<RecordBatch> { + self.ready.pop_front() + } + + /// Borrows the oldest ready output until the next mutation, without draining it. + pub fn peek_ready(&self) -> Option<&RecordBatch> { + self.ready.front() + } + + /// Takes ownership of one decoded input. A flush followed by a passthrough can make two + /// outputs ready; drain both before reading another input. Drop this helper on error. + pub fn push(&mut self, batch: RecordBatch) -> Result<()> { + if !self.ready.is_empty() { + return Err(DataFusionError::Execution( + "Drain shuffle grouping output before reading another input".into(), + )); + } + let rows = batch.num_rows(); + let bytes = self.eligible_bytes(&batch); + + if !self.batches.is_empty() + && (bytes.is_none_or(|bytes| bytes > self.max_bytes - self.bytes) + || self.batches[0].schema() != batch.schema() + || rows > self.max_rows - self.rows) + { + self.flush()?; + } + if let Some(bytes) = bytes { + self.rows += rows; + self.bytes += bytes; + self.batches.push(batch); + if self.rows == self.max_rows + || self.bytes == self.max_bytes + || self.batches.len() == MAX_GROUP_BATCHES + { + self.flush()?; + } + } else { + self.ready.push_back(batch); + } + Ok(()) + } + + /// Flushes the final group. Repeated calls do not copy or emit it again. + pub fn finish(&mut self) -> Result<()> { + self.flush() + } + + /// Returns the larger of retained and logical copy bytes for a groupable input. + /// Examines flat-column buffer metadata and string/binary endpoint offsets, + /// never rows. The common shared IPC allocation is charged once. Other shared + /// owners can be overcounted safely, and slices retain their full backing + /// allocation charge. RecordBatch guarantees matching field and array types, + /// so dispatch borrows schema fields without cloning or inspecting array + /// trait objects twice. None means pass through without copying. + fn eligible_bytes(&self, batch: &RecordBatch) -> Option<usize> { Review Comment: Replaced the per-type byte calculations with `ArrayData::get_slice_memory_size()` and buffer traversal in `b97daacb45`. Retained backing allocations are still charged separately from each column's copied bytes, including the original null bitmap when `to_data()` drops an all-valid mask. The existing five-type scope is unchanged; this removes the custom calculations without adding new input layouts to this PR. I also retained the backing/alias regression tests. ########## 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 { Review Comment: Trimmed these comments in `b97daacb45` to the accounting and ownership invariants: charge the shared IPC body once, count each output column's copy, and allow conservative overcounting of other shared buffers. The static-dispatch/allocation-free optimization commentary is gone. ########## spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala: ########## @@ -110,9 +115,24 @@ case class CometBroadcastExchangeExec( private lazy val maxBroadcastRows = 512000000L private def getByteArrayRdd(plan: SparkPlan): RDD[(Long, ChunkedByteBuffer)] = { - plan.executeColumnar().mapPartitionsInternal { iter => - Utils.serializeBatches(iter) + // Capture the read mode on the driver. A new RDD keeps it local to this broadcast consumer, + // including when AQE or exchange reuse shares the original shuffle with another operator. + val input = plan.executeColumnar() + val supported = plan.output.nonEmpty && plan.output.forall { attribute => Review Comment: I kept the JVM type check as a cheap routing gate and added a short explanation in `b97daacb45`. Unsupported schemas would otherwise allocate a coalescer and use the peek/append/take JNI path for every frame, even though native would pass all of them through. Native still checks physical layouts, so a dictionary-encoded STRING may correctly pass the logical gate and then bypass grouping. Removing the gate would broaden the changed path without evidence of a benefit; this change keeps the already-tested schema scope. ########## spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala: ########## @@ -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( Review Comment: Removed `nativeProducerBatching` in `b97daacb45`. Integration tests now use the existing serialized-stream count comparison and native-shuffle plan checks, rather than adding a UI counter for mode selection. The new local-query benchmark uses the same evidence. Also, this gate requires a local `CometNativeShuffleHandle`; Celeborn selects a different handle during registration, so there is no local-to-Celeborn takeover after this gate. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
