viirya commented on code in PR #7258:
URL: https://github.com/apache/datafusion-comet/pull/7258#discussion_r4392381712
##########
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:
This adds a 0/1 sum metric to every broadcast exchange, and while the
feature is off by default it will always show 0. It also counts the gate, not
the effect: it reports 1 when Celeborn takes over at read time or when no batch
qualifies. If we want a metric, the number of input batches actually grouped
(or the output count) would say whether batching did anything. Otherwise the
tests could assert on the RDD instead.
##########
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:
The hand-written per-type accounting is what limits this to five types.
`ArrayData::get_slice_memory_size()` already computes the bytes a slice copies
for any layout. For retained bytes, you could keep the current `data_ptr` dedup
over `ArrayData::buffers()` plus nulls. That would cover fixed-width types like
Float64/Date32/Timestamp/Decimal128 without a `match`, and the JVM-side type
gate could go away. Is the per-column `to_data()` allocation the reason for
avoiding that? Compared with the IPC decode that has just run for the same
batch, I'd expect it to be negligible. If you measured otherwise, it would be
good to note that here.
##########
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:
The comments here are hard to follow, and much of the detail is about
micro-optimizations rather than invariants. Examples are "Static dispatch
avoids repeating trait-object lookups after the array downcast" and
"Allocation-free accounting… discarded after admission". Could these be trimmed
to what a maintainer needs: what is counted, why the first owner is
deduplicated, and that overcounting is safe?
##########
native/core/src/execution/jni_api.rs:
##########
@@ -2105,6 +2106,145 @@ pub unsafe extern "system" fn
Java_org_apache_comet_Native_decodeShuffleBlock(
})
}
+/// Decode a local frame into owned arrays. The direct buffer can be reused
after this call.
+fn append_shuffle_batch(
+ coalescer: &mut DecodedBatchCoalescer,
+ block: &[u8],
+ num_columns: usize,
+) -> CometResult<()> {
+ let batch = read_ipc_compressed(block)?;
+ if batch.num_columns() != num_columns {
+ return Err(CometError::Internal(
+ "Shuffle frame column count mismatch".into(),
+ ));
+ }
+ coalescer.push(batch)?;
+ Ok(())
+}
+
+/// The reader owns this handle and serializes native calls with cleanup.
+/// # Safety
+/// A nonzero handle must refer to a live coalescer exclusively borrowed for
this call.
+unsafe fn shuffle_batch_coalescer<'a>(handle: jlong) -> CometResult<&'a mut
DecodedBatchCoalescer> {
+ unsafe { (handle as *mut DecodedBatchCoalescer).as_mut() }
+ .ok_or_else(|| CometError::Internal("Closed shuffle batch
coalescer".into()))
+}
+
+fn shuffle_batch_ready_columns(coalescer: &DecodedBatchCoalescer) -> jint {
+ coalescer
+ .peek_ready()
+ .map_or(-1, |batch| batch.num_columns() as jint)
+}
+
+#[no_mangle]
+pub extern "system" fn
Java_org_apache_comet_Native_createShuffleBatchCoalescer(
+ e: EnvUnowned,
+ _class: JClass,
+ max_rows: jint,
+) -> jlong {
+ try_unwrap_or_throw(&e, |_| {
+ let coalescer = DecodedBatchCoalescer::try_new(max_rows as usize)?;
+ Ok(Box::into_raw(Box::new(coalescer)) as jlong)
+ })
+}
+
+#[no_mangle]
+/// Decode and retain one local frame, returning ready columns or -1 for more
input.
+/// # Safety
+/// The handle must be live and exclusively owned by this reader. On error,
close the reader.
+pub unsafe extern "system" fn Java_org_apache_comet_Native_appendShuffleBatch(
+ e: EnvUnowned,
+ _class: JClass,
+ handle: jlong,
+ byte_buffer: JByteBuffer,
+ length: jint,
+ num_columns: jint,
+ tracing_enabled: jboolean,
+) -> jint {
+ try_unwrap_or_throw(&e, |env| {
+ with_trace("decodeShuffleBlock", tracing_enabled != JNI_FALSE, || {
Review Comment:
`appendShuffleBatch` reuses the `"decodeShuffleBlock"` trace name, so the
two paths can't be told apart in a trace. Maybe `"appendShuffleBatch"` to match
`exportShuffleBatch` / `coalesceShuffleBatches`?
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala:
##########
@@ -53,31 +54,44 @@ case class NativeBatchDecoderIterator(
private var batch: Option[ColumnarBatch] = None
private val validateRemoteFrames =
in.isInstanceOf[CometShuffleReadFailureHandler]
private var remoteDecoderHandle = 0L
+ private var coalescerHandle = 0L
+ private var endOfInput = false
require(
!validateRemoteFrames || expectedSchema.exists(_ != null),
"Remote shuffle decoding requires the expected Spark schema")
+ require(!validateRemoteFrames || coalescing.isEmpty, "Batch coalescing is
local-shuffle only")
import NativeBatchDecoderIterator._
override def hasNext: Boolean = {
- synchronized {
- if (in == null || isClosed) {
- return false
- }
- if (batch.isDefined) {
- return true
- }
+ try {
+ synchronized {
+ if (in == null || isClosed) {
+ return false
+ }
+ if (batch.isDefined) {
+ return true
+ }
- // Clear ownership before cleanup so a failed close cannot release this
batch twice.
- if (currentBatch != null) {
- val previous = currentBatch
- currentBatch = null
- previous.close()
+ // Clear ownership before cleanup so a failed close cannot release
this batch twice.
+ if (currentBatch != null) {
+ val previous = currentBatch
+ currentBatch = null
+ previous.close()
+ }
}
- }
- fetchNext()
+ if (coalescing.isDefined) fetchCoalesced() else fetchNext()
+ } catch {
+ case NonFatal(failure) if coalescing.isDefined =>
Review Comment:
In coalescing mode, any failure in `hasNext` now closes the iterator, but
the non-coalescing path keeps its existing behaviour. Is there a reason the two
modes need different failure semantics? If closing early is right here, it
seems right for the per-block path too. If it isn't needed, the task completion
listener already covers cleanup.
##########
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:
The Spark-type check duplicates the native `eligible_bytes` check, and the
two can disagree (e.g. dictionary-encoded strings). Since unsupported batches
already pass through unchanged natively, could the JVM gate be reduced to the
shuffle/reader checks?
##########
native/shuffle/src/decoded_batch_coalescer.rs:
##########
@@ -0,0 +1,582 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Group small decoded shuffle batches before exporting them to the JVM.
+//!
+//! Groups are bounded by rows, retained Arrow bytes, and input batch count. An
+//! oversized input passes through unchanged instead of being split or copied.
+//!
+//! Keep inputs until a group is ready, then concatenate once at its actual
size.
+
+use std::collections::VecDeque;
+use std::ptr::NonNull;
+
+use arrow::array::{Array, AsArray, RecordBatch};
+use arrow::buffer::Buffer;
+use arrow::compute::concat_batches;
+use arrow::datatypes::{DataType, Int32Type, Int64Type};
+use datafusion::common::{DataFusionError, Result};
+
+const MAX_ROWS: usize = 65_536;
+const MAX_GROUP_BATCHES: usize = 1024;
+const MAX_GROUP_BYTES: usize = 8 * 1024 * 1024;
+
+/// Allocation-free accounting for one decoded batch, discarded after
admission.
+/// IPC columns normally share their message-body allocation. Remember that
first
+/// owner so its slices are charged once; other owners may be overcounted
safely.
+/// Logical copy bytes are tracked separately because columns may alias values
+/// that concatenation will materialize independently.
+#[derive(Default)]
+struct BatchBytes {
+ first_owner: Option<NonNull<u8>>,
+ retained: usize,
+ copied: usize,
+}
+
+impl BatchBytes {
+ /// Charges a borrowed buffer's allocation and the bytes its output will
copy.
+ /// Uses the allocation base, not the slice address, to deduplicate the
common
+ /// IPC owner before inspecting its capacity again. Saturating counters
+ /// conservatively reject overflow. No ownership is retained and no memory
+ /// is allocated or read from the buffer contents.
+ fn add_buffer(&mut self, buffer: &Buffer, copied: usize) {
+ self.copied = self.copied.saturating_add(copied);
+ let owner = buffer.data_ptr();
+ if self.first_owner == Some(owner) {
+ return;
+ }
+ let capacity = buffer.capacity().max(buffer.len());
+ if capacity != 0 {
+ if self.first_owner.is_none() {
+ self.first_owner = Some(owner);
+ }
+ self.retained = self.retained.saturating_add(capacity);
+ }
+ }
+
+ /// Charges a borrowed concrete array's null buffer and object metadata
after
+ /// its values. The caller supplies the rounded output bitmap size. Static
+ /// dispatch avoids repeating trait-object lookups after the array
downcast;
+ /// ownership, allocation-free accounting, and saturation match
`add_buffer`.
+ #[inline]
+ fn add_nulls_and_metadata<T: Array>(&mut self, array: &T,
copied_null_bytes: usize) {
+ if let Some(nulls) = array.nulls() {
+ self.add_buffer(nulls.buffer(), copied_null_bytes);
+ }
+ let array_metadata = std::mem::size_of::<T>();
+ self.retained = self.retained.saturating_add(array_metadata);
+ self.copied = self.copied.saturating_add(array_metadata);
+ }
+}
+
+/// Ordered grouping of batches returned by the native shuffle IPC decoder.
+///
+/// `push` and `finish` enqueue output. Inspect it with `peek_ready`, and drain
+/// it with `take_ready` before pushing another input. Call `finish` at EOF.
+/// A singleton, an oversized input, or an input at least half the row target
is
+/// never copied. The byte budget bounds retained group inputs, not the whole
+/// reader: concatenation also owns its output and may have one decoded
lookahead
+/// input. Callers must drop the helper on any error.
+#[derive(Debug)]
+pub struct DecodedBatchCoalescer {
+ max_rows: usize,
+ max_bytes: usize,
+ batches: Vec<RecordBatch>,
+ rows: usize,
+ bytes: usize,
+ ready: VecDeque<RecordBatch>,
+}
+
+impl DecodedBatchCoalescer {
+ /// Creates an empty, reader-owned grouper with the production byte budget.
+ /// Rejects row targets outside 1..=65536 without retaining any buffers.
+ pub fn try_new(max_rows: usize) -> Result<Self> {
+ Self::try_new_with_limits(max_rows, MAX_GROUP_BYTES)
+ }
+
+ /// Creates an empty grouper with explicit limits, allowing small test
budgets.
+ /// Invalid limits return a configuration error. Keeping the byte budget at
+ /// most 8 MiB also keeps every concatenated STRING/BINARY offset within
i32.
+ fn try_new_with_limits(max_rows: usize, max_bytes: usize) -> Result<Self> {
+ if !(1..=MAX_ROWS).contains(&max_rows) {
+ return Err(DataFusionError::Configuration(
+ "Decoded shuffle grouping requires 1..65536 rows".into(),
+ ));
+ }
+ if !(1..=MAX_GROUP_BYTES).contains(&max_bytes) {
+ return Err(DataFusionError::Configuration(
+ "Decoded shuffle grouping requires 1..8388608 bytes".into(),
+ ));
+ }
+ Ok(Self {
+ max_rows,
+ max_bytes,
+ batches: Vec::new(),
+ rows: 0,
+ bytes: 0,
+ ready: VecDeque::new(),
+ })
+ }
+
+ /// Transfers the oldest ready output to the caller without copying
buffers.
+ /// Returns None when empty and leaves any pending group unchanged.
+ pub fn take_ready(&mut self) -> Option<RecordBatch> {
+ self.ready.pop_front()
+ }
+
+ /// Borrows the oldest ready output until the next mutation, without
draining it.
+ pub fn peek_ready(&self) -> Option<&RecordBatch> {
+ self.ready.front()
+ }
+
+ /// Takes ownership of one decoded input. A flush followed by a
passthrough can make two
+ /// outputs ready; drain both before reading another input. Drop this
helper on error.
+ pub fn push(&mut self, batch: RecordBatch) -> Result<()> {
+ if !self.ready.is_empty() {
+ return Err(DataFusionError::Execution(
+ "Drain shuffle grouping output before reading another
input".into(),
+ ));
+ }
+ let rows = batch.num_rows();
+ let bytes = self.eligible_bytes(&batch);
+
+ if !self.batches.is_empty()
+ && (bytes.is_none_or(|bytes| bytes > self.max_bytes - self.bytes)
+ || self.batches[0].schema() != batch.schema()
+ || rows > self.max_rows - self.rows)
+ {
+ self.flush()?;
+ }
+ if let Some(bytes) = bytes {
+ self.rows += rows;
+ self.bytes += bytes;
+ self.batches.push(batch);
+ if self.rows == self.max_rows
+ || self.bytes == self.max_bytes
+ || self.batches.len() == MAX_GROUP_BATCHES
+ {
+ self.flush()?;
+ }
+ } else {
+ self.ready.push_back(batch);
+ }
+ Ok(())
+ }
+
+ /// Flushes the final group. Repeated calls do not copy or emit it again.
+ pub fn finish(&mut self) -> Result<()> {
+ self.flush()
+ }
+
+ /// Returns the larger of retained and logical copy bytes for a groupable
input.
+ /// Examines flat-column buffer metadata and string/binary endpoint
offsets,
+ /// never rows. The common shared IPC allocation is charged once. Other
shared
+ /// owners can be overcounted safely, and slices retain their full backing
+ /// allocation charge. RecordBatch guarantees matching field and array
types,
+ /// so dispatch borrows schema fields without cloning or inspecting array
+ /// trait objects twice. None means pass through without copying.
+ fn eligible_bytes(&self, batch: &RecordBatch) -> Option<usize> {
+ if batch.num_rows() == 0
+ || batch.num_rows() >= self.max_rows.div_ceil(2)
+ || batch.num_columns() == 0
+ {
+ return None;
+ }
+ let rows = batch.num_rows();
+ // Arrow bitmap builders round their allocation to 64-byte blocks.
+ let bitmap_bytes = rows.div_ceil(8).next_multiple_of(64);
+ let mut bytes = BatchBytes::default();
+ for (field, column) in
batch.schema_ref().fields().iter().zip(batch.columns()) {
+ match field.data_type() {
+ DataType::Int32 => {
+ let array = column.as_primitive::<Int32Type>();
+ bytes.add_buffer(array.values().inner(), rows * 4);
+ bytes.add_nulls_and_metadata(array, bitmap_bytes);
+ }
+ DataType::Int64 => {
+ let array = column.as_primitive::<Int64Type>();
+ bytes.add_buffer(array.values().inner(), rows * 8);
+ bytes.add_nulls_and_metadata(array, bitmap_bytes);
+ }
+ DataType::Boolean => {
+ let array = column.as_boolean();
+ bytes.add_buffer(array.values().inner(), bitmap_bytes);
+ bytes.add_nulls_and_metadata(array, bitmap_bytes);
+ }
+ DataType::Utf8 => {
+ let array = column.as_string::<i32>();
+ let offsets = array.value_offsets();
+ bytes.add_buffer(array.offsets().inner().inner(), (rows +
1) * 4);
+ bytes.add_buffer(array.values(), (offsets[rows] -
offsets[0]) as usize);
+ bytes.add_nulls_and_metadata(array, bitmap_bytes);
+ }
+ DataType::Binary => {
+ let array = column.as_binary::<i32>();
+ let offsets = array.value_offsets();
+ bytes.add_buffer(array.offsets().inner().inner(), (rows +
1) * 4);
+ bytes.add_buffer(array.values(), (offsets[rows] -
offsets[0]) as usize);
+ bytes.add_nulls_and_metadata(array, bitmap_bytes);
+ }
+ _ => return None,
+ }
+ if bytes.retained.max(bytes.copied) > self.max_bytes {
+ return None;
+ }
+ }
+ Some(bytes.retained.max(bytes.copied))
+ }
+
+ fn flush(&mut self) -> Result<()> {
+ let batches = std::mem::take(&mut self.batches);
+ self.rows = 0;
+ self.bytes = 0;
+ if batches.len() < 2 {
+ self.ready.extend(batches);
+ } else {
+ self.ready
+ .push_back(concat_batches(&batches[0].schema(), &batches)?);
+ }
+ Ok(())
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use std::sync::Arc;
+
+ use arrow::array::{ArrayRef, BinaryArray, BooleanArray, Int32Array,
Int64Array, StringArray};
+ use arrow::datatypes::{Field, Schema};
+ use arrow::ipc::writer::StreamWriter;
+
+ fn one_column(array: ArrayRef) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "x",
+ array.data_type().clone(),
+ true,
+ )]));
+ RecordBatch::try_new(schema, vec![array]).unwrap()
+ }
+
+ fn ints(values: impl IntoIterator<Item = i32>) -> RecordBatch {
+ one_column(Arc::new(Int32Array::from_iter_values(values)))
+ }
+
+ fn encode_frame(batch: &RecordBatch) -> Vec<u8> {
+ let mut frame = b"NONE".to_vec();
+ let mut writer = StreamWriter::try_new(&mut frame,
batch.schema().as_ref()).unwrap();
+ writer.write(batch).unwrap();
+ writer.finish().unwrap();
+ frame
+ }
+
+ fn collect(max_rows: usize, batches: Vec<RecordBatch>) -> Vec<RecordBatch>
{
+ collect_with_limits(max_rows, MAX_GROUP_BYTES, batches)
+ }
+
+ fn collect_with_limits(
+ max_rows: usize,
+ max_bytes: usize,
+ batches: Vec<RecordBatch>,
+ ) -> Vec<RecordBatch> {
+ let mut coalescer =
+ DecodedBatchCoalescer::try_new_with_limits(max_rows,
max_bytes).unwrap();
+ let mut output = Vec::new();
+ for batch in batches {
+ coalescer.push(batch).unwrap();
+ while let Some(batch) = coalescer.take_ready() {
+ output.push(batch);
+ }
+ }
+ coalescer.finish().unwrap();
+ coalescer.finish().unwrap();
+ while let Some(batch) = coalescer.take_ready() {
+ output.push(batch);
+ }
+ coalescer.finish().unwrap();
+ assert!(coalescer.peek_ready().is_none());
+ output
+ }
+
+ #[test]
+ fn singleton_large_and_near_full_inputs_are_not_copied() {
+ for batches in [
+ vec![ints(0..1)],
+ vec![ints(0..1), ints(1..7)],
+ vec![ints(0..6), ints(6..7)],
+ vec![ints(0..1), ints(1..10)],
+ ] {
+ let originals: Vec<_> = batches
+ .iter()
+ .map(|batch| batch.column(0).clone())
+ .collect();
+ let output = collect(8, batches);
+ assert_eq!(output.len(), originals.len());
+ for (batch, original) in output.iter().zip(originals) {
+ assert!(Arc::ptr_eq(batch.column(0), &original));
+ }
+ }
+ }
+
+ #[test]
+ fn row_and_input_count_limits_preserve_order() {
+ let batches = (0..4).map(|i| ints(i * 2..i * 2 + 2)).collect();
+ assert_eq!(collect(8, batches), vec![ints(0..8)]);
+ assert_eq!(
+ collect(8, vec![ints(0..3), ints(3..6), ints(6..9)]),
+ vec![ints(0..6), ints(6..9)]
+ );
+ let batches = (0..1025).map(|i| ints(i..i + 1)).collect();
+ assert_eq!(
+ collect(MAX_ROWS, batches),
+ vec![ints(0..1024), ints(1024..1025)]
+ );
+ }
+
+ #[test]
+ fn byte_limit_flushes_at_and_before_crossing_the_boundary() {
+ let inputs: Vec<_> = (0..3).map(|i| ints(i..i + 1)).collect();
+ let two_batches = inputs[0].get_array_memory_size() +
inputs[1].get_array_memory_size();
+ assert_eq!(
+ collect_with_limits(32, two_batches, inputs.clone()),
+ vec![ints(0..2), ints(2..3)]
+ );
+ let output = collect_with_limits(32, two_batches - 1, inputs.clone());
+ assert_eq!(output, inputs);
+ for (output, input) in output.iter().zip(&inputs) {
+ assert!(Arc::ptr_eq(output.column(0), input.column(0)));
+ }
+
+ let mut coalescer = DecodedBatchCoalescer::try_new_with_limits(32,
two_batches).unwrap();
+ coalescer.push(inputs[0].clone()).unwrap();
+ assert!(coalescer.peek_ready().is_none());
+ coalescer.push(inputs[1].clone()).unwrap();
+ assert_eq!(coalescer.take_ready(), Some(ints(0..2)));
+ assert_eq!(coalescer.bytes, 0);
+ }
+
+ #[test]
+ fn wide_json_values_pass_through_without_copying() {
+ let small_json = format!("{{\"payload\":\"{}\"}}", "x".repeat(32));
+ let large_json = format!("{{\"payload\":\"{}\"}}", "x".repeat(2048));
+ let small =
one_column(Arc::new(StringArray::from(vec![small_json.as_str()])));
+ let large =
one_column(Arc::new(StringArray::from(vec![large_json.as_str()])));
+ let original = large.column(0).clone();
+ let max_bytes = 2 * small.get_array_memory_size();
+ assert!(large.get_array_memory_size() > max_bytes);
+ let output = collect_with_limits(
+ 32,
+ max_bytes,
+ vec![small.clone(), large.clone(), small.clone(), small, large],
+ );
+ let expected = one_column(Arc::new(StringArray::from(vec![
+ small_json.as_str(),
+ small_json.as_str(),
+ ])));
+ assert_eq!(output.len(), 4);
+ assert_eq!(output[0].num_rows(), 1);
+ assert!(Arc::ptr_eq(output[1].column(0), &original));
+ assert_eq!(output[2], expected);
+ assert!(Arc::ptr_eq(output[3].column(0), &original));
+ }
+
+ #[test]
+ fn sliced_backing_buffers_count_toward_the_byte_budget() {
+ let wide = "x".repeat(2048);
+ let backing = one_column(Arc::new(StringArray::from(vec!["a",
wide.as_str()])));
+ let slice = backing.slice(0, 1);
+ let small = one_column(Arc::new(StringArray::from(vec!["b"])));
+ let max_bytes = 2 * small.get_array_memory_size();
+ assert!(slice.get_array_memory_size() > max_bytes);
+ let original = slice.column(0).clone();
+ let output = collect_with_limits(32, max_bytes, vec![small.clone(),
slice, small]);
+ assert_eq!(output.len(), 3);
+ assert!(Arc::ptr_eq(output[1].column(0), &original));
+ }
+
+ #[test]
+ fn shared_ipc_body_is_charged_once_for_wide_batches() {
+ let schema = Arc::new(Schema::new(
+ (0..128)
+ .map(|column| Field::new(format!("c{column}"),
DataType::Int32, true))
+ .collect::<Vec<_>>(),
+ ));
+ let arrays: Vec<ArrayRef> = (0..128)
+ .map(|column| {
+ Arc::new(Int32Array::from_iter((0..64).map(|row| {
+ ((row + column + 1) % 17 != 0).then_some(row * 31 + column)
+ }))) as ArrayRef
+ })
+ .collect();
+ let input = RecordBatch::try_new(schema, arrays).unwrap();
+ let frame = encode_frame(&input);
+ let decoded = crate::read_ipc_compressed(&frame).unwrap();
+ let mut coalescer = DecodedBatchCoalescer::try_new(8192).unwrap();
+ let charged = coalescer.eligible_bytes(&decoded).unwrap();
+ let first_buffer = decoded
+ .column(0)
+ .as_primitive::<Int32Type>()
+ .values()
+ .inner();
+ let last_buffer = decoded
+ .column(127)
+ .as_primitive::<Int32Type>()
+ .values()
+ .inner();
+ assert_eq!(first_buffer.data_ptr(), last_buffer.data_ptr());
+ assert!(charged < 2 * first_buffer.capacity());
+ assert!(charged * 128 <= MAX_GROUP_BYTES);
+ assert!(decoded.get_array_memory_size() > MAX_GROUP_BYTES);
+
+ for _ in 0..128 {
+ coalescer
+ .push(crate::read_ipc_compressed(&frame).unwrap())
+ .unwrap();
+ }
+ let output = coalescer.take_ready().unwrap();
+ assert_eq!(output.num_rows(), 8192);
+ assert!(coalescer.take_ready().is_none());
+ for offset in (0..8192).step_by(64) {
+ assert_eq!(output.slice(offset, 64), input);
+ }
+ }
+
+ #[test]
+ fn aliased_columns_are_charged_for_each_output_copy() {
+ let array = ints(0..160).column(0).clone();
+ let schema = Arc::new(Schema::new(
+ (0..16)
+ .map(|column| Field::new(format!("c{column}"),
DataType::Int32, false))
+ .collect::<Vec<_>>(),
+ ));
+ let input = RecordBatch::try_new(schema, vec![array.clone();
16]).unwrap();
+ let metadata = 16 * std::mem::size_of_val(array.as_ref());
+ let copied = metadata + 16 * 160 * 4;
+ let retained = metadata
+ + array
+ .as_primitive::<Int32Type>()
+ .values()
+ .inner()
+ .capacity();
+ assert!(copied > retained);
+ let coalescer = DecodedBatchCoalescer::try_new(8192).unwrap();
+ assert_eq!(coalescer.eligible_bytes(&input), Some(copied));
+ let output = collect_with_limits(8192, copied * 2 - 1,
vec![input.clone(), input.clone()]);
+ assert_eq!(output, vec![input.clone(), input]);
+ for batch in output {
+ for column in batch.columns() {
+ assert!(Arc::ptr_eq(column, &array));
+ }
+ }
+ }
+
+ #[test]
+ fn invalid_byte_budgets_are_configuration_errors() {
Review Comment:
A few tests overlap. `invalid_byte_budgets_are_configuration_errors` and
`empty_inputs_and_invalid_limits` both cover invalid limits.
`collect_with_limits` calls `finish()` three times, which mixes an idempotence
check into a general helper. Several assertions also pin internal state
(`coalescer.bytes == 0`). Could the idempotence check become its own test, with
the rest asserting observable output only?
--
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]
