comphead commented on code in PR #7258: URL: https://github.com/apache/datafusion-comet/pull/7258#discussion_r4393302367
########## native/shuffle/src/decoded_batch_coalescer.rs: ########## @@ -0,0 +1,582 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Group small decoded shuffle batches before exporting them to the JVM. +//! +//! Groups are bounded by rows, retained Arrow bytes, and input batch count. An +//! oversized input passes through unchanged instead of being split or copied. +//! +//! Keep inputs until a group is ready, then concatenate once at its actual size. + +use std::collections::VecDeque; +use std::ptr::NonNull; + +use arrow::array::{Array, AsArray, RecordBatch}; +use arrow::buffer::Buffer; +use arrow::compute::concat_batches; +use arrow::datatypes::{DataType, Int32Type, Int64Type}; +use datafusion::common::{DataFusionError, Result}; + +const MAX_ROWS: usize = 65_536; +const MAX_GROUP_BATCHES: usize = 1024; +const MAX_GROUP_BYTES: usize = 8 * 1024 * 1024; + +/// Allocation-free accounting for one decoded batch, discarded after admission. +/// IPC columns normally share their message-body allocation. Remember that first +/// owner so its slices are charged once; other owners may be overcounted safely. +/// Logical copy bytes are tracked separately because columns may alias values +/// that concatenation will materialize independently. +#[derive(Default)] +struct BatchBytes { + first_owner: Option<NonNull<u8>>, + retained: usize, + copied: usize, +} + +impl BatchBytes { + /// Charges a borrowed buffer's allocation and the bytes its output will copy. + /// Uses the allocation base, not the slice address, to deduplicate the common + /// IPC owner before inspecting its capacity again. Saturating counters + /// conservatively reject overflow. No ownership is retained and no memory + /// is allocated or read from the buffer contents. + fn add_buffer(&mut self, buffer: &Buffer, copied: usize) { + self.copied = self.copied.saturating_add(copied); + let owner = buffer.data_ptr(); + if self.first_owner == Some(owner) { + return; + } + let capacity = buffer.capacity().max(buffer.len()); + if capacity != 0 { + if self.first_owner.is_none() { + self.first_owner = Some(owner); + } + self.retained = self.retained.saturating_add(capacity); + } + } + + /// Charges a borrowed concrete array's null buffer and object metadata after + /// its values. The caller supplies the rounded output bitmap size. Static + /// dispatch avoids repeating trait-object lookups after the array downcast; + /// ownership, allocation-free accounting, and saturation match `add_buffer`. + #[inline] + fn add_nulls_and_metadata<T: Array>(&mut self, array: &T, copied_null_bytes: usize) { + if let Some(nulls) = array.nulls() { + self.add_buffer(nulls.buffer(), copied_null_bytes); + } + let array_metadata = std::mem::size_of::<T>(); + self.retained = self.retained.saturating_add(array_metadata); + self.copied = self.copied.saturating_add(array_metadata); + } +} + +/// Ordered grouping of batches returned by the native shuffle IPC decoder. +/// +/// `push` and `finish` enqueue output. Inspect it with `peek_ready`, and drain +/// it with `take_ready` before pushing another input. Call `finish` at EOF. +/// A singleton, an oversized input, or an input at least half the row target is +/// never copied. The byte budget bounds retained group inputs, not the whole +/// reader: concatenation also owns its output and may have one decoded lookahead +/// input. Callers must drop the helper on any error. +#[derive(Debug)] +pub struct DecodedBatchCoalescer { Review Comment: `BufBatchWriter` in this crate already uses Arrow's `BatchCoalescer` with `with_biggest_coalesce_batch_size` for the same job of merging small batches and passing large ones through. Could this reuse it? `push_batch`, `finish_buffered_batch` and `next_completed_batch` would replace `push`, `finish` and `take_ready`. It would likely also remove `BatchBytes`, `eligible_bytes`, the five-type gate, the 1,024-input cap and the tests that pin that accounting (shared IPC body, aliased columns, sliced backing buffers). I expect `BatchCoalescer` to copy each input into its in-progress output as it arrives, so inputs would not stay alive until a group is ready. That is the cost the retained-bytes accounting works around. I have not measured it. Two differences to check. The bypass fires only for batches strictly larger than the limit, so the half-row rule would be `max_rows.div_ceil(2) - 1`. A large input that arrives after a small buffered group would be merged instead of passed through, which also removes the flush-plus-passthrough state with two ready outputs. If the 8 MiB bound is still wanted for wide STRING and BINARY rows, is there a way to express it on top of `BatchCoalescer`? ########## spark/src/main/scala/org/apache/comet/Native.scala: ########## @@ -289,6 +289,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 Review Comment: `drainShuffleBatch` flushes at EOF and reports how many columns the next ready batch has. It does not remove anything, and `takeShuffleBatch` is the call that does. A name such as `pollShuffleBatch` (matching `peek_ready` on the Rust side) would describe it better and keep `drain` for the call that empties the queue. ########## native/shuffle/src/decoded_batch_coalescer.rs: ########## @@ -0,0 +1,582 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Group small decoded shuffle batches before exporting them to the JVM. +//! +//! Groups are bounded by rows, retained Arrow bytes, and input batch count. An +//! oversized input passes through unchanged instead of being split or copied. +//! +//! Keep inputs until a group is ready, then concatenate once at its actual size. + +use std::collections::VecDeque; +use std::ptr::NonNull; + +use arrow::array::{Array, AsArray, RecordBatch}; +use arrow::buffer::Buffer; +use arrow::compute::concat_batches; +use arrow::datatypes::{DataType, Int32Type, Int64Type}; +use datafusion::common::{DataFusionError, Result}; + +const MAX_ROWS: usize = 65_536; +const MAX_GROUP_BATCHES: usize = 1024; +const MAX_GROUP_BYTES: usize = 8 * 1024 * 1024; + +/// Allocation-free accounting for one decoded batch, discarded after admission. +/// IPC columns normally share their message-body allocation. Remember that first +/// owner so its slices are charged once; other owners may be overcounted safely. +/// Logical copy bytes are tracked separately because columns may alias values +/// that concatenation will materialize independently. +#[derive(Default)] +struct BatchBytes { + first_owner: Option<NonNull<u8>>, + retained: usize, + copied: usize, +} + +impl BatchBytes { + /// Charges a borrowed buffer's allocation and the bytes its output will copy. + /// Uses the allocation base, not the slice address, to deduplicate the common + /// IPC owner before inspecting its capacity again. Saturating counters + /// conservatively reject overflow. No ownership is retained and no memory + /// is allocated or read from the buffer contents. + fn add_buffer(&mut self, buffer: &Buffer, copied: usize) { + self.copied = self.copied.saturating_add(copied); + let owner = buffer.data_ptr(); + if self.first_owner == Some(owner) { + return; + } + let capacity = buffer.capacity().max(buffer.len()); + if capacity != 0 { + if self.first_owner.is_none() { + self.first_owner = Some(owner); + } + self.retained = self.retained.saturating_add(capacity); + } + } + + /// Charges a borrowed concrete array's null buffer and object metadata after + /// its values. The caller supplies the rounded output bitmap size. Static + /// dispatch avoids repeating trait-object lookups after the array downcast; + /// ownership, allocation-free accounting, and saturation match `add_buffer`. + #[inline] + fn add_nulls_and_metadata<T: Array>(&mut self, array: &T, copied_null_bytes: usize) { + if let Some(nulls) = array.nulls() { + self.add_buffer(nulls.buffer(), copied_null_bytes); + } + let array_metadata = std::mem::size_of::<T>(); + self.retained = self.retained.saturating_add(array_metadata); + self.copied = self.copied.saturating_add(array_metadata); + } +} + +/// Ordered grouping of batches returned by the native shuffle IPC decoder. +/// +/// `push` and `finish` enqueue output. Inspect it with `peek_ready`, and drain +/// it with `take_ready` before pushing another input. Call `finish` at EOF. +/// A singleton, an oversized input, or an input at least half the row target is +/// never copied. The byte budget bounds retained group inputs, not the whole +/// reader: concatenation also owns its output and may have one decoded lookahead +/// input. Callers must drop the helper on any error. +#[derive(Debug)] +pub struct DecodedBatchCoalescer { + max_rows: usize, + max_bytes: usize, + batches: Vec<RecordBatch>, + rows: usize, + bytes: usize, + ready: VecDeque<RecordBatch>, +} + +impl DecodedBatchCoalescer { + /// Creates an empty, reader-owned grouper with the production byte budget. + /// Rejects row targets outside 1..=65536 without retaining any buffers. + pub fn try_new(max_rows: usize) -> Result<Self> { + Self::try_new_with_limits(max_rows, MAX_GROUP_BYTES) + } + + /// Creates an empty grouper with explicit limits, allowing small test budgets. + /// Invalid limits return a configuration error. Keeping the byte budget at + /// most 8 MiB also keeps every concatenated STRING/BINARY offset within i32. + fn try_new_with_limits(max_rows: usize, max_bytes: usize) -> Result<Self> { + if !(1..=MAX_ROWS).contains(&max_rows) { + return Err(DataFusionError::Configuration( + "Decoded shuffle grouping requires 1..65536 rows".into(), + )); + } + if !(1..=MAX_GROUP_BYTES).contains(&max_bytes) { + return Err(DataFusionError::Configuration( + "Decoded shuffle grouping requires 1..8388608 bytes".into(), + )); + } + Ok(Self { + max_rows, + max_bytes, + batches: Vec::new(), + rows: 0, + bytes: 0, + ready: VecDeque::new(), + }) + } + + /// Transfers the oldest ready output to the caller without copying buffers. + /// Returns None when empty and leaves any pending group unchanged. + pub fn take_ready(&mut self) -> Option<RecordBatch> { + self.ready.pop_front() + } + + /// Borrows the oldest ready output until the next mutation, without draining it. + pub fn peek_ready(&self) -> Option<&RecordBatch> { + self.ready.front() + } + + /// Takes ownership of one decoded input. A flush followed by a passthrough can make two + /// outputs ready; drain both before reading another input. Drop this helper on error. + pub fn push(&mut self, batch: RecordBatch) -> Result<()> { + if !self.ready.is_empty() { + return Err(DataFusionError::Execution( + "Drain shuffle grouping output before reading another input".into(), + )); + } + let rows = batch.num_rows(); + let bytes = self.eligible_bytes(&batch); + + if !self.batches.is_empty() + && (bytes.is_none_or(|bytes| bytes > self.max_bytes - self.bytes) + || self.batches[0].schema() != batch.schema() + || rows > self.max_rows - self.rows) + { + self.flush()?; + } + if let Some(bytes) = bytes { + self.rows += rows; + self.bytes += bytes; + self.batches.push(batch); + if self.rows == self.max_rows + || self.bytes == self.max_bytes + || self.batches.len() == MAX_GROUP_BATCHES + { + self.flush()?; + } + } else { + self.ready.push_back(batch); + } + Ok(()) + } + + /// Flushes the final group. Repeated calls do not copy or emit it again. + pub fn finish(&mut self) -> Result<()> { + self.flush() + } + + /// Returns the larger of retained and logical copy bytes for a groupable input. + /// Examines flat-column buffer metadata and string/binary endpoint offsets, + /// never rows. The common shared IPC allocation is charged once. Other shared + /// owners can be overcounted safely, and slices retain their full backing + /// allocation charge. RecordBatch guarantees matching field and array types, + /// so dispatch borrows schema fields without cloning or inspecting array + /// trait objects twice. None means pass through without copying. + fn eligible_bytes(&self, batch: &RecordBatch) -> Option<usize> { + if batch.num_rows() == 0 + || batch.num_rows() >= self.max_rows.div_ceil(2) + || batch.num_columns() == 0 + { + return None; + } + let rows = batch.num_rows(); + // Arrow bitmap builders round their allocation to 64-byte blocks. + let bitmap_bytes = rows.div_ceil(8).next_multiple_of(64); + let mut bytes = BatchBytes::default(); + for (field, column) in batch.schema_ref().fields().iter().zip(batch.columns()) { + match field.data_type() { + DataType::Int32 => { + let array = column.as_primitive::<Int32Type>(); + bytes.add_buffer(array.values().inner(), rows * 4); + bytes.add_nulls_and_metadata(array, bitmap_bytes); + } + DataType::Int64 => { + let array = column.as_primitive::<Int64Type>(); + bytes.add_buffer(array.values().inner(), rows * 8); + bytes.add_nulls_and_metadata(array, bitmap_bytes); + } + DataType::Boolean => { + let array = column.as_boolean(); + bytes.add_buffer(array.values().inner(), bitmap_bytes); + bytes.add_nulls_and_metadata(array, bitmap_bytes); + } + DataType::Utf8 => { + let array = column.as_string::<i32>(); + let offsets = array.value_offsets(); + bytes.add_buffer(array.offsets().inner().inner(), (rows + 1) * 4); + bytes.add_buffer(array.values(), (offsets[rows] - offsets[0]) as usize); + bytes.add_nulls_and_metadata(array, bitmap_bytes); + } + DataType::Binary => { + let array = column.as_binary::<i32>(); + let offsets = array.value_offsets(); + bytes.add_buffer(array.offsets().inner().inner(), (rows + 1) * 4); + bytes.add_buffer(array.values(), (offsets[rows] - offsets[0]) as usize); + bytes.add_nulls_and_metadata(array, bitmap_bytes); + } + _ => return None, + } + if bytes.retained.max(bytes.copied) > self.max_bytes { + return None; + } + } + Some(bytes.retained.max(bytes.copied)) + } + + fn flush(&mut self) -> Result<()> { + let batches = std::mem::take(&mut self.batches); Review Comment: `mem::take` drops the `Vec`'s capacity at every flush, so each group of up to 1,024 inputs regrows it from zero. Calling `concat_batches(&schema, &self.batches)` and then `self.batches.clear()` (or `self.ready.extend(self.batches.drain(..))` for the pass-through case) would keep the allocation for the life of the reader. This only matters if the hand-written coalescer stays. ########## spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala: ########## @@ -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") { Review Comment: The result-correctness half of this test could be a fixture next to `join/broadcast_shared_decode.sql`, with `-- Config: spark.comet.exec.broadcastExchange.producerCoalesce.enabled=true` and an AQE on/off `ConfigMatrix`. A fixture can also cover BIGINT, BOOLEAN and BINARY with NULLs through the JVM path, which today only the Rust unit tests exercise. The Scala test would then keep only the `numCoalescedBatches` comparison, and the unsupported-schema test below could become a DECIMAL column in the same fixture. ########## spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala: ########## @@ -110,9 +115,24 @@ case class CometBroadcastExchangeExec( private lazy val maxBroadcastRows = 512000000L private def getByteArrayRdd(plan: SparkPlan): RDD[(Long, ChunkedByteBuffer)] = { - plan.executeColumnar().mapPartitionsInternal { iter => - Utils.serializeBatches(iter) + // Capture the read mode on the driver. A new RDD keeps it local to this broadcast consumer, + // including when AQE or exchange reuse shares the original shuffle with another operator. + val input = plan.executeColumnar() + val supported = plan.output.nonEmpty && plan.output.forall { attribute => + attribute.dataType match { + case IntegerType | LongType | BooleanType | StringType | BinaryType => true + case _ => false + } + } + val readInput = input match { + case shuffle: CometShuffledBatchRDD + if CometConf.COMET_BROADCAST_PRODUCER_COALESCE_ENABLED.get(conf) && + supported && shuffle.supportsNativeBatchCoalescing => + longMetric("nativeProducerBatching") += 1L + shuffle.withNativeBatchCoalescing(math.min(CometConf.COMET_BATCH_SIZE.get(conf), 65536)) Review Comment: `65536` repeats `MAX_ROWS` in the native crate, and the two have to agree because `try_new` returns an error above it, which would fail the task. Could the clamp live in one place, for example by having native `try_new` use `max_rows.min(MAX_ROWS)` instead of rejecting, so the JVM only passes `spark.comet.batchSize`? -- 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]
