Skip to main content

arrow_select/
coalesce.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! [`BatchCoalescer`]  concatenates multiple [`RecordBatch`]es after
19//! operations such as [`filter`] and [`take`].
20//!
21//! [`filter`]: crate::filter::filter
22//! [`take`]: crate::take::take
23use crate::filter::{FilterBuilder, FilterPredicate, FilterSelection};
24use crate::take::take_record_batch;
25use arrow_array::types::{BinaryViewType, StringViewType};
26use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, downcast_primitive};
27use arrow_schema::{ArrowError, DataType, SchemaRef};
28use std::collections::VecDeque;
29use std::sync::Arc;
30// Originally From DataFusion's coalesce module:
31// https://github.com/apache/datafusion/blob/9d2f04996604e709ee440b65f41e7b882f50b788/datafusion/physical-plan/src/coalesce/mod.rs#L26-L25
32
33mod byte_view;
34mod generic;
35mod primitive;
36
37use byte_view::InProgressByteViewArray;
38use generic::GenericInProgressArray;
39use primitive::InProgressPrimitiveArray;
40
41fn has_sparse_filter_copy(data_type: &DataType) -> bool {
42    data_type.is_primitive() || matches!(data_type, DataType::Utf8View | DataType::BinaryView)
43}
44
45/// Maximum selected row fraction for the fused sparse-filter copy path.
46///
47/// Shared benchmark results show this path helps low-selectivity filters, but
48/// can regress once the filter becomes denser. Keep this as a cheap integer
49/// threshold on the hot path: `selected_count <= filter_len / 16`.
50const SPARSE_FILTER_COPY_MAX_SELECTIVITY_DENOMINATOR: usize = 16;
51
52fn should_use_sparse_filter_copy(filter_len: usize, selected_count: usize) -> bool {
53    selected_count <= filter_len / SPARSE_FILTER_COPY_MAX_SELECTIVITY_DENOMINATOR
54}
55
56/// Concatenate multiple [`RecordBatch`]es
57///
58/// Implements the common pattern of incrementally creating output
59/// [`RecordBatch`]es of a specific size from an input stream of
60/// [`RecordBatch`]es.
61///
62/// This is useful after operations such as [`filter`] and [`take`] that produce
63/// smaller batches, and we want to coalesce them into larger batches for
64/// further processing.
65///
66/// # Motivation
67///
68/// If we use [`concat_batches`] to implement the same functionality, there are 2 potential issues:
69/// 1. At least 2x peak memory (holding the input and output of concat)
70/// 2. 2 copies of the data (to create the output of filter and then create the output of concat)
71///
72/// See: <https://github.com/apache/arrow-rs/issues/6692> for more discussions
73/// about the motivation.
74///
75/// [`filter`]: crate::filter::filter
76/// [`take`]: crate::take::take
77/// [`concat_batches`]: crate::concat::concat_batches
78///
79/// # Example
80/// ```
81/// use arrow_array::record_batch;
82/// use arrow_select::coalesce::{BatchCoalescer};
83/// let batch1 = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
84/// let batch2 = record_batch!(("a", Int32, [4, 5])).unwrap();
85///
86/// // Create a `BatchCoalescer` that will produce batches with at least 4 rows
87/// let target_batch_size = 4;
88/// let mut coalescer = BatchCoalescer::new(batch1.schema(), 4);
89///
90/// // push the batches
91/// coalescer.push_batch(batch1).unwrap();
92/// // only pushed 3 rows (not yet 4, enough to produce a batch)
93/// assert!(coalescer.next_completed_batch().is_none());
94/// coalescer.push_batch(batch2).unwrap();
95/// // now we have 5 rows, so we can produce a batch
96/// let finished = coalescer.next_completed_batch().unwrap();
97/// // 4 rows came out (target batch size is 4)
98/// let expected = record_batch!(("a", Int32, [1, 2, 3, 4])).unwrap();
99/// assert_eq!(finished, expected);
100///
101/// // Have no more input, but still have an in-progress batch
102/// assert!(coalescer.next_completed_batch().is_none());
103/// // We can finish the batch, which will produce the remaining rows
104/// coalescer.finish_buffered_batch().unwrap();
105/// let expected = record_batch!(("a", Int32, [5])).unwrap();
106/// assert_eq!(coalescer.next_completed_batch().unwrap(), expected);
107///
108/// // The coalescer is now empty
109/// assert!(coalescer.next_completed_batch().is_none());
110/// ```
111///
112/// # Background
113///
114/// Generally speaking, larger [`RecordBatch`]es are more efficient to process
115/// than smaller [`RecordBatch`]es (until the CPU cache is exceeded) because
116/// there is fixed processing overhead per batch. This coalescer builds up these
117/// larger batches incrementally.
118///
119/// ```text
120/// ┌────────────────────┐
121/// │    RecordBatch     │
122/// │   num_rows = 100   │
123/// └────────────────────┘                 ┌────────────────────┐
124///                                        │                    │
125/// ┌────────────────────┐     Coalesce    │                    │
126/// │                    │      Batches    │                    │
127/// │    RecordBatch     │                 │                    │
128/// │   num_rows = 200   │  ─ ─ ─ ─ ─ ─ ▶  │                    │
129/// │                    │                 │    RecordBatch     │
130/// │                    │                 │   num_rows = 400   │
131/// └────────────────────┘                 │                    │
132///                                        │                    │
133/// ┌────────────────────┐                 │                    │
134/// │                    │                 │                    │
135/// │    RecordBatch     │                 │                    │
136/// │   num_rows = 100   │                 └────────────────────┘
137/// │                    │
138/// └────────────────────┘
139/// ```
140///
141/// # Notes:
142///
143/// 1. Output rows are produced in the same order as the input rows
144///
145/// 2. The output is a sequence of batches, with all but the last being at exactly
146///    `target_batch_size` rows.
147#[derive(Debug)]
148pub struct BatchCoalescer {
149    /// The input schema
150    schema: SchemaRef,
151    /// The target batch size (and thus size for views allocation). This is a
152    /// hard limit: the output batch will be exactly `target_batch_size`,
153    /// rather than possibly being slightly above.
154    target_batch_size: usize,
155    /// In-progress arrays
156    in_progress_arrays: Vec<Box<dyn InProgressArray>>,
157    /// True if some column still needs the materialized filter path.
158    has_non_specialized_filter_columns: bool,
159    /// Buffered row count. Always less than `batch_size`
160    buffered_rows: usize,
161    /// Completed batches
162    completed: VecDeque<RecordBatch>,
163    /// Biggest coalesce batch size. See [`Self::with_biggest_coalesce_batch_size`]
164    biggest_coalesce_batch_size: Option<usize>,
165}
166
167impl BatchCoalescer {
168    /// Create a new `BatchCoalescer`
169    ///
170    /// # Arguments
171    /// - `schema` - the schema of the output batches
172    /// - `target_batch_size` - the number of rows in each output batch.
173    ///   Typical values are `4096` or `8192` rows.
174    ///
175    pub fn new(schema: SchemaRef, target_batch_size: usize) -> Self {
176        let has_non_specialized_filter_columns = schema
177            .fields()
178            .iter()
179            .any(|field| !has_sparse_filter_copy(field.data_type()));
180        let in_progress_arrays = schema
181            .fields()
182            .iter()
183            .map(|field| create_in_progress_array(field.data_type(), target_batch_size))
184            .collect::<Vec<_>>();
185
186        Self {
187            schema,
188            target_batch_size,
189            in_progress_arrays,
190            has_non_specialized_filter_columns,
191            // We will for sure store at least one completed batch
192            completed: VecDeque::with_capacity(1),
193            buffered_rows: 0,
194            biggest_coalesce_batch_size: None,
195        }
196    }
197
198    /// Set the coalesce batch size limit (default `None`)
199    ///
200    /// This limit determine when batches should bypass coalescing. Intuitively,
201    /// batches that are already large are costly to coalesce and are efficient
202    /// enough to process directly without coalescing.
203    ///
204    /// If `Some(limit)`, batches larger than this limit will bypass coalescing
205    /// when there is no buffered data, or when the previously buffered data
206    /// already exceeds this limit.
207    ///
208    /// If `None`, all batches will be coalesced according to the
209    /// target_batch_size.
210    pub fn with_biggest_coalesce_batch_size(mut self, limit: Option<usize>) -> Self {
211        self.biggest_coalesce_batch_size = limit;
212        self
213    }
214
215    /// Get the current biggest coalesce batch size limit
216    ///
217    /// See [`Self::with_biggest_coalesce_batch_size`] for details
218    pub fn biggest_coalesce_batch_size(&self) -> Option<usize> {
219        self.biggest_coalesce_batch_size
220    }
221
222    /// Set the biggest coalesce batch size limit
223    ///
224    /// See [`Self::with_biggest_coalesce_batch_size`] for details
225    pub fn set_biggest_coalesce_batch_size(&mut self, limit: Option<usize>) {
226        self.biggest_coalesce_batch_size = limit;
227    }
228
229    /// Return the schema of the output batches
230    pub fn schema(&self) -> SchemaRef {
231        Arc::clone(&self.schema)
232    }
233
234    /// Push a batch into the Coalescer after applying a filter
235    ///
236    /// This is semantically equivalent of calling [`Self::push_batch`]
237    /// with the results from [`crate::filter::filter_record_batch`]
238    ///
239    /// # Example
240    /// ```
241    /// # use arrow_array::{record_batch, BooleanArray};
242    /// # use arrow_select::coalesce::BatchCoalescer;
243    /// let batch1 = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
244    /// let batch2 = record_batch!(("a", Int32, [4, 5, 6])).unwrap();
245    /// // Apply a filter to each batch to pick the first and last row
246    /// let filter = BooleanArray::from(vec![true, false, true]);
247    /// // create a new Coalescer that targets creating 1000 row batches
248    /// let mut coalescer = BatchCoalescer::new(batch1.schema(), 1000);
249    /// coalescer.push_batch_with_filter(batch1, &filter);
250    /// coalescer.push_batch_with_filter(batch2, &filter);
251    /// // finsh and retrieve the created batch
252    /// coalescer.finish_buffered_batch().unwrap();
253    /// let completed_batch = coalescer.next_completed_batch().unwrap();
254    /// // filtered out 2 and 5:
255    /// let expected_batch = record_batch!(("a", Int32, [1, 3, 4, 6])).unwrap();
256    /// assert_eq!(completed_batch, expected_batch);
257    /// ```
258    pub fn push_batch_with_filter(
259        &mut self,
260        batch: RecordBatch,
261        filter: &BooleanArray,
262    ) -> Result<(), ArrowError> {
263        self.push_batch_with_filtered_columns(batch, filter)
264    }
265
266    /// Push a batch into the Coalescer after applying a set of indices
267    /// This is semantically equivalent of calling [`Self::push_batch`]
268    /// with the results from  [`take_record_batch`]
269    ///
270    /// # Example
271    /// ```
272    /// # use arrow_array::{record_batch, UInt64Array};
273    /// # use arrow_select::coalesce::BatchCoalescer;
274    /// let batch1 = record_batch!(("a", Int32, [0, 0, 0])).unwrap();
275    /// let batch2 = record_batch!(("a", Int32, [1, 1, 4, 5, 1, 4])).unwrap();
276    /// // Sorted indices to create a sorted output, this can be obtained with
277    /// // `arrow-ord`'s sort_to_indices operation
278    /// let indices = UInt64Array::from(vec![0, 1, 4, 2, 5, 3]);
279    /// // create a new Coalescer that targets creating 1000 row batches
280    /// let mut coalescer = BatchCoalescer::new(batch1.schema(), 1000);
281    /// coalescer.push_batch(batch1);
282    /// coalescer.push_batch_with_indices(batch2, &indices);
283    /// // finsh and retrieve the created batch
284    /// coalescer.finish_buffered_batch().unwrap();
285    /// let completed_batch = coalescer.next_completed_batch().unwrap();
286    /// let expected_batch = record_batch!(("a", Int32, [0, 0, 0, 1, 1, 1, 4, 4, 5])).unwrap();
287    /// assert_eq!(completed_batch, expected_batch);
288    /// ```
289    pub fn push_batch_with_indices(
290        &mut self,
291        batch: RecordBatch,
292        indices: &dyn Array,
293    ) -> Result<(), ArrowError> {
294        // todo: optimize this to avoid materializing (copying the results of take indices to a new batch)
295        let taken_batch = take_record_batch(&batch, indices)?;
296        self.push_batch(taken_batch)
297    }
298
299    /// Push all the rows from `batch` into the Coalescer
300    ///
301    /// When buffered data plus incoming rows reach `target_batch_size` ,
302    /// completed batches are generated eagerly and can be retrieved via
303    /// [`Self::next_completed_batch()`].
304    /// Output batches contain exactly `target_batch_size` rows, so the tail of
305    /// the input batch may remain buffered.
306    /// Remaining partial data either waits for future input batches or can be
307    /// materialized immediately by calling [`Self::finish_buffered_batch()`].
308    ///
309    /// # Example
310    /// ```
311    /// # use arrow_array::record_batch;
312    /// # use arrow_select::coalesce::BatchCoalescer;
313    /// let batch1 = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
314    /// let batch2 = record_batch!(("a", Int32, [4, 5, 6])).unwrap();
315    /// // create a new Coalescer that targets creating 1000 row batches
316    /// let mut coalescer = BatchCoalescer::new(batch1.schema(), 1000);
317    /// coalescer.push_batch(batch1);
318    /// coalescer.push_batch(batch2);
319    /// // finsh and retrieve the created batch
320    /// coalescer.finish_buffered_batch().unwrap();
321    /// let completed_batch = coalescer.next_completed_batch().unwrap();
322    /// let expected_batch = record_batch!(("a", Int32, [1, 2, 3, 4, 5, 6])).unwrap();
323    /// assert_eq!(completed_batch, expected_batch);
324    /// ```
325    pub fn push_batch(&mut self, batch: RecordBatch) -> Result<(), ArrowError> {
326        // Large batch bypass optimization:
327        // When biggest_coalesce_batch_size is configured and a batch exceeds this limit,
328        // we can avoid expensive split-and-merge operations by passing it through directly.
329        //
330        // IMPORTANT: This optimization is OPTIONAL and only active when biggest_coalesce_batch_size
331        // is explicitly set via with_biggest_coalesce_batch_size(Some(limit)).
332        // If not set (None), ALL batches follow normal coalescing behavior regardless of size.
333
334        // =============================================================================
335        // CASE 1: No buffer + large batch → Direct bypass
336        // =============================================================================
337        // Example scenario (target_batch_size=1000, biggest_coalesce_batch_size=Some(500)):
338        // Input sequence: [600, 1200, 300]
339        //
340        // With biggest_coalesce_batch_size=Some(500) (optimization enabled):
341        //   600 → large batch detected! buffered_rows=0 → Case 1: direct bypass
342        //        → output: [600] (bypass, preserves large batch)
343        //   1200 → large batch detected! buffered_rows=0 → Case 1: direct bypass
344        //         → output: [1200] (bypass, preserves large batch)
345        //   300 → normal batch, buffer: [300]
346        //   Result: [600], [1200], [300] - large batches preserved, mixed sizes
347
348        // =============================================================================
349        // CASE 2: Buffer too large + large batch → Flush first, then bypass
350        // =============================================================================
351        // This case prevents creating extremely large merged batches that would
352        // significantly exceed both target_batch_size and biggest_coalesce_batch_size.
353        //
354        // Example 1: Buffer exceeds limit before large batch arrives
355        // target_batch_size=1000, biggest_coalesce_batch_size=Some(400)
356        // Input: [350, 200, 800]
357        //
358        // Step 1: push_batch([350])
359        //   → batch_size=350 <= 400, normal path
360        //   → buffer: [350], buffered_rows=350
361        //
362        // Step 2: push_batch([200])
363        //   → batch_size=200 <= 400, normal path
364        //   → buffer: [350, 200], buffered_rows=550
365        //
366        // Step 3: push_batch([800])
367        //   → batch_size=800 > 400, large batch path
368        //   → buffered_rows=550 > 400 → Case 2: flush first
369        //   → flush: output [550] (combined [350, 200])
370        //   → then bypass: output [800]
371        //   Result: [550], [800] - buffer flushed to prevent oversized merge
372        //
373        // Example 2: Multiple small batches accumulate before large batch
374        // target_batch_size=1000, biggest_coalesce_batch_size=Some(300)
375        // Input: [150, 100, 80, 900]
376        //
377        // Step 1-3: Accumulate small batches
378        //   150 → buffer: [150], buffered_rows=150
379        //   100 → buffer: [150, 100], buffered_rows=250
380        //   80  → buffer: [150, 100, 80], buffered_rows=330
381        //
382        // Step 4: push_batch([900])
383        //   → batch_size=900 > 300, large batch path
384        //   → buffered_rows=330 > 300 → Case 2: flush first
385        //   → flush: output [330] (combined [150, 100, 80])
386        //   → then bypass: output [900]
387        //   Result: [330], [900] - prevents merge into [1230] which would be too large
388
389        // =============================================================================
390        // CASE 3: Small buffer + large batch → Normal coalescing (no bypass)
391        // =============================================================================
392        // When buffer is small enough, we still merge to maintain efficiency
393        // Example: target_batch_size=1000, biggest_coalesce_batch_size=Some(500)
394        // Input: [300, 1200]
395        //
396        // Step 1: push_batch([300])
397        //   → batch_size=300 <= 500, normal path
398        //   → buffer: [300], buffered_rows=300
399        //
400        // Step 2: push_batch([1200])
401        //   → batch_size=1200 > 500, large batch path
402        //   → buffered_rows=300 <= 500 → Case 3: normal merge
403        //   → buffer: [300, 1200] (1500 total)
404        //   → 1500 > target_batch_size → split: output [1000], buffer [500]
405        //   Result: [1000], [500] - normal split/merge behavior maintained
406
407        // =============================================================================
408        // Comparison: Default vs Optimized Behavior
409        // =============================================================================
410        // target_batch_size=1000, biggest_coalesce_batch_size=Some(500)
411        // Input: [600, 1200, 300]
412        //
413        // DEFAULT BEHAVIOR (biggest_coalesce_batch_size=None):
414        //   600 → buffer: [600]
415        //   1200 → buffer: [600, 1200] (1800 rows total)
416        //         → split: output [1000 rows], buffer [800 rows remaining]
417        //   300 → buffer: [800, 300] (1100 rows total)
418        //        → split: output [1000 rows], buffer [100 rows remaining]
419        //   Result: [1000], [1000], [100] - all outputs respect target_batch_size
420        //
421        // OPTIMIZED BEHAVIOR (biggest_coalesce_batch_size=Some(500)):
422        //   600 → Case 1: direct bypass → output: [600]
423        //   1200 → Case 1: direct bypass → output: [1200]
424        //   300 → normal path → buffer: [300]
425        //   Result: [600], [1200], [300] - large batches preserved
426
427        // =============================================================================
428        // Benefits and Trade-offs
429        // =============================================================================
430        // Benefits of the optimization:
431        // - Large batches stay intact (better for downstream vectorized processing)
432        // - Fewer split/merge operations (better CPU performance)
433        // - More predictable memory usage patterns
434        // - Maintains streaming efficiency while preserving batch boundaries
435        //
436        // Trade-offs:
437        // - Output batch sizes become variable (not always target_batch_size)
438        // - May produce smaller partial batches when flushing before large batches
439        // - Requires tuning biggest_coalesce_batch_size parameter for optimal performance
440
441        // TODO, for unsorted batches, we may can filter all large batches, and coalesce all
442        // small batches together?
443
444        let batch_size = batch.num_rows();
445
446        // Fast path: skip empty batches
447        if batch_size == 0 {
448            return Ok(());
449        }
450
451        // Large batch optimization: bypass coalescing for oversized batches
452        if let Some(limit) = self.biggest_coalesce_batch_size
453            && batch_size > limit
454        {
455            // Case 1: No buffered data - emit large batch directly
456            // Example: [] + [1200] → output [1200], buffer []
457            if self.buffered_rows == 0 {
458                self.completed.push_back(batch);
459                return Ok(());
460            }
461
462            // Case 2: Buffer too large - flush then emit to avoid oversized merge
463            // Example: [850] + [1200] → output [850], then output [1200]
464            // This prevents creating batches much larger than both target_batch_size
465            // and biggest_coalesce_batch_size, which could cause memory issues
466            if self.buffered_rows > limit {
467                self.finish_buffered_batch()?;
468                self.completed.push_back(batch);
469                return Ok(());
470            }
471
472            // Case 3: Small buffer - proceed with normal coalescing
473            // Example: [300] + [1200] → split and merge normally
474            // This ensures small batches still get properly coalesced
475            // while allowing some controlled growth beyond the limit
476        }
477
478        let (_schema, arrays, mut num_rows) = batch.into_parts();
479
480        // Validate column count matches the expected schema
481        if arrays.len() != self.in_progress_arrays.len() {
482            return Err(ArrowError::InvalidArgumentError(format!(
483                "Batch has {} columns but BatchCoalescer expects {}",
484                arrays.len(),
485                self.in_progress_arrays.len()
486            )));
487        }
488        self.in_progress_arrays
489            .iter_mut()
490            .zip(arrays)
491            .for_each(|(in_progress, array)| {
492                in_progress.set_source(Some(array));
493            });
494
495        // If pushing this batch would exceed the target batch size,
496        // finish the current batch and start a new one
497        let mut offset = 0;
498        while num_rows > (self.target_batch_size - self.buffered_rows) {
499            let remaining_rows = self.target_batch_size - self.buffered_rows;
500            debug_assert!(remaining_rows > 0);
501
502            // Copy remaining_rows from each array
503            for in_progress in self.in_progress_arrays.iter_mut() {
504                in_progress.copy_rows(offset, remaining_rows)?;
505            }
506
507            self.buffered_rows += remaining_rows;
508            offset += remaining_rows;
509            num_rows -= remaining_rows;
510
511            self.finish_buffered_batch()?;
512        }
513
514        // Add any the remaining rows to the buffer
515        self.buffered_rows += num_rows;
516        if num_rows > 0 {
517            for in_progress in self.in_progress_arrays.iter_mut() {
518                in_progress.copy_rows(offset, num_rows)?;
519            }
520        }
521
522        // If we have reached the target batch size, finalize the buffered batch
523        if self.buffered_rows >= self.target_batch_size {
524            self.finish_buffered_batch()?;
525        }
526
527        // clear in progress sources (to allow the memory to be freed)
528        for in_progress in self.in_progress_arrays.iter_mut() {
529            in_progress.set_source(None);
530        }
531
532        Ok(())
533    }
534
535    /// Returns the number of buffered rows
536    pub fn get_buffered_rows(&self) -> usize {
537        self.buffered_rows
538    }
539
540    /// Concatenates any buffered batches into a single `RecordBatch` and
541    /// clears any output buffers
542    ///
543    /// Normally this is called when the input stream is exhausted, and
544    /// we want to finalize the last batch of rows.
545    ///
546    /// See [`Self::next_completed_batch()`] for the completed batches.
547    pub fn finish_buffered_batch(&mut self) -> Result<(), ArrowError> {
548        if self.buffered_rows == 0 {
549            return Ok(());
550        }
551        let new_arrays = self
552            .in_progress_arrays
553            .iter_mut()
554            .map(|array| array.finish())
555            .collect::<Result<Vec<_>, ArrowError>>()?;
556
557        for (array, field) in new_arrays.iter().zip(self.schema.fields().iter()) {
558            debug_assert_eq!(array.data_type(), field.data_type());
559            debug_assert_eq!(array.len(), self.buffered_rows);
560        }
561
562        // SAFETY: each array was created of the correct type and length.
563        let batch = unsafe {
564            RecordBatch::new_unchecked(Arc::clone(&self.schema), new_arrays, self.buffered_rows)
565        };
566
567        self.buffered_rows = 0;
568        self.completed.push_back(batch);
569        Ok(())
570    }
571
572    /// Returns true if there is any buffered data
573    pub fn is_empty(&self) -> bool {
574        self.buffered_rows == 0 && self.completed.is_empty()
575    }
576
577    /// Returns true if there are any completed batches
578    pub fn has_completed_batch(&self) -> bool {
579        !self.completed.is_empty()
580    }
581
582    /// Removes and returns the next completed batch, if any.
583    pub fn next_completed_batch(&mut self) -> Option<RecordBatch> {
584        self.completed.pop_front()
585    }
586
587    /// Returns the number of bytes used by this data structure.
588    pub fn size(&self) -> usize {
589        self.in_progress_arrays.capacity() * size_of::<Box<dyn InProgressArray>>()
590            + self
591                .in_progress_arrays
592                .iter()
593                .map(|array| array.size())
594                .sum::<usize>()
595            + self.completed.capacity() * size_of::<RecordBatch>()
596            + self
597                .completed
598                .iter()
599                .map(|batch| batch.get_array_memory_size())
600                .sum::<usize>()
601    }
602}
603
604impl BatchCoalescer {
605    fn filter_predicate_for_batch(
606        batch: &RecordBatch,
607        filter: &BooleanArray,
608        selected_count: usize,
609    ) -> FilterPredicate {
610        let mut filter_builder = FilterBuilder::new_with_count(filter, selected_count);
611        if batch.num_columns() > 1
612            || (batch.num_columns() > 0
613                && FilterBuilder::is_optimize_beneficial(batch.schema_ref().field(0).data_type()))
614        {
615            filter_builder = filter_builder.optimize();
616        }
617        filter_builder.build()
618    }
619
620    fn push_batch_with_filtered_columns(
621        &mut self,
622        batch: RecordBatch,
623        filter: &BooleanArray,
624    ) -> Result<(), ArrowError> {
625        let filter_len = filter.len();
626        let batch_num_rows = batch.num_rows();
627        let batch_num_columns = batch.num_columns();
628
629        if filter_len > batch_num_rows {
630            return Err(ArrowError::InvalidArgumentError(format!(
631                "Filter predicate of length {} is larger than target array of length {}",
632                filter_len, batch_num_rows
633            )));
634        }
635
636        let selected_count = filter.true_count();
637        if selected_count == 0 {
638            return Ok(());
639        }
640
641        if selected_count == batch_num_rows && filter_len == batch_num_rows {
642            return self.push_batch(batch);
643        }
644
645        if batch_num_columns != self.in_progress_arrays.len() {
646            return Err(ArrowError::InvalidArgumentError(format!(
647                "Batch has {} columns but BatchCoalescer expects {}",
648                batch_num_columns,
649                self.in_progress_arrays.len()
650            )));
651        }
652
653        let exceeds_coalesce_limit = self
654            .biggest_coalesce_batch_size
655            .is_some_and(|limit| selected_count > limit);
656        let does_not_fit_buffer = selected_count > self.target_batch_size - self.buffered_rows;
657        let should_materialize_filter = exceeds_coalesce_limit
658            || self.has_non_specialized_filter_columns
659            || does_not_fit_buffer
660            || !should_use_sparse_filter_copy(filter_len, selected_count);
661
662        if should_materialize_filter {
663            // Use materialized filtering when sparse per-column copying is unavailable.
664            let predicate = Self::filter_predicate_for_batch(&batch, filter, selected_count);
665            let filtered_batch = predicate.filter_record_batch(&batch)?;
666            return self.push_batch(filtered_batch);
667        }
668
669        let predicate = Self::filter_predicate_for_batch(&batch, filter, selected_count);
670        let (_schema, arrays, _num_rows) = batch.into_parts();
671
672        for (in_progress, array) in self.in_progress_arrays.iter_mut().zip(arrays) {
673            in_progress.copy_rows_by_filter_from(array, &predicate)?;
674        }
675
676        self.buffered_rows += selected_count;
677        if self.buffered_rows >= self.target_batch_size {
678            self.finish_buffered_batch()?;
679        }
680
681        Ok(())
682    }
683}
684
685/// Return a new `InProgressArray` for the given data type
686fn create_in_progress_array(data_type: &DataType, batch_size: usize) -> Box<dyn InProgressArray> {
687    macro_rules! instantiate_primitive {
688        ($t:ty) => {
689            Box::new(InProgressPrimitiveArray::<$t>::new(
690                batch_size,
691                data_type.clone(),
692            ))
693        };
694    }
695
696    downcast_primitive! {
697        // Instantiate InProgressPrimitiveArray for each primitive type
698        data_type => (instantiate_primitive),
699        DataType::Utf8View => Box::new(InProgressByteViewArray::<StringViewType>::new(batch_size)),
700        DataType::BinaryView => {
701            Box::new(InProgressByteViewArray::<BinaryViewType>::new(batch_size))
702        }
703        _ => Box::new(GenericInProgressArray::new()),
704    }
705}
706
707/// Incrementally builds up arrays
708///
709/// [`GenericInProgressArray`] is the default implementation that buffers
710/// arrays, uses other kernels, and concatenates them when finished.
711///
712/// Some types have specialized, faster implementations (e.g.,
713/// [`StringViewArray`], etc.).
714///
715/// [`StringViewArray`]: arrow_array::StringViewArray
716trait InProgressArray: std::fmt::Debug + Send + Sync {
717    /// Set the source array.
718    ///
719    /// Calls to [`Self::copy_rows`] will copy rows from this array into the
720    /// current in-progress array
721    fn set_source(&mut self, source: Option<ArrayRef>);
722
723    /// Copy rows from the current source array into the in-progress array
724    ///
725    /// Note: The source array is set by [`Self::set_source`].
726    ///
727    /// Return an error if the source array is not set
728    fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError>;
729
730    /// Copy rows selected by `filter` from the current source array.
731    ///
732    /// The default implementation calls [`Self::copy_rows_by_selection`]
733    fn copy_rows_by_filter(&mut self, filter: &FilterPredicate) -> Result<(), ArrowError> {
734        self.copy_rows_by_selection(filter.selection())
735    }
736
737    /// Copy rows selected by a [`FilterPredicate`] from `source`.
738    ///
739    /// Unlike the other copy methods, the source array is passed in directly,
740    /// which allows implementations more flexibility. The default
741    /// implementation simply sets `source` via [`Self::set_source`] and then
742    /// calls [`Self::copy_rows_by_filter`].
743    fn copy_rows_by_filter_from(
744        &mut self,
745        source: ArrayRef,
746        filter: &FilterPredicate,
747    ) -> Result<(), ArrowError> {
748        self.set_source(Some(source));
749        let result = self.copy_rows_by_filter(filter);
750        self.set_source(None);
751        result
752    }
753
754    /// Copy rows described by a [`FilterSelection`] from the current source array.
755    ///
756    /// You typically get a [`FilterSelection`] from [`FilterPredicate::selection`].
757    ///
758    /// Note: The source array is set by [`Self::set_source`].
759    fn copy_rows_by_selection(&mut self, selection: FilterSelection<'_>) -> Result<(), ArrowError> {
760        match selection {
761            FilterSelection::None => Ok(()),
762            FilterSelection::All { len } => self.copy_rows(0, len),
763            FilterSelection::Slices(slices) => {
764                slices.try_for_each(|(start, end)| self.copy_rows(start, end - start))
765            }
766            FilterSelection::Indices(indices) => indices.try_for_each(|idx| self.copy_rows(idx, 1)),
767        }
768    }
769
770    /// Finish the currently in-progress array and return it as an `ArrayRef`
771    fn finish(&mut self) -> Result<ArrayRef, ArrowError>;
772
773    /// Get the number of bytes this array is using
774    fn size(&self) -> usize;
775}
776
777#[cfg(test)]
778mod tests {
779    use super::*;
780    use crate::concat::concat_batches;
781    use crate::filter::filter_record_batch;
782    use arrow_array::builder::StringViewBuilder;
783    use arrow_array::cast::AsArray;
784    use arrow_array::types::Int32Type;
785    use arrow_array::{
786        BinaryViewArray, Int32Array, Int64Array, RecordBatchOptions, StringArray, StringViewArray,
787        TimestampNanosecondArray, UInt32Array, UInt64Array, make_array,
788    };
789    use arrow_buffer::BooleanBufferBuilder;
790    use arrow_schema::{DataType, Field, Schema};
791    use rand::{RngExt, SeedableRng};
792    use std::ops::Range;
793
794    #[test]
795    fn test_coalesce() {
796        let batch = uint32_batch(0..8);
797        Test::new("coalesce")
798            .with_batches(std::iter::repeat_n(batch, 10))
799            // expected output is exactly 21 rows (except for the final batch)
800            .with_batch_size(21)
801            .with_expected_output_sizes(vec![21, 21, 21, 17])
802            .run();
803    }
804
805    #[test]
806    fn test_coalesce_one_by_one() {
807        let batch = uint32_batch(0..1); // single row input
808        Test::new("coalesce_one_by_one")
809            .with_batches(std::iter::repeat_n(batch, 97))
810            // expected output is exactly 20 rows (except for the final batch)
811            .with_batch_size(20)
812            .with_expected_output_sizes(vec![20, 20, 20, 20, 17])
813            .run();
814    }
815
816    #[test]
817    fn test_coalesce_empty() {
818        let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]));
819
820        Test::new("coalesce_empty")
821            .with_batches(vec![])
822            .with_schema(schema)
823            .with_batch_size(21)
824            .with_expected_output_sizes(vec![])
825            .run();
826    }
827
828    #[test]
829    fn test_sparse_filter_copy_threshold() {
830        assert!(should_use_sparse_filter_copy(8192, 8));
831        assert!(should_use_sparse_filter_copy(8192, 81));
832        assert!(!should_use_sparse_filter_copy(8192, 819));
833        assert!(!should_use_sparse_filter_copy(8192, 6553));
834    }
835
836    #[test]
837    fn test_single_large_batch_greater_than_target() {
838        // test a single large batch
839        let batch = uint32_batch(0..4096);
840        Test::new("coalesce_single_large_batch_greater_than_target")
841            .with_batch(batch)
842            .with_batch_size(1000)
843            .with_expected_output_sizes(vec![1000, 1000, 1000, 1000, 96])
844            .run();
845    }
846
847    #[test]
848    fn test_single_large_batch_smaller_than_target() {
849        // test a single large batch
850        let batch = uint32_batch(0..4096);
851        Test::new("coalesce_single_large_batch_smaller_than_target")
852            .with_batch(batch)
853            .with_batch_size(8192)
854            .with_expected_output_sizes(vec![4096])
855            .run();
856    }
857
858    #[test]
859    fn test_single_large_batch_equal_to_target() {
860        // test a single large batch
861        let batch = uint32_batch(0..4096);
862        Test::new("coalesce_single_large_batch_equal_to_target")
863            .with_batch(batch)
864            .with_batch_size(4096)
865            .with_expected_output_sizes(vec![4096])
866            .run();
867    }
868
869    #[test]
870    fn test_single_large_batch_equally_divisible_in_target() {
871        // test a single large batch
872        let batch = uint32_batch(0..4096);
873        Test::new("coalesce_single_large_batch_equally_divisible_in_target")
874            .with_batch(batch)
875            .with_batch_size(1024)
876            .with_expected_output_sizes(vec![1024, 1024, 1024, 1024])
877            .run();
878    }
879
880    #[test]
881    fn test_empty_schema() {
882        let schema = Schema::empty();
883        let batch = RecordBatch::new_empty(schema.into());
884        Test::new("coalesce_empty_schema")
885            .with_batch(batch)
886            .with_expected_output_sizes(vec![])
887            .run();
888    }
889
890    /// Coalesce multiple batches, 80k rows, with a 0.1% selectivity filter
891    #[test]
892    #[cfg_attr(miri, ignore)] // Takes too long
893    fn test_coalesce_filtered_001() {
894        let mut filter_builder = RandomFilterBuilder {
895            num_rows: 8000,
896            selectivity: 0.001,
897            seed: 0,
898        };
899
900        // add 10 batches of 8000 rows each
901        // 80k rows, selecting 0.1% means 80 rows
902        // not exactly 80 as the rows are random;
903        let mut test = Test::new("coalesce_filtered_001");
904        for _ in 0..10 {
905            test = test
906                .with_batch(multi_column_batch(0..8000))
907                .with_filter(filter_builder.next_filter())
908        }
909        test.with_batch_size(15)
910            .with_expected_output_sizes(vec![15, 15, 15, 13])
911            .run();
912    }
913
914    /// Coalesce multiple batches, 80k rows, with a 1% selectivity filter
915    #[test]
916    #[cfg_attr(miri, ignore)] // Takes too long
917    fn test_coalesce_filtered_01() {
918        let mut filter_builder = RandomFilterBuilder {
919            num_rows: 8000,
920            selectivity: 0.01,
921            seed: 0,
922        };
923
924        // add 10 batches of 8000 rows each
925        // 80k rows, selecting 1% means 800 rows
926        // not exactly 800 as the rows are random;
927        let mut test = Test::new("coalesce_filtered_01");
928        for _ in 0..10 {
929            test = test
930                .with_batch(multi_column_batch(0..8000))
931                .with_filter(filter_builder.next_filter())
932        }
933        test.with_batch_size(128)
934            .with_expected_output_sizes(vec![128, 128, 128, 128, 128, 128, 15])
935            .run();
936    }
937
938    /// Coalesce multiple batches, 80k rows, with a 10% selectivity filter
939    #[test]
940    #[cfg_attr(miri, ignore)] // Takes too long
941    fn test_coalesce_filtered_10() {
942        let mut filter_builder = RandomFilterBuilder {
943            num_rows: 8000,
944            selectivity: 0.1,
945            seed: 0,
946        };
947
948        // add 10 batches of 8000 rows each
949        // 80k rows, selecting 10% means 8000 rows
950        // not exactly 800 as the rows are random;
951        let mut test = Test::new("coalesce_filtered_10");
952        for _ in 0..10 {
953            test = test
954                .with_batch(multi_column_batch(0..8000))
955                .with_filter(filter_builder.next_filter())
956        }
957        test.with_batch_size(1024)
958            .with_expected_output_sizes(vec![1024, 1024, 1024, 1024, 1024, 1024, 1024, 840])
959            .run();
960    }
961
962    /// Coalesce multiple batches, 8k rows, with a 90% selectivity filter
963    #[test]
964    #[cfg_attr(miri, ignore)] // Takes too long
965    fn test_coalesce_filtered_90() {
966        let mut filter_builder = RandomFilterBuilder {
967            num_rows: 800,
968            selectivity: 0.90,
969            seed: 0,
970        };
971
972        // add 10 batches of 800 rows each
973        // 8k rows, selecting 99% means 7200 rows
974        // not exactly 7200 as the rows are random;
975        let mut test = Test::new("coalesce_filtered_90");
976        for _ in 0..10 {
977            test = test
978                .with_batch(multi_column_batch(0..800))
979                .with_filter(filter_builder.next_filter())
980        }
981        test.with_batch_size(1024)
982            .with_expected_output_sizes(vec![1024, 1024, 1024, 1024, 1024, 1024, 1024, 13])
983            .run();
984    }
985
986    /// Coalesce multiple batches, 8k rows, with mixed filers, including 100%
987    #[test]
988    #[cfg_attr(miri, ignore)] // Takes too long
989    fn test_coalesce_filtered_mixed() {
990        let mut filter_builder = RandomFilterBuilder {
991            num_rows: 800,
992            selectivity: 0.90,
993            seed: 0,
994        };
995
996        let mut test = Test::new("coalesce_filtered_mixed");
997        for _ in 0..3 {
998            // also add in a batch that selects almost all rows and when
999            // sliced will have some batches that are entirely used
1000            let mut all_filter_builder = BooleanBufferBuilder::new(1000);
1001            all_filter_builder.append_n(500, true);
1002            all_filter_builder.append_n(1, false);
1003            all_filter_builder.append_n(499, false);
1004            let all_filter = all_filter_builder.build();
1005
1006            test = test
1007                .with_batch(multi_column_batch(0..1000))
1008                .with_filter(BooleanArray::from(all_filter))
1009                .with_batch(multi_column_batch(0..800))
1010                .with_filter(filter_builder.next_filter());
1011            // decrease selectivity
1012            filter_builder.selectivity *= 0.6;
1013        }
1014
1015        // use a small batch size to ensure the filter is appended in slices
1016        // and some of those slides will select the entire thing.
1017        test.with_batch_size(250)
1018            .with_expected_output_sizes(vec![
1019                250, 250, 250, 250, 250, 250, 250, 250, 250, 250, 250, 179,
1020            ])
1021            .run();
1022    }
1023
1024    #[test]
1025    fn test_coalesce_non_null() {
1026        Test::new("coalesce_non_null")
1027            // 4040 rows of unit32
1028            .with_batch(uint32_batch_non_null(0..3000))
1029            .with_batch(uint32_batch_non_null(0..1040))
1030            .with_batch_size(1024)
1031            .with_expected_output_sizes(vec![1024, 1024, 1024, 968])
1032            .run();
1033    }
1034    #[test]
1035    #[cfg_attr(miri, ignore)] // Takes too long
1036    fn test_utf8_split() {
1037        Test::new("coalesce_utf8")
1038            // 4040 rows of utf8 strings in total, split into batches of 1024
1039            .with_batch(utf8_batch(0..3000))
1040            .with_batch(utf8_batch(0..1040))
1041            .with_batch_size(1024)
1042            .with_expected_output_sizes(vec![1024, 1024, 1024, 968])
1043            .run();
1044    }
1045
1046    #[test]
1047    fn test_string_view_no_views() {
1048        let output_batches = Test::new("coalesce_string_view_no_views")
1049            // both input batches have no views, so no need to compact
1050            .with_batch(stringview_batch([Some("foo"), Some("bar")]))
1051            .with_batch(stringview_batch([Some("baz"), Some("qux")]))
1052            .with_expected_output_sizes(vec![4])
1053            .run();
1054
1055        expect_buffer_layout(
1056            col_as_string_view("c0", output_batches.first().unwrap()),
1057            vec![],
1058        );
1059    }
1060
1061    #[test]
1062    fn test_string_view_batch_small_no_compact() {
1063        // view with only short strings (no buffers) --> no need to compact
1064        let batch = stringview_batch_repeated(1000, [Some("a"), Some("b"), Some("c")]);
1065        let output_batches = Test::new("coalesce_string_view_batch_small_no_compact")
1066            .with_batch(batch.clone())
1067            .with_expected_output_sizes(vec![1000])
1068            .run();
1069
1070        let array = col_as_string_view("c0", &batch);
1071        let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1072        assert_eq!(array.data_buffers().len(), 0);
1073        assert_eq!(array.data_buffers().len(), gc_array.data_buffers().len()); // no compaction
1074
1075        expect_buffer_layout(gc_array, vec![]);
1076    }
1077
1078    #[test]
1079    #[cfg_attr(miri, ignore)] // Takes too long
1080    fn test_string_view_batch_large_no_compact() {
1081        // view with large strings (has buffers) but full --> no need to compact
1082        let batch = stringview_batch_repeated(1000, [Some("This string is longer than 12 bytes")]);
1083        let output_batches = Test::new("coalesce_string_view_batch_large_no_compact")
1084            .with_batch(batch.clone())
1085            .with_batch_size(1000)
1086            .with_expected_output_sizes(vec![1000])
1087            .run();
1088
1089        let array = col_as_string_view("c0", &batch);
1090        let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1091        assert_eq!(array.data_buffers().len(), 5);
1092        assert_eq!(array.data_buffers().len(), gc_array.data_buffers().len()); // no compaction
1093
1094        expect_buffer_layout(
1095            gc_array,
1096            vec![
1097                ExpectedLayout {
1098                    len: 8190,
1099                    capacity: 8192,
1100                },
1101                ExpectedLayout {
1102                    len: 8190,
1103                    capacity: 8192,
1104                },
1105                ExpectedLayout {
1106                    len: 8190,
1107                    capacity: 8192,
1108                },
1109                ExpectedLayout {
1110                    len: 8190,
1111                    capacity: 8192,
1112                },
1113                ExpectedLayout {
1114                    len: 2240,
1115                    capacity: 8192,
1116                },
1117            ],
1118        );
1119    }
1120
1121    #[test]
1122    fn test_string_view_batch_small_with_buffers_no_compact() {
1123        // view with buffers but only short views
1124        let short_strings = std::iter::repeat(Some("SmallString"));
1125        let long_strings = std::iter::once(Some("This string is longer than 12 bytes"));
1126        // 20 short strings, then a long ones
1127        let values = short_strings.take(20).chain(long_strings);
1128        let batch = stringview_batch_repeated(1000, values)
1129            // take only 10 short strings (no long ones)
1130            .slice(5, 10);
1131        let output_batches = Test::new("coalesce_string_view_batch_small_with_buffers_no_compact")
1132            .with_batch(batch.clone())
1133            .with_batch_size(1000)
1134            .with_expected_output_sizes(vec![10])
1135            .run();
1136
1137        let array = col_as_string_view("c0", &batch);
1138        let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1139        assert_eq!(array.data_buffers().len(), 1); // input has one buffer
1140        assert_eq!(gc_array.data_buffers().len(), 0); // output has no buffers as only short strings
1141    }
1142
1143    #[test]
1144    fn test_string_view_batch_large_slice_compact() {
1145        // view with large strings (has buffers) and only partially used  --> no need to compact
1146        let batch = stringview_batch_repeated(1000, [Some("This string is longer than 12 bytes")])
1147            // slice only 22 rows, so most of the buffer is not used
1148            .slice(11, 22);
1149
1150        let output_batches = Test::new("coalesce_string_view_batch_large_slice_compact")
1151            .with_batch(batch.clone())
1152            .with_batch_size(1000)
1153            .with_expected_output_sizes(vec![22])
1154            .run();
1155
1156        let array = col_as_string_view("c0", &batch);
1157        let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1158        assert_eq!(array.data_buffers().len(), 5);
1159
1160        expect_buffer_layout(
1161            gc_array,
1162            vec![ExpectedLayout {
1163                len: 770,
1164                capacity: 8192,
1165            }],
1166        );
1167    }
1168
1169    #[test]
1170    #[cfg_attr(miri, ignore)] // Takes too long
1171    fn test_string_view_mixed() {
1172        let large_view_batch =
1173            stringview_batch_repeated(1000, [Some("This string is longer than 12 bytes")]);
1174        let small_view_batch = stringview_batch_repeated(1000, [Some("SmallString")]);
1175        let mixed_batch = stringview_batch_repeated(
1176            1000,
1177            [Some("This string is longer than 12 bytes"), Some("Small")],
1178        );
1179        let mixed_batch_nulls = stringview_batch_repeated(
1180            1000,
1181            [
1182                Some("This string is longer than 12 bytes"),
1183                Some("Small"),
1184                None,
1185            ],
1186        );
1187
1188        // Several batches with mixed inline / non inline
1189        // 4k rows in
1190        let output_batches = Test::new("coalesce_string_view_mixed")
1191            .with_batch(large_view_batch.clone())
1192            .with_batch(small_view_batch)
1193            // this batch needs to be compacted (less than 1/2 full)
1194            .with_batch(large_view_batch.slice(10, 20))
1195            .with_batch(mixed_batch_nulls)
1196            // this batch needs to be compacted (less than 1/2 full)
1197            .with_batch(large_view_batch.slice(10, 20))
1198            .with_batch(mixed_batch)
1199            .with_expected_output_sizes(vec![1024, 1024, 1024, 968])
1200            .run();
1201
1202        expect_buffer_layout(
1203            col_as_string_view("c0", output_batches.first().unwrap()),
1204            vec![
1205                ExpectedLayout {
1206                    len: 8190,
1207                    capacity: 8192,
1208                },
1209                ExpectedLayout {
1210                    len: 8190,
1211                    capacity: 8192,
1212                },
1213                ExpectedLayout {
1214                    len: 8190,
1215                    capacity: 8192,
1216                },
1217                ExpectedLayout {
1218                    len: 8190,
1219                    capacity: 8192,
1220                },
1221                ExpectedLayout {
1222                    len: 2240,
1223                    capacity: 8192,
1224                },
1225            ],
1226        );
1227    }
1228
1229    #[test]
1230    #[cfg_attr(miri, ignore)] // Takes too long
1231    fn test_string_view_many_small_compact() {
1232        // 200 rows alternating long (28) and short (≤12) strings.
1233        // Only the 100 long strings go into data buffers: 100 × 28 = 2800.
1234        let batch = stringview_batch_repeated(
1235            200,
1236            [Some("This string is 28 bytes long"), Some("small string")],
1237        );
1238        let output_batches = Test::new("coalesce_string_view_many_small_compact")
1239            // First allocated buffer is 8kb.
1240            // Appending 10 batches of 2800 bytes will use 2800 * 10 = 14kb (8kb, an 16kb and 32kbkb)
1241            .with_batch(batch.clone())
1242            .with_batch(batch.clone())
1243            .with_batch(batch.clone())
1244            .with_batch(batch.clone())
1245            .with_batch(batch.clone())
1246            .with_batch(batch.clone())
1247            .with_batch(batch.clone())
1248            .with_batch(batch.clone())
1249            .with_batch(batch.clone())
1250            .with_batch(batch.clone())
1251            .with_batch_size(8000)
1252            .with_expected_output_sizes(vec![2000]) // only 1000 rows total
1253            .run();
1254
1255        // expect a nice even distribution of buffers
1256        expect_buffer_layout(
1257            col_as_string_view("c0", output_batches.first().unwrap()),
1258            vec![
1259                ExpectedLayout {
1260                    len: 8176,
1261                    capacity: 8192,
1262                },
1263                ExpectedLayout {
1264                    len: 16380,
1265                    capacity: 16384,
1266                },
1267                ExpectedLayout {
1268                    len: 3444,
1269                    capacity: 32768,
1270                },
1271            ],
1272        );
1273    }
1274
1275    #[test]
1276    #[cfg_attr(miri, ignore)] // Takes too long
1277    fn test_string_view_many_small_boundary() {
1278        // The strings are designed to exactly fit into buffers that are powers of 2 long
1279        let batch = stringview_batch_repeated(100, [Some("This string is a power of two=32")]);
1280        let output_batches = Test::new("coalesce_string_view_many_small_boundary")
1281            .with_batches(std::iter::repeat_n(batch, 20))
1282            .with_batch_size(900)
1283            .with_expected_output_sizes(vec![900, 900, 200])
1284            .run();
1285
1286        // expect each buffer to be entirely full except the last one
1287        expect_buffer_layout(
1288            col_as_string_view("c0", output_batches.first().unwrap()),
1289            vec![
1290                ExpectedLayout {
1291                    len: 8192,
1292                    capacity: 8192,
1293                },
1294                ExpectedLayout {
1295                    len: 16384,
1296                    capacity: 16384,
1297                },
1298                ExpectedLayout {
1299                    len: 4224,
1300                    capacity: 32768,
1301                },
1302            ],
1303        );
1304    }
1305
1306    #[test]
1307    #[cfg_attr(miri, ignore)] // Takes too long
1308    fn test_string_view_large_small() {
1309        // The strings are 37 bytes long, so each batch has 100 * 28 = 2800 bytes
1310        let mixed_batch = stringview_batch_repeated(
1311            200,
1312            [Some("This string is 28 bytes long"), Some("small string")],
1313        );
1314        // These strings aren't copied, this array has an 8k buffer
1315        let all_large = stringview_batch_repeated(
1316            50,
1317            [Some(
1318                "This buffer has only large strings in it so there are no buffer copies",
1319            )],
1320        );
1321
1322        let output_batches = Test::new("coalesce_string_view_large_small")
1323            // First allocated buffer is 8kb.
1324            // Appending five batches of 2800 bytes will use 2800 * 10 = 28kb (8kb, an 16kb and 32kbkb)
1325            .with_batch(mixed_batch.clone())
1326            .with_batch(mixed_batch.clone())
1327            .with_batch(all_large.clone())
1328            .with_batch(mixed_batch.clone())
1329            .with_batch(all_large.clone())
1330            .with_batch(mixed_batch.clone())
1331            .with_batch(mixed_batch.clone())
1332            .with_batch(all_large.clone())
1333            .with_batch(mixed_batch.clone())
1334            .with_batch(all_large.clone())
1335            .with_batch_size(8000)
1336            .with_expected_output_sizes(vec![1400])
1337            .run();
1338
1339        expect_buffer_layout(
1340            col_as_string_view("c0", output_batches.first().unwrap()),
1341            vec![
1342                ExpectedLayout {
1343                    len: 8190,
1344                    capacity: 8192,
1345                },
1346                ExpectedLayout {
1347                    len: 16366,
1348                    capacity: 16384,
1349                },
1350                ExpectedLayout {
1351                    len: 6244,
1352                    capacity: 32768,
1353                },
1354            ],
1355        );
1356    }
1357
1358    #[test]
1359    #[cfg_attr(miri, ignore)] // Takes too long
1360    fn test_binary_view() {
1361        let values: Vec<Option<&[u8]>> = vec![
1362            Some(b"foo"),
1363            None,
1364            Some(b"A longer string that is more than 12 bytes"),
1365        ];
1366
1367        let binary_view =
1368            BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1369        let batch =
1370            RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1371
1372        Test::new("coalesce_binary_view")
1373            .with_batch(batch.clone())
1374            .with_batch(batch.clone())
1375            .with_batch_size(512)
1376            .with_expected_output_sizes(vec![512, 512, 512, 464])
1377            .run();
1378    }
1379
1380    #[test]
1381    fn test_binary_view_filtered() {
1382        let values: Vec<Option<&[u8]>> = vec![
1383            Some(b"foo"),
1384            None,
1385            Some(b"A longer string that is more than 12 bytes"),
1386        ];
1387
1388        let binary_view =
1389            BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1390        let batch =
1391            RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1392        let filter = sparse_filter(1000);
1393
1394        Test::new("coalesce_binary_view_filtered")
1395            .with_batch(batch.clone())
1396            .with_filter(filter.clone())
1397            .with_batch(batch)
1398            .with_filter(filter)
1399            .with_batch_size(256)
1400            .with_expected_output_sizes(vec![250])
1401            .run();
1402    }
1403
1404    #[test]
1405    fn test_binary_view_filtered_inline() {
1406        let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1407
1408        let binary_view =
1409            BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1410        let batch =
1411            RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1412        let filter = sparse_filter(1000);
1413
1414        Test::new("coalesce_binary_view_filtered_inline")
1415            .with_batch(batch.clone())
1416            .with_filter(filter.clone())
1417            .with_batch(batch)
1418            .with_filter(filter)
1419            .with_batch_size(300)
1420            .with_expected_output_sizes(vec![250])
1421            .run();
1422    }
1423
1424    #[test]
1425    fn test_string_view_filtered_inline() {
1426        let values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1427
1428        let string_view =
1429            StringViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1430        let batch =
1431            RecordBatch::try_from_iter(vec![("c0", Arc::new(string_view) as ArrayRef)]).unwrap();
1432        let filter = sparse_filter(1000);
1433
1434        Test::new("coalesce_string_view_filtered_inline")
1435            .with_batch(batch.clone())
1436            .with_filter(filter.clone())
1437            .with_batch(batch)
1438            .with_filter(filter)
1439            .with_batch_size(300)
1440            .with_expected_output_sizes(vec![250])
1441            .run();
1442    }
1443
1444    #[test]
1445    fn test_mixed_inline_binary_view_filtered() {
1446        let int_values =
1447            Int32Array::from_iter((0..1000).map(|v| if v % 5 == 0 { None } else { Some(v) }));
1448        let float_values = arrow_array::Float64Array::from_iter((0..1000).map(|v| Some(v as f64)));
1449        let binary_values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1450        let binary_view = BinaryViewArray::from_iter(
1451            std::iter::repeat(binary_values.iter()).flatten().take(1000),
1452        );
1453
1454        let batch = RecordBatch::try_from_iter(vec![
1455            ("i", Arc::new(int_values) as ArrayRef),
1456            ("f", Arc::new(float_values) as ArrayRef),
1457            ("b", Arc::new(binary_view) as ArrayRef),
1458        ])
1459        .unwrap();
1460
1461        let filter = sparse_filter(1000);
1462
1463        Test::new("coalesce_mixed_inline_binary_view_filtered")
1464            .with_batch(batch.clone())
1465            .with_filter(filter.clone())
1466            .with_batch(batch)
1467            .with_filter(filter)
1468            .with_batch_size(300)
1469            .with_expected_output_sizes(vec![250])
1470            .run();
1471    }
1472
1473    #[test]
1474    fn test_mixed_inline_string_view_filtered() {
1475        let int_values =
1476            Int32Array::from_iter((0..1000).map(|v| if v % 5 == 0 { None } else { Some(v) }));
1477        let float_values = arrow_array::Float64Array::from_iter((0..1000).map(|v| Some(v as f64)));
1478        let string_values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1479        let string_view = StringViewArray::from_iter(
1480            std::iter::repeat(string_values.iter()).flatten().take(1000),
1481        );
1482
1483        let batch = RecordBatch::try_from_iter(vec![
1484            ("i", Arc::new(int_values) as ArrayRef),
1485            ("f", Arc::new(float_values) as ArrayRef),
1486            ("s", Arc::new(string_view) as ArrayRef),
1487        ])
1488        .unwrap();
1489
1490        let filter = sparse_filter(1000);
1491
1492        Test::new("coalesce_mixed_inline_string_view_filtered")
1493            .with_batch(batch.clone())
1494            .with_filter(filter.clone())
1495            .with_batch(batch)
1496            .with_filter(filter)
1497            .with_batch_size(300)
1498            .with_expected_output_sizes(vec![250])
1499            .run();
1500    }
1501
1502    #[test]
1503    fn test_inline_binary_view_sparse() {
1504        // All-inline binary views (<=12 bytes) + 5% selectivity
1505        let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1506        let binary_view =
1507            BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1508        let batch =
1509            RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1510        let filter = very_sparse_filter(1000);
1511
1512        Test::new("inline_binary_view_sparse")
1513            .with_batch(batch.clone())
1514            .with_filter(filter.clone())
1515            .with_batch(batch)
1516            .with_filter(filter)
1517            .with_batch_size(1024)
1518            .with_expected_output_sizes(vec![100])
1519            .run();
1520    }
1521
1522    #[test]
1523    fn test_inline_string_view_sparse() {
1524        let values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1525        let string_view =
1526            StringViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1527        let batch =
1528            RecordBatch::try_from_iter(vec![("c0", Arc::new(string_view) as ArrayRef)]).unwrap();
1529        let filter = very_sparse_filter(1000);
1530
1531        Test::new("inline_string_view_sparse")
1532            .with_batch(batch.clone())
1533            .with_filter(filter.clone())
1534            .with_batch(batch)
1535            .with_filter(filter)
1536            .with_batch_size(1024)
1537            .with_expected_output_sizes(vec![100])
1538            .run();
1539    }
1540
1541    #[test]
1542    fn test_inline_mixed_sparse() {
1543        // Mixing primitives with inline views exercises materialized `Indices`
1544        // selection for both the primitive and the byte-view.
1545        let int_values =
1546            Int32Array::from_iter((0..1000).map(|v| if v % 5 == 0 { None } else { Some(v) }));
1547        let float_values = arrow_array::Float64Array::from_iter((0..1000).map(|v| Some(v as f64)));
1548        let string_values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1549        let string_view = StringViewArray::from_iter(
1550            std::iter::repeat(string_values.iter()).flatten().take(1000),
1551        );
1552        let binary_values: Vec<Option<&[u8]>> = vec![Some(b"x"), None, Some(b"abcdef")];
1553        let binary_view = BinaryViewArray::from_iter(
1554            std::iter::repeat(binary_values.iter()).flatten().take(1000),
1555        );
1556
1557        let batch = RecordBatch::try_from_iter(vec![
1558            ("i", Arc::new(int_values) as ArrayRef),
1559            ("f", Arc::new(float_values) as ArrayRef),
1560            ("s", Arc::new(string_view) as ArrayRef),
1561            ("b", Arc::new(binary_view) as ArrayRef),
1562        ])
1563        .unwrap();
1564        let filter = very_sparse_filter(1000);
1565
1566        Test::new("inline_mixed_sparse")
1567            .with_batch(batch.clone())
1568            .with_filter(filter.clone())
1569            .with_batch(batch)
1570            .with_filter(filter)
1571            .with_batch_size(1024)
1572            .with_expected_output_sizes(vec![100])
1573            .run();
1574    }
1575
1576    #[test]
1577    fn test_inline_crosses_target_batch_size() {
1578        // Each 1000-row batch selects 50 rows; target_batch_size 100
1579        // exercises intermediate batch flushing batch (not just the final
1580        // flush).
1581        let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1582        let make_batch = || {
1583            let binary_view =
1584                BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1585            RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap()
1586        };
1587        let filter = very_sparse_filter(1000);
1588
1589        Test::new("inline_crosses_target_batch_size")
1590            .with_batch(make_batch())
1591            .with_filter(filter.clone())
1592            .with_batch(make_batch())
1593            .with_filter(filter.clone())
1594            .with_batch(make_batch())
1595            .with_filter(filter)
1596            .with_batch_size(100)
1597            .with_expected_output_sizes(vec![100, 50])
1598            .run();
1599    }
1600
1601    #[test]
1602    fn test_inline_filter_rejects_filter_longer_than_batch() {
1603        let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), Some(b"bar")];
1604        let binary_view = BinaryViewArray::from_iter(values);
1605        let batch =
1606            RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1607        let filter = BooleanArray::from(vec![true, false, true]);
1608
1609        let mut coalescer = BatchCoalescer::new(batch.schema(), 100);
1610        let result = coalescer.push_batch_with_filter(batch, &filter);
1611        assert!(result.is_err());
1612        let err = result.unwrap_err().to_string();
1613        assert!(
1614            err.contains("Filter predicate of length 3 is larger than target array of length 2"),
1615            "unexpected error: {err}"
1616        );
1617    }
1618
1619    #[test]
1620    fn test_mixed_boolean_inline_string_view_filtered() {
1621        let bool_values = BooleanArray::from_iter((0..1000).map(|v| Some(v % 3 == 0)));
1622        let string_values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1623        let string_view = StringViewArray::from_iter(
1624            std::iter::repeat(string_values.iter()).flatten().take(1000),
1625        );
1626
1627        let batch = RecordBatch::try_from_iter(vec![
1628            ("b", Arc::new(bool_values) as ArrayRef),
1629            ("s", Arc::new(string_view) as ArrayRef),
1630        ])
1631        .unwrap();
1632
1633        let filter = sparse_filter(1000);
1634
1635        Test::new("coalesce_mixed_boolean_inline_string_view_filtered")
1636            .with_batch(batch.clone())
1637            .with_filter(filter.clone())
1638            .with_batch(batch)
1639            .with_filter(filter)
1640            .with_batch_size(300)
1641            .with_expected_output_sizes(vec![250])
1642            .run();
1643    }
1644
1645    #[test]
1646    fn test_filter_fast_path_schema_capability() {
1647        let supported = Arc::new(Schema::new(vec![
1648            Field::new("primitive", DataType::UInt32, false),
1649            Field::new("utf8_view", DataType::Utf8View, true),
1650            Field::new("binary_view", DataType::BinaryView, true),
1651        ]));
1652        let coalescer = BatchCoalescer::new(supported, 100);
1653        assert!(!coalescer.has_non_specialized_filter_columns);
1654
1655        let utf8 = Arc::new(Schema::new(vec![Field::new("utf8", DataType::Utf8, true)]));
1656        let coalescer = BatchCoalescer::new(utf8, 100);
1657        assert!(coalescer.has_non_specialized_filter_columns);
1658
1659        let boolean = Arc::new(Schema::new(vec![Field::new(
1660            "boolean",
1661            DataType::Boolean,
1662            true,
1663        )]));
1664        let coalescer = BatchCoalescer::new(boolean, 100);
1665        assert!(coalescer.has_non_specialized_filter_columns);
1666    }
1667
1668    #[derive(Debug, Clone, PartialEq)]
1669    struct ExpectedLayout {
1670        len: usize,
1671        capacity: usize,
1672    }
1673
1674    /// Asserts that the buffer layout of the specified StringViewArray matches the expected layout
1675    fn expect_buffer_layout(array: &StringViewArray, expected: Vec<ExpectedLayout>) {
1676        let actual = array
1677            .data_buffers()
1678            .iter()
1679            .map(|b| ExpectedLayout {
1680                len: b.len(),
1681                capacity: b.capacity(),
1682            })
1683            .collect::<Vec<_>>();
1684
1685        assert_eq!(
1686            actual, expected,
1687            "Expected buffer layout {expected:#?} but got {actual:#?}"
1688        );
1689    }
1690
1691    /// Test for [`BatchCoalescer`]
1692    ///
1693    /// Pushes the input batches to the coalescer and verifies that the resulting
1694    /// batches have the
1695    /// 1. expected number of rows
1696    /// 2. The same results when the batches are filtered using the filter kernel
1697    #[derive(Debug, Clone)]
1698    struct Test {
1699        /// A human readable name to assist in debugging
1700        name: String,
1701        /// Batches to feed to the coalescer.
1702        input_batches: Vec<RecordBatch>,
1703        /// Filters to apply to the corresponding input batches.
1704        ///
1705        /// If there are no filters for the input batches, the batch will be
1706        /// pushed as is.
1707        filters: Vec<BooleanArray>,
1708        /// The schema. If not provided, the first batch's schema is used.
1709        schema: Option<SchemaRef>,
1710        /// Expected output sizes of the resulting batches
1711        expected_output_sizes: Vec<usize>,
1712        /// target batch size (default to 1024)
1713        target_batch_size: usize,
1714    }
1715
1716    impl Default for Test {
1717        fn default() -> Self {
1718            Self {
1719                name: String::new(),
1720                input_batches: vec![],
1721                filters: vec![],
1722                schema: None,
1723                expected_output_sizes: vec![],
1724                target_batch_size: 1024,
1725            }
1726        }
1727    }
1728
1729    impl Test {
1730        fn new(name: impl Into<String>) -> Self {
1731            Self {
1732                name: name.into(),
1733                ..Self::default()
1734            }
1735        }
1736
1737        /// Append the description to the test name
1738        fn with_description(mut self, description: &str) -> Self {
1739            self.name.push_str(": ");
1740            self.name.push_str(description);
1741            self
1742        }
1743
1744        /// Set the target batch size
1745        fn with_batch_size(mut self, target_batch_size: usize) -> Self {
1746            self.target_batch_size = target_batch_size;
1747            self
1748        }
1749
1750        /// Extend the input batches with `batch`
1751        fn with_batch(mut self, batch: RecordBatch) -> Self {
1752            self.input_batches.push(batch);
1753            self
1754        }
1755
1756        /// Extend the filters with `filter`
1757        fn with_filter(mut self, filter: BooleanArray) -> Self {
1758            self.filters.push(filter);
1759            self
1760        }
1761
1762        /// Replaces the input batches with `batches`
1763        fn with_batches(mut self, batches: impl IntoIterator<Item = RecordBatch>) -> Self {
1764            self.input_batches = batches.into_iter().collect();
1765            self
1766        }
1767
1768        /// Specifies the schema for the test
1769        fn with_schema(mut self, schema: SchemaRef) -> Self {
1770            self.schema = Some(schema);
1771            self
1772        }
1773
1774        /// Extends `sizes` to expected output sizes
1775        fn with_expected_output_sizes(mut self, sizes: impl IntoIterator<Item = usize>) -> Self {
1776            self.expected_output_sizes.extend(sizes);
1777            self
1778        }
1779
1780        /// Runs the test -- see documentation on [`Test`] for details
1781        ///
1782        /// Returns the resulting output batches
1783        fn run(self) -> Vec<RecordBatch> {
1784            // Test several permutations of input batches:
1785            // 1. Removing nulls from some batches (test non-null fast paths)
1786            // 2. Empty batches
1787            // 3. One column (from the batch)
1788            let mut extra_tests = vec![];
1789            extra_tests.push(self.clone().make_half_non_nullable());
1790            extra_tests.push(self.clone().insert_empty_batches());
1791            let single_column_tests = self.make_single_column_tests();
1792            for test in single_column_tests {
1793                extra_tests.push(test.clone().make_half_non_nullable());
1794                extra_tests.push(test);
1795            }
1796
1797            // Run original test case first, so any obvious errors are caught
1798            // by an easier to understand test case
1799            let results = self.run_inner();
1800            // Run the extra cases to expand coverage
1801            for extra in extra_tests {
1802                extra.run_inner();
1803            }
1804
1805            results
1806        }
1807
1808        /// Runs the current test instance
1809        fn run_inner(self) -> Vec<RecordBatch> {
1810            let expected_output = self.expected_output();
1811            let schema = self.schema();
1812
1813            let Self {
1814                name,
1815                input_batches,
1816                filters,
1817                schema: _,
1818                target_batch_size,
1819                expected_output_sizes,
1820            } = self;
1821
1822            println!("Running test '{name}'");
1823
1824            let had_input = input_batches.iter().any(|b| b.num_rows() > 0);
1825
1826            let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), target_batch_size);
1827
1828            // feed input batches and filters to the coalescer
1829            let mut filters = filters.into_iter();
1830            for batch in input_batches {
1831                if let Some(filter) = filters.next() {
1832                    coalescer.push_batch_with_filter(batch, &filter).unwrap();
1833                } else {
1834                    coalescer.push_batch(batch).unwrap();
1835                }
1836            }
1837            assert_eq!(schema, coalescer.schema());
1838
1839            if had_input {
1840                assert!(!coalescer.is_empty(), "Coalescer should not be empty");
1841            } else {
1842                assert!(coalescer.is_empty(), "Coalescer should be empty");
1843            }
1844
1845            coalescer.finish_buffered_batch().unwrap();
1846            if had_input {
1847                assert!(
1848                    coalescer.has_completed_batch(),
1849                    "Coalescer should have completed batches"
1850                );
1851            }
1852
1853            let mut output_batches = vec![];
1854            while let Some(batch) = coalescer.next_completed_batch() {
1855                output_batches.push(batch);
1856            }
1857
1858            // make sure we got the expected number of output batches and content
1859            let mut starting_idx = 0;
1860            let actual_output_sizes: Vec<usize> =
1861                output_batches.iter().map(|b| b.num_rows()).collect();
1862            assert_eq!(
1863                expected_output_sizes, actual_output_sizes,
1864                "Unexpected number of rows in output batches\n\
1865                Expected\n{expected_output_sizes:#?}\nActual:{actual_output_sizes:#?}"
1866            );
1867            let iter = expected_output_sizes
1868                .iter()
1869                .zip(output_batches.iter())
1870                .enumerate();
1871
1872            // Verify that the actual contents of each output batch matches the expected output
1873            for (i, (expected_size, batch)) in iter {
1874                // compare the contents of the batch after normalization (using
1875                // `==` compares the underlying memory layout too)
1876                let expected_batch = expected_output.slice(starting_idx, *expected_size);
1877                let expected_batch = normalize_batch(expected_batch);
1878                let batch = normalize_batch(batch.clone());
1879                assert_eq!(
1880                    expected_batch, batch,
1881                    "Unexpected content in batch {i}:\
1882                    \n\nExpected:\n{expected_batch:#?}\n\nActual:\n{batch:#?}"
1883                );
1884                starting_idx += *expected_size;
1885            }
1886            output_batches
1887        }
1888
1889        /// Return the expected output schema. If not overridden by `with_schema`, it
1890        /// returns the schema of the first input batch.
1891        fn schema(&self) -> SchemaRef {
1892            self.schema
1893                .clone()
1894                .unwrap_or_else(|| Arc::clone(&self.input_batches[0].schema()))
1895        }
1896
1897        /// Returns the expected output as a single `RecordBatch`
1898        fn expected_output(&self) -> RecordBatch {
1899            let schema = self.schema();
1900            if self.filters.is_empty() {
1901                return concat_batches(&schema, &self.input_batches).unwrap();
1902            }
1903
1904            let mut filters = self.filters.iter();
1905            let filtered_batches = self
1906                .input_batches
1907                .iter()
1908                .map(|batch| {
1909                    if let Some(filter) = filters.next() {
1910                        filter_record_batch(batch, filter).unwrap()
1911                    } else {
1912                        batch.clone()
1913                    }
1914                })
1915                .collect::<Vec<_>>();
1916            concat_batches(&schema, &filtered_batches).unwrap()
1917        }
1918
1919        /// Return a copy of self where every other batch has had its nulls removed
1920        /// (there are often fast paths that are used when there are no nulls)
1921        fn make_half_non_nullable(mut self) -> Self {
1922            // remove the nulls from every other batch
1923            self.input_batches = self
1924                .input_batches
1925                .iter()
1926                .enumerate()
1927                .map(|(i, batch)| {
1928                    if i % 2 == 1 {
1929                        batch.clone()
1930                    } else {
1931                        Self::remove_nulls_from_batch(batch)
1932                    }
1933                })
1934                .collect();
1935            self.with_description("non-nullable")
1936        }
1937
1938        /// Insert several empty batches into the input before each existing input
1939        fn insert_empty_batches(mut self) -> Self {
1940            let empty_batch = RecordBatch::new_empty(self.schema());
1941            self.input_batches = self
1942                .input_batches
1943                .into_iter()
1944                .flat_map(|batch| [empty_batch.clone(), batch])
1945                .collect();
1946            let empty_filters = BooleanArray::builder(0).finish();
1947            self.filters = self
1948                .filters
1949                .into_iter()
1950                .flat_map(|filter| [empty_filters.clone(), filter])
1951                .collect();
1952            self.with_description("empty batches inserted")
1953        }
1954
1955        /// Sets one batch to be non-nullable by removing nulls from all columns
1956        fn remove_nulls_from_batch(batch: &RecordBatch) -> RecordBatch {
1957            let new_columns = batch
1958                .columns()
1959                .iter()
1960                .map(Self::remove_nulls_from_array)
1961                .collect::<Vec<_>>();
1962            let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
1963            RecordBatch::try_new_with_options(batch.schema(), new_columns, &options).unwrap()
1964        }
1965
1966        fn remove_nulls_from_array(array: &ArrayRef) -> ArrayRef {
1967            make_array(array.to_data().into_builder().nulls(None).build().unwrap())
1968        }
1969
1970        /// Returns a set of tests where each test that is the sae as self, but
1971        /// has a single column from the original input batch
1972        ///
1973        /// This can be useful to single column optimizations, specifically
1974        /// filter optimization.
1975        fn make_single_column_tests(&self) -> Vec<Self> {
1976            let original_schema = self.schema();
1977            let mut new_tests = vec![];
1978            for column in original_schema.fields() {
1979                let single_column_schema = Arc::new(Schema::new(vec![column.clone()]));
1980
1981                let single_column_batches = self.input_batches.iter().map(|batch| {
1982                    let single_column = batch.column_by_name(column.name()).unwrap();
1983                    RecordBatch::try_new(
1984                        Arc::clone(&single_column_schema),
1985                        vec![single_column.clone()],
1986                    )
1987                    .unwrap()
1988                });
1989
1990                let single_column_test = self
1991                    .clone()
1992                    .with_schema(Arc::clone(&single_column_schema))
1993                    .with_batches(single_column_batches)
1994                    .with_description("single column")
1995                    .with_description(column.name());
1996
1997                new_tests.push(single_column_test);
1998            }
1999            new_tests
2000        }
2001    }
2002
2003    /// Return a RecordBatch with a UInt32Array with the specified range and
2004    /// every third value is null.
2005    fn uint32_batch<T: std::iter::Iterator<Item = u32>>(range: T) -> RecordBatch {
2006        let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, true)]));
2007
2008        let array = UInt32Array::from_iter(range.map(|i| if i % 3 == 0 { None } else { Some(i) }));
2009        RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2010    }
2011
2012    /// Return a RecordBatch with a UInt32Array with no nulls specified range
2013    fn uint32_batch_non_null<T: std::iter::Iterator<Item = u32>>(range: T) -> RecordBatch {
2014        let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]));
2015
2016        let array = UInt32Array::from_iter_values(range);
2017        RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2018    }
2019
2020    /// Return a RecordBatch with a UInt64Array with no nulls specified range
2021    fn uint64_batch_non_null<T: std::iter::Iterator<Item = u64>>(range: T) -> RecordBatch {
2022        let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt64, false)]));
2023
2024        let array = UInt64Array::from_iter_values(range);
2025        RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2026    }
2027
2028    /// Return a RecordBatch with a StringArrary with values `value0`, `value1`, ...
2029    /// and every third value is `None`.
2030    fn utf8_batch(range: Range<u32>) -> RecordBatch {
2031        let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Utf8, true)]));
2032
2033        let array = StringArray::from_iter(range.map(|i| {
2034            if i % 3 == 0 {
2035                None
2036            } else {
2037                Some(format!("value{i}"))
2038            }
2039        }));
2040
2041        RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2042    }
2043
2044    /// Return a RecordBatch with a StringViewArray with (only) the specified values
2045    fn stringview_batch<'a>(values: impl IntoIterator<Item = Option<&'a str>>) -> RecordBatch {
2046        let schema = Arc::new(Schema::new(vec![Field::new(
2047            "c0",
2048            DataType::Utf8View,
2049            false,
2050        )]));
2051
2052        let array = StringViewArray::from_iter(values);
2053        RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2054    }
2055
2056    /// Return a RecordBatch with a StringViewArray with num_rows by repeating
2057    /// values over and over.
2058    fn stringview_batch_repeated<'a>(
2059        num_rows: usize,
2060        values: impl IntoIterator<Item = Option<&'a str>>,
2061    ) -> RecordBatch {
2062        let schema = Arc::new(Schema::new(vec![Field::new(
2063            "c0",
2064            DataType::Utf8View,
2065            true,
2066        )]));
2067
2068        // Repeat the values to a total of num_rows
2069        let values: Vec<_> = values.into_iter().collect();
2070        let values_iter = std::iter::repeat(values.iter())
2071            .flatten()
2072            .cloned()
2073            .take(num_rows);
2074
2075        let mut builder = StringViewBuilder::with_capacity(100).with_fixed_block_size(8192);
2076        for val in values_iter {
2077            builder.append_option(val);
2078        }
2079
2080        let array = builder.finish();
2081        RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2082    }
2083
2084    /// Return a RecordBatch of 100 rows
2085    fn multi_column_batch(range: Range<i32>) -> RecordBatch {
2086        let int64_array = Int64Array::from_iter(
2087            range
2088                .clone()
2089                .map(|v| if v % 5 == 0 { None } else { Some(v as i64) }),
2090        );
2091        let string_view_array = StringViewArray::from_iter(range.clone().map(|v| {
2092            if v % 5 == 0 {
2093                None
2094            } else if v % 7 == 0 {
2095                Some(format!("This is a string longer than 12 bytes{v}"))
2096            } else {
2097                Some(format!("Short {v}"))
2098            }
2099        }));
2100        let string_array = StringArray::from_iter(range.clone().map(|v| {
2101            if v % 11 == 0 {
2102                None
2103            } else {
2104                Some(format!("Value {v}"))
2105            }
2106        }));
2107        let timestamp_array = TimestampNanosecondArray::from_iter(range.map(|v| {
2108            if v % 3 == 0 {
2109                None
2110            } else {
2111                Some(v as i64 * 1000) // simulate a timestamp in milliseconds
2112            }
2113        }))
2114        .with_timezone("America/New_York");
2115
2116        RecordBatch::try_from_iter(vec![
2117            ("int64", Arc::new(int64_array) as ArrayRef),
2118            ("stringview", Arc::new(string_view_array) as ArrayRef),
2119            ("string", Arc::new(string_array) as ArrayRef),
2120            ("timestamp", Arc::new(timestamp_array) as ArrayRef),
2121        ])
2122        .unwrap()
2123    }
2124
2125    /// Return a boolean array that filters out randomly selected rows
2126    /// from the input batch with a `selectivity`.
2127    ///
2128    /// For example a `selectivity` of 0.1 will filter out
2129    /// 90% of the rows.
2130    #[derive(Debug)]
2131    struct RandomFilterBuilder {
2132        /// Number of rows to add to each filter
2133        num_rows: usize,
2134        /// selectivity of the filter (between 0.0 and 1.0)
2135        /// 0 selects no rows, 1.0 selects all rows
2136        selectivity: f64,
2137        /// seed for random number generator, increases by one each time
2138        /// `next_filter` is called
2139        seed: u64,
2140    }
2141    impl RandomFilterBuilder {
2142        /// Build the next filter with the current seed and increment the seed
2143        /// by one.
2144        fn next_filter(&mut self) -> BooleanArray {
2145            assert!(self.selectivity >= 0.0 && self.selectivity <= 1.0);
2146            let mut rng = rand::rngs::StdRng::seed_from_u64(self.seed);
2147            self.seed += 1;
2148            BooleanArray::from_iter(
2149                (0..self.num_rows)
2150                    .map(|_| rng.random_bool(self.selectivity))
2151                    .map(Some),
2152            )
2153        }
2154    }
2155
2156    /// Returns the named column as a StringViewArray
2157    fn col_as_string_view<'b>(name: &str, batch: &'b RecordBatch) -> &'b StringViewArray {
2158        batch
2159            .column_by_name(name)
2160            .expect("column not found")
2161            .as_string_view_opt()
2162            .expect("column is not a string view")
2163    }
2164
2165    /// Filter that selects 12.5% of the rows (`idx % 8 == 0`)
2166    fn sparse_filter(len: usize) -> BooleanArray {
2167        BooleanArray::from_iter((0..len).map(|idx| Some(idx % 8 == 0)))
2168    }
2169
2170    /// Filter that selects 5% of the rows (`idx % 20 == 0`)
2171    ///
2172    /// Different code paths are taken for very sparse filters
2173    fn very_sparse_filter(len: usize) -> BooleanArray {
2174        BooleanArray::from_iter((0..len).map(|idx| Some(idx % 20 == 0)))
2175    }
2176
2177    /// Normalize the `RecordBatch` so that the memory layout is consistent
2178    /// (e.g. StringArray is compacted).
2179    fn normalize_batch(batch: RecordBatch) -> RecordBatch {
2180        // Only need to normalize StringViews (as == also tests for memory layout)
2181        let (schema, mut columns, row_count) = batch.into_parts();
2182
2183        for column in columns.iter_mut() {
2184            if let Some(string_view) = column.as_string_view_opt() {
2185                // Re-create the StringViewArray to ensure memory layout is
2186                // consistent
2187                let mut builder = StringViewBuilder::new();
2188                for s in string_view.iter() {
2189                    builder.append_option(s);
2190                }
2191                *column = Arc::new(builder.finish());
2192                continue;
2193            }
2194
2195            if let Some(binary_view) = column.as_binary_view_opt() {
2196                *column = Arc::new(BinaryViewArray::from_iter(binary_view.iter()));
2197            }
2198        }
2199
2200        let options = RecordBatchOptions::new().with_row_count(Some(row_count));
2201        RecordBatch::try_new_with_options(schema, columns, &options).unwrap()
2202    }
2203
2204    /// Helper function to create a test batch with specified number of rows
2205    fn create_test_batch(num_rows: usize) -> RecordBatch {
2206        let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)]));
2207        let array = Int32Array::from_iter_values(0..num_rows as i32);
2208        RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
2209    }
2210    #[test]
2211    fn test_biggest_coalesce_batch_size_none_default() {
2212        // Test that default behavior (None) coalesces all batches
2213        let mut coalescer = BatchCoalescer::new(
2214            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2215            100,
2216        );
2217
2218        // Push a large batch (1000 rows) - should be coalesced normally
2219        let large_batch = create_test_batch(1000);
2220        coalescer.push_batch(large_batch).unwrap();
2221
2222        // Should produce multiple batches of target size (100)
2223        let mut output_batches = vec![];
2224        while let Some(batch) = coalescer.next_completed_batch() {
2225            output_batches.push(batch);
2226        }
2227
2228        coalescer.finish_buffered_batch().unwrap();
2229        while let Some(batch) = coalescer.next_completed_batch() {
2230            output_batches.push(batch);
2231        }
2232
2233        // Should have 10 batches of 100 rows each
2234        assert_eq!(output_batches.len(), 10);
2235        for batch in output_batches {
2236            assert_eq!(batch.num_rows(), 100);
2237        }
2238    }
2239
2240    #[test]
2241    fn test_biggest_coalesce_batch_size_bypass_large_batch() {
2242        // Test that batches larger than biggest_coalesce_batch_size bypass coalescing
2243        let mut coalescer = BatchCoalescer::new(
2244            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2245            100,
2246        );
2247        coalescer.set_biggest_coalesce_batch_size(Some(500));
2248
2249        // Push a large batch (1000 rows) - should bypass coalescing
2250        let large_batch = create_test_batch(1000);
2251        coalescer.push_batch(large_batch.clone()).unwrap();
2252
2253        // Should have one completed batch immediately (the original large batch)
2254        assert!(coalescer.has_completed_batch());
2255        let output_batch = coalescer.next_completed_batch().unwrap();
2256        assert_eq!(output_batch.num_rows(), 1000);
2257
2258        // Should be no more completed batches
2259        assert!(!coalescer.has_completed_batch());
2260        assert_eq!(coalescer.get_buffered_rows(), 0);
2261    }
2262
2263    #[test]
2264    fn test_biggest_coalesce_batch_size_coalesce_small_batch() {
2265        // Test that batches smaller than biggest_coalesce_batch_size are coalesced normally
2266        let mut coalescer = BatchCoalescer::new(
2267            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2268            100,
2269        )
2270        .with_biggest_coalesce_batch_size(Some(500));
2271
2272        // Push small batches that should be coalesced
2273        let small_batch = create_test_batch(50);
2274        coalescer.push_batch(small_batch.clone()).unwrap();
2275
2276        // Should not have completed batch yet (only 50 rows, target is 100)
2277        assert!(!coalescer.has_completed_batch());
2278        assert_eq!(coalescer.get_buffered_rows(), 50);
2279
2280        // Push another small batch
2281        coalescer.push_batch(small_batch).unwrap();
2282
2283        // Now should have a completed batch (100 rows total)
2284        assert!(coalescer.has_completed_batch());
2285        let output_batch = coalescer.next_completed_batch().unwrap();
2286        let size = output_batch
2287            .column(0)
2288            .as_primitive::<Int32Type>()
2289            .get_buffer_memory_size();
2290        assert_eq!(size, 400); // 100 rows * 4 bytes each
2291        assert_eq!(output_batch.num_rows(), 100);
2292
2293        assert_eq!(coalescer.get_buffered_rows(), 0);
2294    }
2295
2296    #[test]
2297    fn test_biggest_coalesce_batch_size_equal_boundary() {
2298        // Test behavior when batch size equals biggest_coalesce_batch_size
2299        let mut coalescer = BatchCoalescer::new(
2300            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2301            100,
2302        );
2303        coalescer.set_biggest_coalesce_batch_size(Some(500));
2304
2305        // Push a batch exactly equal to the limit
2306        let boundary_batch = create_test_batch(500);
2307        coalescer.push_batch(boundary_batch).unwrap();
2308
2309        // Should be coalesced (not bypass) since it's equal, not greater
2310        let mut output_count = 0;
2311        while coalescer.next_completed_batch().is_some() {
2312            output_count += 1;
2313        }
2314
2315        coalescer.finish_buffered_batch().unwrap();
2316        while coalescer.next_completed_batch().is_some() {
2317            output_count += 1;
2318        }
2319
2320        // Should have 5 batches of 100 rows each
2321        assert_eq!(output_count, 5);
2322    }
2323
2324    #[test]
2325    fn test_biggest_coalesce_batch_size_first_large_then_consecutive_bypass() {
2326        // Test the new consecutive large batch bypass behavior
2327        // Pattern: small batches -> first large batch (coalesced) -> consecutive large batches (bypass)
2328        let mut coalescer = BatchCoalescer::new(
2329            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2330            100,
2331        );
2332        coalescer.set_biggest_coalesce_batch_size(Some(200));
2333
2334        let small_batch = create_test_batch(50);
2335
2336        // Push small batch first to create buffered data
2337        coalescer.push_batch(small_batch).unwrap();
2338        assert_eq!(coalescer.get_buffered_rows(), 50);
2339        assert!(!coalescer.has_completed_batch());
2340
2341        // Push first large batch - should go through normal coalescing due to buffered data
2342        let large_batch1 = create_test_batch(250);
2343        coalescer.push_batch(large_batch1).unwrap();
2344
2345        // 50 + 250 = 300 -> 3 complete batches of 100, 0 rows buffered
2346        let mut completed_batches = vec![];
2347        while let Some(batch) = coalescer.next_completed_batch() {
2348            completed_batches.push(batch);
2349        }
2350        assert_eq!(completed_batches.len(), 3);
2351        assert_eq!(coalescer.get_buffered_rows(), 0);
2352
2353        // Now push consecutive large batches - they should bypass
2354        let large_batch2 = create_test_batch(300);
2355        let large_batch3 = create_test_batch(400);
2356
2357        // Push second large batch - should bypass since it's consecutive and buffer is empty
2358        coalescer.push_batch(large_batch2).unwrap();
2359        assert!(coalescer.has_completed_batch());
2360        let output = coalescer.next_completed_batch().unwrap();
2361        assert_eq!(output.num_rows(), 300); // bypassed with original size
2362        assert_eq!(coalescer.get_buffered_rows(), 0);
2363
2364        // Push third large batch - should also bypass
2365        coalescer.push_batch(large_batch3).unwrap();
2366        assert!(coalescer.has_completed_batch());
2367        let output = coalescer.next_completed_batch().unwrap();
2368        assert_eq!(output.num_rows(), 400); // bypassed with original size
2369        assert_eq!(coalescer.get_buffered_rows(), 0);
2370    }
2371
2372    #[test]
2373    fn test_biggest_coalesce_batch_size_empty_batch() {
2374        // Test that empty batches don't trigger the bypass logic
2375        let mut coalescer = BatchCoalescer::new(
2376            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2377            100,
2378        );
2379        coalescer.set_biggest_coalesce_batch_size(Some(50));
2380
2381        let empty_batch = create_test_batch(0);
2382        coalescer.push_batch(empty_batch).unwrap();
2383
2384        // Empty batch should be handled normally (no effect)
2385        assert!(!coalescer.has_completed_batch());
2386        assert_eq!(coalescer.get_buffered_rows(), 0);
2387    }
2388
2389    #[test]
2390    fn test_biggest_coalesce_batch_size_with_buffered_data_no_bypass() {
2391        // Test that when there is buffered data, large batches do NOT bypass (unless consecutive)
2392        let mut coalescer = BatchCoalescer::new(
2393            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2394            100,
2395        );
2396        coalescer.set_biggest_coalesce_batch_size(Some(200));
2397
2398        // Add some buffered data first
2399        let small_batch = create_test_batch(30);
2400        coalescer.push_batch(small_batch.clone()).unwrap();
2401        coalescer.push_batch(small_batch).unwrap();
2402        assert_eq!(coalescer.get_buffered_rows(), 60);
2403
2404        // Push large batch that would normally bypass, but shouldn't because buffered_rows > 0
2405        let large_batch = create_test_batch(250);
2406        coalescer.push_batch(large_batch).unwrap();
2407
2408        // The large batch should be processed through normal coalescing logic
2409        // Total: 60 (buffered) + 250 (new) = 310 rows
2410        // Output: 3 complete batches of 100 rows each, 10 rows remain buffered
2411
2412        let mut completed_batches = vec![];
2413        while let Some(batch) = coalescer.next_completed_batch() {
2414            completed_batches.push(batch);
2415        }
2416
2417        assert_eq!(completed_batches.len(), 3);
2418        for batch in &completed_batches {
2419            assert_eq!(batch.num_rows(), 100);
2420        }
2421        assert_eq!(coalescer.get_buffered_rows(), 10);
2422    }
2423
2424    #[test]
2425    fn test_biggest_coalesce_batch_size_zero_limit() {
2426        // Test edge case where limit is 0 (all batches bypass when no buffered data)
2427        let mut coalescer = BatchCoalescer::new(
2428            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2429            100,
2430        );
2431        coalescer.set_biggest_coalesce_batch_size(Some(0));
2432
2433        // Even a 1-row batch should bypass when there's no buffered data
2434        let tiny_batch = create_test_batch(1);
2435        coalescer.push_batch(tiny_batch).unwrap();
2436
2437        assert!(coalescer.has_completed_batch());
2438        let output = coalescer.next_completed_batch().unwrap();
2439        assert_eq!(output.num_rows(), 1);
2440    }
2441
2442    #[test]
2443    fn test_biggest_coalesce_batch_size_bypass_only_when_no_buffer() {
2444        // Test that bypass only occurs when buffered_rows == 0
2445        let mut coalescer = BatchCoalescer::new(
2446            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2447            100,
2448        );
2449        coalescer.set_biggest_coalesce_batch_size(Some(200));
2450
2451        // First, push a large batch with no buffered data - should bypass
2452        let large_batch = create_test_batch(300);
2453        coalescer.push_batch(large_batch.clone()).unwrap();
2454
2455        assert!(coalescer.has_completed_batch());
2456        let output = coalescer.next_completed_batch().unwrap();
2457        assert_eq!(output.num_rows(), 300); // bypassed
2458        assert_eq!(coalescer.get_buffered_rows(), 0);
2459
2460        // Now add some buffered data
2461        let small_batch = create_test_batch(50);
2462        coalescer.push_batch(small_batch).unwrap();
2463        assert_eq!(coalescer.get_buffered_rows(), 50);
2464
2465        // Push the same large batch again - should NOT bypass this time (not consecutive)
2466        coalescer.push_batch(large_batch).unwrap();
2467
2468        // Should process through normal coalescing: 50 + 300 = 350 rows
2469        // Output: 3 complete batches of 100 rows, 50 rows buffered
2470        let mut completed_batches = vec![];
2471        while let Some(batch) = coalescer.next_completed_batch() {
2472            completed_batches.push(batch);
2473        }
2474
2475        assert_eq!(completed_batches.len(), 3);
2476        for batch in &completed_batches {
2477            assert_eq!(batch.num_rows(), 100);
2478        }
2479        assert_eq!(coalescer.get_buffered_rows(), 50);
2480    }
2481
2482    #[test]
2483    fn test_biggest_coalesce_batch_size_consecutive_large_batches_scenario() {
2484        // Test your exact scenario: 20, 20, 30, 700, 600, 700, 900, 700, 600
2485        let mut coalescer = BatchCoalescer::new(
2486            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2487            1000,
2488        );
2489        coalescer.set_biggest_coalesce_batch_size(Some(500));
2490
2491        // Push small batches first
2492        coalescer.push_batch(create_test_batch(20)).unwrap();
2493        coalescer.push_batch(create_test_batch(20)).unwrap();
2494        coalescer.push_batch(create_test_batch(30)).unwrap();
2495
2496        assert_eq!(coalescer.get_buffered_rows(), 70);
2497        assert!(!coalescer.has_completed_batch());
2498
2499        // Push first large batch (700) - should coalesce due to buffered data
2500        coalescer.push_batch(create_test_batch(700)).unwrap();
2501
2502        // 70 + 700 = 770 rows, not enough for 1000, so all stay buffered
2503        assert_eq!(coalescer.get_buffered_rows(), 770);
2504        assert!(!coalescer.has_completed_batch());
2505
2506        // Push second large batch (600) - should bypass since previous was large
2507        coalescer.push_batch(create_test_batch(600)).unwrap();
2508
2509        // Should flush buffer (770 rows) and bypass the 600
2510        let mut outputs = vec![];
2511        while let Some(batch) = coalescer.next_completed_batch() {
2512            outputs.push(batch);
2513        }
2514        assert_eq!(outputs.len(), 2); // one flushed buffer batch (770) + one bypassed (600)
2515        assert_eq!(outputs[0].num_rows(), 770);
2516        assert_eq!(outputs[1].num_rows(), 600);
2517        assert_eq!(coalescer.get_buffered_rows(), 0);
2518
2519        // Push remaining large batches - should all bypass
2520        let remaining_batches = [700, 900, 700, 600];
2521        for &size in &remaining_batches {
2522            coalescer.push_batch(create_test_batch(size)).unwrap();
2523
2524            assert!(coalescer.has_completed_batch());
2525            let output = coalescer.next_completed_batch().unwrap();
2526            assert_eq!(output.num_rows(), size);
2527            assert_eq!(coalescer.get_buffered_rows(), 0);
2528        }
2529    }
2530
2531    #[test]
2532    fn test_biggest_coalesce_batch_size_truly_consecutive_large_bypass() {
2533        // Test truly consecutive large batches that should all bypass
2534        // This test ensures buffer is completely empty between large batches
2535        let mut coalescer = BatchCoalescer::new(
2536            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2537            100,
2538        );
2539        coalescer.set_biggest_coalesce_batch_size(Some(200));
2540
2541        // Push consecutive large batches with no prior buffered data
2542        let large_batches = vec![
2543            create_test_batch(300),
2544            create_test_batch(400),
2545            create_test_batch(350),
2546            create_test_batch(500),
2547        ];
2548
2549        let mut all_outputs = vec![];
2550
2551        for (i, large_batch) in large_batches.into_iter().enumerate() {
2552            let expected_size = large_batch.num_rows();
2553
2554            // Buffer should be empty before each large batch
2555            assert_eq!(
2556                coalescer.get_buffered_rows(),
2557                0,
2558                "Buffer should be empty before batch {}",
2559                i
2560            );
2561
2562            coalescer.push_batch(large_batch).unwrap();
2563
2564            // Each large batch should bypass and produce exactly one output batch
2565            assert!(
2566                coalescer.has_completed_batch(),
2567                "Should have completed batch after pushing batch {}",
2568                i
2569            );
2570
2571            let output = coalescer.next_completed_batch().unwrap();
2572            assert_eq!(
2573                output.num_rows(),
2574                expected_size,
2575                "Batch {} should have bypassed with original size",
2576                i
2577            );
2578
2579            // Should be no more batches and buffer should be empty
2580            assert!(
2581                !coalescer.has_completed_batch(),
2582                "Should have no more completed batches after batch {}",
2583                i
2584            );
2585            assert_eq!(
2586                coalescer.get_buffered_rows(),
2587                0,
2588                "Buffer should be empty after batch {}",
2589                i
2590            );
2591
2592            all_outputs.push(output);
2593        }
2594
2595        // Verify we got exactly 4 output batches with original sizes
2596        assert_eq!(all_outputs.len(), 4);
2597        assert_eq!(all_outputs[0].num_rows(), 300);
2598        assert_eq!(all_outputs[1].num_rows(), 400);
2599        assert_eq!(all_outputs[2].num_rows(), 350);
2600        assert_eq!(all_outputs[3].num_rows(), 500);
2601    }
2602
2603    #[test]
2604    fn test_biggest_coalesce_batch_size_reset_consecutive_on_small_batch() {
2605        // Test that small batches reset the consecutive large batch tracking
2606        let mut coalescer = BatchCoalescer::new(
2607            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2608            100,
2609        );
2610        coalescer.set_biggest_coalesce_batch_size(Some(200));
2611
2612        // Push first large batch - should bypass (no buffered data)
2613        coalescer.push_batch(create_test_batch(300)).unwrap();
2614        let output = coalescer.next_completed_batch().unwrap();
2615        assert_eq!(output.num_rows(), 300);
2616
2617        // Push second large batch - should bypass (consecutive)
2618        coalescer.push_batch(create_test_batch(400)).unwrap();
2619        let output = coalescer.next_completed_batch().unwrap();
2620        assert_eq!(output.num_rows(), 400);
2621
2622        // Push small batch - resets consecutive tracking
2623        coalescer.push_batch(create_test_batch(50)).unwrap();
2624        assert_eq!(coalescer.get_buffered_rows(), 50);
2625
2626        // Push large batch again - should NOT bypass due to buffered data
2627        coalescer.push_batch(create_test_batch(350)).unwrap();
2628
2629        // Should coalesce: 50 + 350 = 400 -> 4 complete batches of 100
2630        let mut outputs = vec![];
2631        while let Some(batch) = coalescer.next_completed_batch() {
2632            outputs.push(batch);
2633        }
2634        assert_eq!(outputs.len(), 4);
2635        for batch in outputs {
2636            assert_eq!(batch.num_rows(), 100);
2637        }
2638        assert_eq!(coalescer.get_buffered_rows(), 0);
2639    }
2640
2641    #[test]
2642    fn test_coalasce_push_batch_with_indices() {
2643        const MID_POINT: u32 = 2333;
2644        const TOTAL_ROWS: u32 = 23333;
2645        let batch1 = uint32_batch_non_null(0..MID_POINT);
2646        let batch2 = uint32_batch_non_null((MID_POINT..TOTAL_ROWS).rev());
2647
2648        let mut coalescer = BatchCoalescer::new(
2649            Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])),
2650            TOTAL_ROWS as usize,
2651        );
2652        coalescer.push_batch(batch1).unwrap();
2653
2654        let rev_indices = (0..((TOTAL_ROWS - MID_POINT) as u64)).rev();
2655        let reversed_indices_batch = uint64_batch_non_null(rev_indices);
2656
2657        let reverse_indices = UInt64Array::from(reversed_indices_batch.column(0).to_data());
2658        coalescer
2659            .push_batch_with_indices(batch2, &reverse_indices)
2660            .unwrap();
2661
2662        coalescer.finish_buffered_batch().unwrap();
2663        let actual = coalescer.next_completed_batch().unwrap();
2664
2665        let expected = uint32_batch_non_null(0..TOTAL_ROWS);
2666
2667        assert_eq!(expected, actual);
2668    }
2669
2670    #[test]
2671    fn test_push_batch_schema_mismatch_fewer_columns() {
2672        // Coalescer expects 0 columns, batch has 1
2673        let empty_schema = Arc::new(Schema::empty());
2674        let mut coalescer = BatchCoalescer::new(empty_schema, 100);
2675        let batch = uint32_batch(0..5);
2676        let result = coalescer.push_batch(batch);
2677        assert!(result.is_err());
2678        let err = result.unwrap_err().to_string();
2679        assert!(
2680            err.contains("Batch has 1 columns but BatchCoalescer expects 0"),
2681            "unexpected error: {err}"
2682        );
2683    }
2684
2685    #[test]
2686    fn test_push_batch_schema_mismatch_more_columns() {
2687        // Coalescer expects 2 columns, batch has 1
2688        let schema = Arc::new(Schema::new(vec![
2689            Field::new("c0", DataType::UInt32, false),
2690            Field::new("c1", DataType::UInt32, false),
2691        ]));
2692        let mut coalescer = BatchCoalescer::new(schema, 100);
2693        let batch = uint32_batch(0..5);
2694        let result = coalescer.push_batch(batch);
2695        assert!(result.is_err());
2696        let err = result.unwrap_err().to_string();
2697        assert!(
2698            err.contains("Batch has 1 columns but BatchCoalescer expects 2"),
2699            "unexpected error: {err}"
2700        );
2701    }
2702
2703    #[test]
2704    fn test_push_batch_schema_mismatch_two_vs_zero() {
2705        // Coalescer expects 0 columns, batch has 2
2706        let empty_schema = Arc::new(Schema::empty());
2707        let mut coalescer = BatchCoalescer::new(empty_schema, 100);
2708        let schema = Arc::new(Schema::new(vec![
2709            Field::new("c0", DataType::UInt32, false),
2710            Field::new("c1", DataType::UInt32, false),
2711        ]));
2712        let batch = RecordBatch::try_new(
2713            schema,
2714            vec![
2715                Arc::new(UInt32Array::from(vec![1, 2, 3])),
2716                Arc::new(UInt32Array::from(vec![4, 5, 6])),
2717            ],
2718        )
2719        .unwrap();
2720        let result = coalescer.push_batch(batch);
2721        assert!(result.is_err());
2722        let err = result.unwrap_err().to_string();
2723        assert!(
2724            err.contains("Batch has 2 columns but BatchCoalescer expects 0"),
2725            "unexpected error: {err}"
2726        );
2727    }
2728
2729    #[test]
2730    fn test_size_grows_with_buffering_and_shrinks_when_draining() {
2731        let batch = uint32_batch(0..8);
2732        let mut coalescer = BatchCoalescer::new(batch.schema(), 21);
2733        let baseline = coalescer.size();
2734        assert!(baseline > 0, "size includes container capacities");
2735
2736        // Buffer rows without completing a batch
2737        coalescer.push_batch(batch.clone()).unwrap();
2738        assert!(coalescer.next_completed_batch().is_none());
2739        let buffered = coalescer.size();
2740        assert!(
2741            buffered > baseline,
2742            "buffering rows should grow size ({buffered} > {baseline})"
2743        );
2744
2745        // Push enough to complete several batches
2746        for _ in 0..10 {
2747            coalescer.push_batch(batch.clone()).unwrap();
2748        }
2749        let peak = coalescer.size();
2750        assert!(peak > buffered);
2751
2752        // Draining completed batches must never grow the reported size
2753        let mut prev = peak;
2754        let mut drained_any = false;
2755        while coalescer.next_completed_batch().is_some() {
2756            drained_any = true;
2757            let now = coalescer.size();
2758            assert!(now <= prev, "size grew while draining: {now} > {prev}");
2759            prev = now;
2760        }
2761        assert!(drained_any);
2762        assert!(
2763            prev < peak,
2764            "draining completed batches should release memory"
2765        );
2766    }
2767
2768    #[test]
2769    fn test_size_string_view_buffers_released_after_drain() {
2770        // Long strings spill into external data buffers, exercising the byte-view
2771        // size accounting through the real coalescer path (including compaction).
2772        let batch = stringview_batch_repeated(
2773            1000,
2774            [Some("this string is definitely longer than 12 bytes")],
2775        );
2776        let mut coalescer = BatchCoalescer::new(batch.schema(), 4096);
2777        let baseline = coalescer.size();
2778
2779        for _ in 0..20 {
2780            coalescer.push_batch(batch.clone()).unwrap();
2781        }
2782        let peak = coalescer.size();
2783        assert!(
2784            peak > baseline,
2785            "buffered string view data should grow size ({peak} > {baseline})"
2786        );
2787
2788        coalescer.finish_buffered_batch().unwrap();
2789        while coalescer.next_completed_batch().is_some() {}
2790
2791        // Once fully drained the in-progress byte-view buffers are released.
2792        let drained = coalescer.size();
2793        assert!(
2794            drained < peak,
2795            "draining should release buffered data ({drained} < {peak})"
2796        );
2797    }
2798
2799    /// Every byte added to the accounting must eventually be removed: running the
2800    /// exact same push/finish/drain sequence twice must report identical sizes.
2801    /// This catches accounting leaks and drift without hard-coding magic numbers.
2802    #[test]
2803    fn test_size_accounting_conserved_across_cycles() {
2804        // Primitive column: internal capacities stabilize after the first cycle
2805        // (unlike byte-view, whose buffer sizer keeps growing), so the readings
2806        // are deterministic across cycles.
2807        let batch = uint32_batch(0..8);
2808        let mut coalescer = BatchCoalescer::new(batch.schema(), 4096);
2809
2810        let run_cycle = |coalescer: &mut BatchCoalescer| {
2811            for _ in 0..20 {
2812                coalescer.push_batch(batch.clone()).unwrap();
2813            }
2814            coalescer.finish_buffered_batch().unwrap();
2815            let peak = coalescer.size();
2816            while coalescer.next_completed_batch().is_some() {}
2817            (peak, coalescer.size())
2818        };
2819
2820        let (peak1, drained1) = run_cycle(&mut coalescer);
2821        let (peak2, drained2) = run_cycle(&mut coalescer);
2822
2823        assert_eq!(
2824            peak1, peak2,
2825            "identical work must report identical peak size"
2826        );
2827        assert_eq!(
2828            drained1, drained2,
2829            "fully-drained size must be stable across cycles (no accounting leak)"
2830        );
2831        assert!(
2832            drained1 < peak1,
2833            "draining must release the accounted memory"
2834        );
2835    }
2836}