sunchao commented on code in PR #7258:
URL: https://github.com/apache/datafusion-comet/pull/7258#discussion_r4406374041
##########
native/core/src/execution/jni_api.rs:
##########
@@ -2105,6 +2106,150 @@ 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, |_| {
+ if max_rows <= 0 {
Review Comment:
Centralized the check in `try_new`: JNI maps negative values to zero before
converting to usize, so nonpositive targets use the same validation and message.
Updated in `48e1024bfc`.
##########
native/core/src/execution/jni_api.rs:
##########
@@ -2105,6 +2106,150 @@ 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, |_| {
+ if max_rows <= 0 {
+ return Err(CometError::Internal(
+ "Decoded shuffle grouping requires a positive row
target".into(),
+ ));
+ }
+ let coalescer = DecodedBatchCoalescer::try_new(max_rows as usize)?;
+ Ok(Box::into_raw(Box::new(coalescer)) as jlong)
+ })
+}
+
+#[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, || {
+ let capacity = env.get_direct_buffer_capacity(&byte_buffer)?;
+ if length < 4 || length as usize > capacity || num_columns < 0 {
Review Comment:
Both JNI decode paths now use the same small direct-buffer helper, including
the capacity check. I also removed the extra checks from `takeShuffleBatch`;
its safety contract now states that both address arrays have the column count
returned by peek, matching how `NativeUtil.getNextBatch` allocates them.
Updated in `48e1024bfc`.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala:
##########
@@ -121,6 +124,45 @@ class CometBlockStoreShuffleReader[K, C](
})
.map(b => (0, b))
+ withReadMetrics(recordIter)
+ }
+
+ /** One decoder groups frames across fetched blocks for a single broadcast
producer task. */
+ private[shuffle] def readWithCoalescing(maxRows: Int): Iterator[Product2[K,
C]] = {
+ val input = readAsRawStream()
+ val nativeUtil =
+ try new NativeUtil()
Review Comment:
Simplified this to direct construction followed by listener registration.
The task-completion listener and decoder runtime-error cleanup remain; the
constructor-only catch blocks are gone.
Updated in `48e1024bfc`.
##########
native/shuffle/src/decoded_batch_coalescer.rs:
##########
@@ -0,0 +1,579 @@
+// 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<NonNull<u8>>,
+ 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<RecordBatch>,
+ rows: usize,
+ bytes: usize,
+ ready: VecDeque<RecordBatch>,
+}
+
+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> {
+ 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 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<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()
+ }
+
+ /// Bounds both retained backing allocations and output copies for flat
columns.
+ /// 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();
+ // ArrayData reports logical bytes; bitmap builders round up to
64-byte blocks.
+ let bitmap_bytes = rows.div_ceil(8);
Review Comment:
I kept the conservative bitmap allowance. You are right that charging
padding per input can overcount, but removing it does not preserve the
output-copy bound: `get_slice_memory_size()` reports logical bitmap bytes,
while concatenation rounds Boolean values and null buffers to 64-byte blocks.
With aliased Boolean columns, retained source buffers are shared but each
output column needs its own bitmaps, so logical bytes plus metadata need not
cover those allocations. I added a short comment explaining why the allowance
stays; `BatchBytes` still keeps that retained-versus-copied accounting together.
Updated in `48e1024bfc`.
##########
native/shuffle/src/decoded_batch_coalescer.rs:
##########
@@ -0,0 +1,579 @@
+// 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<NonNull<u8>>,
+ 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<RecordBatch>,
+ rows: usize,
+ bytes: usize,
+ ready: VecDeque<RecordBatch>,
+}
+
+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> {
+ 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 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<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()
+ }
+
+ /// Bounds both retained backing allocations and output copies for flat
columns.
+ /// 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();
+ // ArrayData reports logical bytes; bitmap builders round up to
64-byte blocks.
+ let bitmap_bytes = rows.div_ceil(8);
+ let padded_bitmap_bytes = bitmap_bytes.next_multiple_of(64);
+ let mut bytes = BatchBytes::default();
+ for 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<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();
+ 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)));
+ }
+
+ let mut coalescer = DecodedBatchCoalescer::try_new_with_limits(32,
two_batches).unwrap();
+ coalescer.push(inputs[0].clone()).unwrap();
+ assert!(coalescer.peek_ready().is_none());
+ coalescer.push(inputs[1].clone()).unwrap();
+ assert_eq!(coalescer.take_ready(), Some(ints(0..2)));
+ coalescer.push(inputs[2].clone()).unwrap();
+ assert!(coalescer.take_ready().is_none());
+ coalescer.finish().unwrap();
+ assert_eq!(coalescer.take_ready(), Some(ints(2..3)));
+ }
+
+ #[test]
+ fn wide_json_values_pass_through_without_copying() {
Review Comment:
Removed the JSON-specific duplicate, the repeated push/peek/take block, and
the weak-reference lifetime test. The sliced-backing-buffer case now has a
second small input after passthrough and checks that grouping resumes. The
remaining native and JVM cases retain the observable flush/passthrough and
ordering coverage.
Updated in `48e1024bfc`.
--
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]
