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            if batch_size > limit {
454                // Case 1: No buffered data - emit large batch directly
455                // Example: [] + [1200] → output [1200], buffer []
456                if self.buffered_rows == 0 {
457                    self.completed.push_back(batch);
458                    return Ok(());
459                }
460
461                // Case 2: Buffer too large - flush then emit to avoid oversized merge
462                // Example: [850] + [1200] → output [850], then output [1200]
463                // This prevents creating batches much larger than both target_batch_size
464                // and biggest_coalesce_batch_size, which could cause memory issues
465                if self.buffered_rows > limit {
466                    self.finish_buffered_batch()?;
467                    self.completed.push_back(batch);
468                    return Ok(());
469                }
470
471                // Case 3: Small buffer - proceed with normal coalescing
472                // Example: [300] + [1200] → split and merge normally
473                // This ensures small batches still get properly coalesced
474                // while allowing some controlled growth beyond the limit
475            }
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 and uses other kernels concatenates them when finished.
711///
712/// Some types have specialized implementations for this array types (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    /// 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    fn copy_rows_by_filter(&mut self, filter: &FilterPredicate) -> Result<(), ArrowError> {
732        self.copy_rows_by_selection(filter.selection())
733    }
734
735    /// Copy rows selected by `filter` from `source`.
736    fn copy_rows_by_filter_from(
737        &mut self,
738        source: ArrayRef,
739        filter: &FilterPredicate,
740    ) -> Result<(), ArrowError> {
741        self.set_source(Some(source));
742        let result = self.copy_rows_by_filter(filter);
743        self.set_source(None);
744        result
745    }
746
747    /// Copy rows described by a [`FilterSelection`] from the current source array.
748    fn copy_rows_by_selection(&mut self, selection: FilterSelection<'_>) -> Result<(), ArrowError> {
749        match selection {
750            FilterSelection::None => Ok(()),
751            FilterSelection::All { len } => self.copy_rows(0, len),
752            FilterSelection::Slices(slices) => {
753                slices.try_for_each(|(start, end)| self.copy_rows(start, end - start))
754            }
755            FilterSelection::Indices(indices) => indices.try_for_each(|idx| self.copy_rows(idx, 1)),
756        }
757    }
758
759    /// Finish the currently in-progress array and return it as an `ArrayRef`
760    fn finish(&mut self) -> Result<ArrayRef, ArrowError>;
761
762    /// Get the number of bytes this array is using
763    fn size(&self) -> usize;
764}
765
766#[cfg(test)]
767mod tests {
768    use super::*;
769    use crate::concat::concat_batches;
770    use crate::filter::filter_record_batch;
771    use arrow_array::builder::StringViewBuilder;
772    use arrow_array::cast::AsArray;
773    use arrow_array::types::Int32Type;
774    use arrow_array::{
775        BinaryViewArray, Int32Array, Int64Array, RecordBatchOptions, StringArray, StringViewArray,
776        TimestampNanosecondArray, UInt32Array, UInt64Array, make_array,
777    };
778    use arrow_buffer::BooleanBufferBuilder;
779    use arrow_schema::{DataType, Field, Schema};
780    use rand::{Rng, SeedableRng};
781    use std::ops::Range;
782
783    #[test]
784    fn test_coalesce() {
785        let batch = uint32_batch(0..8);
786        Test::new("coalesce")
787            .with_batches(std::iter::repeat_n(batch, 10))
788            // expected output is exactly 21 rows (except for the final batch)
789            .with_batch_size(21)
790            .with_expected_output_sizes(vec![21, 21, 21, 17])
791            .run();
792    }
793
794    #[test]
795    fn test_coalesce_one_by_one() {
796        let batch = uint32_batch(0..1); // single row input
797        Test::new("coalesce_one_by_one")
798            .with_batches(std::iter::repeat_n(batch, 97))
799            // expected output is exactly 20 rows (except for the final batch)
800            .with_batch_size(20)
801            .with_expected_output_sizes(vec![20, 20, 20, 20, 17])
802            .run();
803    }
804
805    #[test]
806    fn test_coalesce_empty() {
807        let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]));
808
809        Test::new("coalesce_empty")
810            .with_batches(vec![])
811            .with_schema(schema)
812            .with_batch_size(21)
813            .with_expected_output_sizes(vec![])
814            .run();
815    }
816
817    #[test]
818    fn test_sparse_filter_copy_threshold() {
819        assert!(should_use_sparse_filter_copy(8192, 8));
820        assert!(should_use_sparse_filter_copy(8192, 81));
821        assert!(!should_use_sparse_filter_copy(8192, 819));
822        assert!(!should_use_sparse_filter_copy(8192, 6553));
823    }
824
825    #[test]
826    fn test_single_large_batch_greater_than_target() {
827        // test a single large batch
828        let batch = uint32_batch(0..4096);
829        Test::new("coalesce_single_large_batch_greater_than_target")
830            .with_batch(batch)
831            .with_batch_size(1000)
832            .with_expected_output_sizes(vec![1000, 1000, 1000, 1000, 96])
833            .run();
834    }
835
836    #[test]
837    fn test_single_large_batch_smaller_than_target() {
838        // test a single large batch
839        let batch = uint32_batch(0..4096);
840        Test::new("coalesce_single_large_batch_smaller_than_target")
841            .with_batch(batch)
842            .with_batch_size(8192)
843            .with_expected_output_sizes(vec![4096])
844            .run();
845    }
846
847    #[test]
848    fn test_single_large_batch_equal_to_target() {
849        // test a single large batch
850        let batch = uint32_batch(0..4096);
851        Test::new("coalesce_single_large_batch_equal_to_target")
852            .with_batch(batch)
853            .with_batch_size(4096)
854            .with_expected_output_sizes(vec![4096])
855            .run();
856    }
857
858    #[test]
859    fn test_single_large_batch_equally_divisible_in_target() {
860        // test a single large batch
861        let batch = uint32_batch(0..4096);
862        Test::new("coalesce_single_large_batch_equally_divisible_in_target")
863            .with_batch(batch)
864            .with_batch_size(1024)
865            .with_expected_output_sizes(vec![1024, 1024, 1024, 1024])
866            .run();
867    }
868
869    #[test]
870    fn test_empty_schema() {
871        let schema = Schema::empty();
872        let batch = RecordBatch::new_empty(schema.into());
873        Test::new("coalesce_empty_schema")
874            .with_batch(batch)
875            .with_expected_output_sizes(vec![])
876            .run();
877    }
878
879    /// Coalesce multiple batches, 80k rows, with a 0.1% selectivity filter
880    #[test]
881    fn test_coalesce_filtered_001() {
882        let mut filter_builder = RandomFilterBuilder {
883            num_rows: 8000,
884            selectivity: 0.001,
885            seed: 0,
886        };
887
888        // add 10 batches of 8000 rows each
889        // 80k rows, selecting 0.1% means 80 rows
890        // not exactly 80 as the rows are random;
891        let mut test = Test::new("coalesce_filtered_001");
892        for _ in 0..10 {
893            test = test
894                .with_batch(multi_column_batch(0..8000))
895                .with_filter(filter_builder.next_filter())
896        }
897        test.with_batch_size(15)
898            .with_expected_output_sizes(vec![15, 15, 15, 13])
899            .run();
900    }
901
902    /// Coalesce multiple batches, 80k rows, with a 1% selectivity filter
903    #[test]
904    fn test_coalesce_filtered_01() {
905        let mut filter_builder = RandomFilterBuilder {
906            num_rows: 8000,
907            selectivity: 0.01,
908            seed: 0,
909        };
910
911        // add 10 batches of 8000 rows each
912        // 80k rows, selecting 1% means 800 rows
913        // not exactly 800 as the rows are random;
914        let mut test = Test::new("coalesce_filtered_01");
915        for _ in 0..10 {
916            test = test
917                .with_batch(multi_column_batch(0..8000))
918                .with_filter(filter_builder.next_filter())
919        }
920        test.with_batch_size(128)
921            .with_expected_output_sizes(vec![128, 128, 128, 128, 128, 128, 15])
922            .run();
923    }
924
925    /// Coalesce multiple batches, 80k rows, with a 10% selectivity filter
926    #[test]
927    fn test_coalesce_filtered_10() {
928        let mut filter_builder = RandomFilterBuilder {
929            num_rows: 8000,
930            selectivity: 0.1,
931            seed: 0,
932        };
933
934        // add 10 batches of 8000 rows each
935        // 80k rows, selecting 10% means 8000 rows
936        // not exactly 800 as the rows are random;
937        let mut test = Test::new("coalesce_filtered_10");
938        for _ in 0..10 {
939            test = test
940                .with_batch(multi_column_batch(0..8000))
941                .with_filter(filter_builder.next_filter())
942        }
943        test.with_batch_size(1024)
944            .with_expected_output_sizes(vec![1024, 1024, 1024, 1024, 1024, 1024, 1024, 840])
945            .run();
946    }
947
948    /// Coalesce multiple batches, 8k rows, with a 90% selectivity filter
949    #[test]
950    fn test_coalesce_filtered_90() {
951        let mut filter_builder = RandomFilterBuilder {
952            num_rows: 800,
953            selectivity: 0.90,
954            seed: 0,
955        };
956
957        // add 10 batches of 800 rows each
958        // 8k rows, selecting 99% means 7200 rows
959        // not exactly 7200 as the rows are random;
960        let mut test = Test::new("coalesce_filtered_90");
961        for _ in 0..10 {
962            test = test
963                .with_batch(multi_column_batch(0..800))
964                .with_filter(filter_builder.next_filter())
965        }
966        test.with_batch_size(1024)
967            .with_expected_output_sizes(vec![1024, 1024, 1024, 1024, 1024, 1024, 1024, 13])
968            .run();
969    }
970
971    /// Coalesce multiple batches, 8k rows, with mixed filers, including 100%
972    #[test]
973    fn test_coalesce_filtered_mixed() {
974        let mut filter_builder = RandomFilterBuilder {
975            num_rows: 800,
976            selectivity: 0.90,
977            seed: 0,
978        };
979
980        let mut test = Test::new("coalesce_filtered_mixed");
981        for _ in 0..3 {
982            // also add in a batch that selects almost all rows and when
983            // sliced will have some batches that are entirely used
984            let mut all_filter_builder = BooleanBufferBuilder::new(1000);
985            all_filter_builder.append_n(500, true);
986            all_filter_builder.append_n(1, false);
987            all_filter_builder.append_n(499, false);
988            let all_filter = all_filter_builder.build();
989
990            test = test
991                .with_batch(multi_column_batch(0..1000))
992                .with_filter(BooleanArray::from(all_filter))
993                .with_batch(multi_column_batch(0..800))
994                .with_filter(filter_builder.next_filter());
995            // decrease selectivity
996            filter_builder.selectivity *= 0.6;
997        }
998
999        // use a small batch size to ensure the filter is appended in slices
1000        // and some of those slides will select the entire thing.
1001        test.with_batch_size(250)
1002            .with_expected_output_sizes(vec![
1003                250, 250, 250, 250, 250, 250, 250, 250, 250, 250, 250, 179,
1004            ])
1005            .run();
1006    }
1007
1008    #[test]
1009    fn test_coalesce_non_null() {
1010        Test::new("coalesce_non_null")
1011            // 4040 rows of unit32
1012            .with_batch(uint32_batch_non_null(0..3000))
1013            .with_batch(uint32_batch_non_null(0..1040))
1014            .with_batch_size(1024)
1015            .with_expected_output_sizes(vec![1024, 1024, 1024, 968])
1016            .run();
1017    }
1018    #[test]
1019    fn test_utf8_split() {
1020        Test::new("coalesce_utf8")
1021            // 4040 rows of utf8 strings in total, split into batches of 1024
1022            .with_batch(utf8_batch(0..3000))
1023            .with_batch(utf8_batch(0..1040))
1024            .with_batch_size(1024)
1025            .with_expected_output_sizes(vec![1024, 1024, 1024, 968])
1026            .run();
1027    }
1028
1029    #[test]
1030    fn test_string_view_no_views() {
1031        let output_batches = Test::new("coalesce_string_view_no_views")
1032            // both input batches have no views, so no need to compact
1033            .with_batch(stringview_batch([Some("foo"), Some("bar")]))
1034            .with_batch(stringview_batch([Some("baz"), Some("qux")]))
1035            .with_expected_output_sizes(vec![4])
1036            .run();
1037
1038        expect_buffer_layout(
1039            col_as_string_view("c0", output_batches.first().unwrap()),
1040            vec![],
1041        );
1042    }
1043
1044    #[test]
1045    fn test_string_view_batch_small_no_compact() {
1046        // view with only short strings (no buffers) --> no need to compact
1047        let batch = stringview_batch_repeated(1000, [Some("a"), Some("b"), Some("c")]);
1048        let output_batches = Test::new("coalesce_string_view_batch_small_no_compact")
1049            .with_batch(batch.clone())
1050            .with_expected_output_sizes(vec![1000])
1051            .run();
1052
1053        let array = col_as_string_view("c0", &batch);
1054        let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1055        assert_eq!(array.data_buffers().len(), 0);
1056        assert_eq!(array.data_buffers().len(), gc_array.data_buffers().len()); // no compaction
1057
1058        expect_buffer_layout(gc_array, vec![]);
1059    }
1060
1061    #[test]
1062    fn test_string_view_batch_large_no_compact() {
1063        // view with large strings (has buffers) but full --> no need to compact
1064        let batch = stringview_batch_repeated(1000, [Some("This string is longer than 12 bytes")]);
1065        let output_batches = Test::new("coalesce_string_view_batch_large_no_compact")
1066            .with_batch(batch.clone())
1067            .with_batch_size(1000)
1068            .with_expected_output_sizes(vec![1000])
1069            .run();
1070
1071        let array = col_as_string_view("c0", &batch);
1072        let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1073        assert_eq!(array.data_buffers().len(), 5);
1074        assert_eq!(array.data_buffers().len(), gc_array.data_buffers().len()); // no compaction
1075
1076        expect_buffer_layout(
1077            gc_array,
1078            vec![
1079                ExpectedLayout {
1080                    len: 8190,
1081                    capacity: 8192,
1082                },
1083                ExpectedLayout {
1084                    len: 8190,
1085                    capacity: 8192,
1086                },
1087                ExpectedLayout {
1088                    len: 8190,
1089                    capacity: 8192,
1090                },
1091                ExpectedLayout {
1092                    len: 8190,
1093                    capacity: 8192,
1094                },
1095                ExpectedLayout {
1096                    len: 2240,
1097                    capacity: 8192,
1098                },
1099            ],
1100        );
1101    }
1102
1103    #[test]
1104    fn test_string_view_batch_small_with_buffers_no_compact() {
1105        // view with buffers but only short views
1106        let short_strings = std::iter::repeat(Some("SmallString"));
1107        let long_strings = std::iter::once(Some("This string is longer than 12 bytes"));
1108        // 20 short strings, then a long ones
1109        let values = short_strings.take(20).chain(long_strings);
1110        let batch = stringview_batch_repeated(1000, values)
1111            // take only 10 short strings (no long ones)
1112            .slice(5, 10);
1113        let output_batches = Test::new("coalesce_string_view_batch_small_with_buffers_no_compact")
1114            .with_batch(batch.clone())
1115            .with_batch_size(1000)
1116            .with_expected_output_sizes(vec![10])
1117            .run();
1118
1119        let array = col_as_string_view("c0", &batch);
1120        let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1121        assert_eq!(array.data_buffers().len(), 1); // input has one buffer
1122        assert_eq!(gc_array.data_buffers().len(), 0); // output has no buffers as only short strings
1123    }
1124
1125    #[test]
1126    fn test_string_view_batch_large_slice_compact() {
1127        // view with large strings (has buffers) and only partially used  --> no need to compact
1128        let batch = stringview_batch_repeated(1000, [Some("This string is longer than 12 bytes")])
1129            // slice only 22 rows, so most of the buffer is not used
1130            .slice(11, 22);
1131
1132        let output_batches = Test::new("coalesce_string_view_batch_large_slice_compact")
1133            .with_batch(batch.clone())
1134            .with_batch_size(1000)
1135            .with_expected_output_sizes(vec![22])
1136            .run();
1137
1138        let array = col_as_string_view("c0", &batch);
1139        let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1140        assert_eq!(array.data_buffers().len(), 5);
1141
1142        expect_buffer_layout(
1143            gc_array,
1144            vec![ExpectedLayout {
1145                len: 770,
1146                capacity: 8192,
1147            }],
1148        );
1149    }
1150
1151    #[test]
1152    fn test_string_view_mixed() {
1153        let large_view_batch =
1154            stringview_batch_repeated(1000, [Some("This string is longer than 12 bytes")]);
1155        let small_view_batch = stringview_batch_repeated(1000, [Some("SmallString")]);
1156        let mixed_batch = stringview_batch_repeated(
1157            1000,
1158            [Some("This string is longer than 12 bytes"), Some("Small")],
1159        );
1160        let mixed_batch_nulls = stringview_batch_repeated(
1161            1000,
1162            [
1163                Some("This string is longer than 12 bytes"),
1164                Some("Small"),
1165                None,
1166            ],
1167        );
1168
1169        // Several batches with mixed inline / non inline
1170        // 4k rows in
1171        let output_batches = Test::new("coalesce_string_view_mixed")
1172            .with_batch(large_view_batch.clone())
1173            .with_batch(small_view_batch)
1174            // this batch needs to be compacted (less than 1/2 full)
1175            .with_batch(large_view_batch.slice(10, 20))
1176            .with_batch(mixed_batch_nulls)
1177            // this batch needs to be compacted (less than 1/2 full)
1178            .with_batch(large_view_batch.slice(10, 20))
1179            .with_batch(mixed_batch)
1180            .with_expected_output_sizes(vec![1024, 1024, 1024, 968])
1181            .run();
1182
1183        expect_buffer_layout(
1184            col_as_string_view("c0", output_batches.first().unwrap()),
1185            vec![
1186                ExpectedLayout {
1187                    len: 8190,
1188                    capacity: 8192,
1189                },
1190                ExpectedLayout {
1191                    len: 8190,
1192                    capacity: 8192,
1193                },
1194                ExpectedLayout {
1195                    len: 8190,
1196                    capacity: 8192,
1197                },
1198                ExpectedLayout {
1199                    len: 8190,
1200                    capacity: 8192,
1201                },
1202                ExpectedLayout {
1203                    len: 2240,
1204                    capacity: 8192,
1205                },
1206            ],
1207        );
1208    }
1209
1210    #[test]
1211    fn test_string_view_many_small_compact() {
1212        // 200 rows alternating long (28) and short (≤12) strings.
1213        // Only the 100 long strings go into data buffers: 100 × 28 = 2800.
1214        let batch = stringview_batch_repeated(
1215            200,
1216            [Some("This string is 28 bytes long"), Some("small string")],
1217        );
1218        let output_batches = Test::new("coalesce_string_view_many_small_compact")
1219            // First allocated buffer is 8kb.
1220            // Appending 10 batches of 2800 bytes will use 2800 * 10 = 14kb (8kb, an 16kb and 32kbkb)
1221            .with_batch(batch.clone())
1222            .with_batch(batch.clone())
1223            .with_batch(batch.clone())
1224            .with_batch(batch.clone())
1225            .with_batch(batch.clone())
1226            .with_batch(batch.clone())
1227            .with_batch(batch.clone())
1228            .with_batch(batch.clone())
1229            .with_batch(batch.clone())
1230            .with_batch(batch.clone())
1231            .with_batch_size(8000)
1232            .with_expected_output_sizes(vec![2000]) // only 1000 rows total
1233            .run();
1234
1235        // expect a nice even distribution of buffers
1236        expect_buffer_layout(
1237            col_as_string_view("c0", output_batches.first().unwrap()),
1238            vec![
1239                ExpectedLayout {
1240                    len: 8176,
1241                    capacity: 8192,
1242                },
1243                ExpectedLayout {
1244                    len: 16380,
1245                    capacity: 16384,
1246                },
1247                ExpectedLayout {
1248                    len: 3444,
1249                    capacity: 32768,
1250                },
1251            ],
1252        );
1253    }
1254
1255    #[test]
1256    fn test_string_view_many_small_boundary() {
1257        // The strings are designed to exactly fit into buffers that are powers of 2 long
1258        let batch = stringview_batch_repeated(100, [Some("This string is a power of two=32")]);
1259        let output_batches = Test::new("coalesce_string_view_many_small_boundary")
1260            .with_batches(std::iter::repeat_n(batch, 20))
1261            .with_batch_size(900)
1262            .with_expected_output_sizes(vec![900, 900, 200])
1263            .run();
1264
1265        // expect each buffer to be entirely full except the last one
1266        expect_buffer_layout(
1267            col_as_string_view("c0", output_batches.first().unwrap()),
1268            vec![
1269                ExpectedLayout {
1270                    len: 8192,
1271                    capacity: 8192,
1272                },
1273                ExpectedLayout {
1274                    len: 16384,
1275                    capacity: 16384,
1276                },
1277                ExpectedLayout {
1278                    len: 4224,
1279                    capacity: 32768,
1280                },
1281            ],
1282        );
1283    }
1284
1285    #[test]
1286    fn test_string_view_large_small() {
1287        // The strings are 37 bytes long, so each batch has 100 * 28 = 2800 bytes
1288        let mixed_batch = stringview_batch_repeated(
1289            200,
1290            [Some("This string is 28 bytes long"), Some("small string")],
1291        );
1292        // These strings aren't copied, this array has an 8k buffer
1293        let all_large = stringview_batch_repeated(
1294            50,
1295            [Some(
1296                "This buffer has only large strings in it so there are no buffer copies",
1297            )],
1298        );
1299
1300        let output_batches = Test::new("coalesce_string_view_large_small")
1301            // First allocated buffer is 8kb.
1302            // Appending five batches of 2800 bytes will use 2800 * 10 = 28kb (8kb, an 16kb and 32kbkb)
1303            .with_batch(mixed_batch.clone())
1304            .with_batch(mixed_batch.clone())
1305            .with_batch(all_large.clone())
1306            .with_batch(mixed_batch.clone())
1307            .with_batch(all_large.clone())
1308            .with_batch(mixed_batch.clone())
1309            .with_batch(mixed_batch.clone())
1310            .with_batch(all_large.clone())
1311            .with_batch(mixed_batch.clone())
1312            .with_batch(all_large.clone())
1313            .with_batch_size(8000)
1314            .with_expected_output_sizes(vec![1400])
1315            .run();
1316
1317        expect_buffer_layout(
1318            col_as_string_view("c0", output_batches.first().unwrap()),
1319            vec![
1320                ExpectedLayout {
1321                    len: 8190,
1322                    capacity: 8192,
1323                },
1324                ExpectedLayout {
1325                    len: 16366,
1326                    capacity: 16384,
1327                },
1328                ExpectedLayout {
1329                    len: 6244,
1330                    capacity: 32768,
1331                },
1332            ],
1333        );
1334    }
1335
1336    #[test]
1337    fn test_binary_view() {
1338        let values: Vec<Option<&[u8]>> = vec![
1339            Some(b"foo"),
1340            None,
1341            Some(b"A longer string that is more than 12 bytes"),
1342        ];
1343
1344        let binary_view =
1345            BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1346        let batch =
1347            RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1348
1349        Test::new("coalesce_binary_view")
1350            .with_batch(batch.clone())
1351            .with_batch(batch.clone())
1352            .with_batch_size(512)
1353            .with_expected_output_sizes(vec![512, 512, 512, 464])
1354            .run();
1355    }
1356
1357    #[test]
1358    fn test_binary_view_filtered() {
1359        let values: Vec<Option<&[u8]>> = vec![
1360            Some(b"foo"),
1361            None,
1362            Some(b"A longer string that is more than 12 bytes"),
1363        ];
1364
1365        let binary_view =
1366            BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1367        let batch =
1368            RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1369        let filter = sparse_filter(1000);
1370
1371        Test::new("coalesce_binary_view_filtered")
1372            .with_batch(batch.clone())
1373            .with_filter(filter.clone())
1374            .with_batch(batch)
1375            .with_filter(filter)
1376            .with_batch_size(256)
1377            .with_expected_output_sizes(vec![250])
1378            .run();
1379    }
1380
1381    #[test]
1382    fn test_binary_view_filtered_inline() {
1383        let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1384
1385        let binary_view =
1386            BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1387        let batch =
1388            RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1389        let filter = sparse_filter(1000);
1390
1391        Test::new("coalesce_binary_view_filtered_inline")
1392            .with_batch(batch.clone())
1393            .with_filter(filter.clone())
1394            .with_batch(batch)
1395            .with_filter(filter)
1396            .with_batch_size(300)
1397            .with_expected_output_sizes(vec![250])
1398            .run();
1399    }
1400
1401    #[test]
1402    fn test_string_view_filtered_inline() {
1403        let values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1404
1405        let string_view =
1406            StringViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1407        let batch =
1408            RecordBatch::try_from_iter(vec![("c0", Arc::new(string_view) as ArrayRef)]).unwrap();
1409        let filter = sparse_filter(1000);
1410
1411        Test::new("coalesce_string_view_filtered_inline")
1412            .with_batch(batch.clone())
1413            .with_filter(filter.clone())
1414            .with_batch(batch)
1415            .with_filter(filter)
1416            .with_batch_size(300)
1417            .with_expected_output_sizes(vec![250])
1418            .run();
1419    }
1420
1421    #[test]
1422    fn test_mixed_inline_binary_view_filtered() {
1423        let int_values =
1424            Int32Array::from_iter((0..1000).map(|v| if v % 5 == 0 { None } else { Some(v) }));
1425        let float_values = arrow_array::Float64Array::from_iter((0..1000).map(|v| Some(v as f64)));
1426        let binary_values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1427        let binary_view = BinaryViewArray::from_iter(
1428            std::iter::repeat(binary_values.iter()).flatten().take(1000),
1429        );
1430
1431        let batch = RecordBatch::try_from_iter(vec![
1432            ("i", Arc::new(int_values) as ArrayRef),
1433            ("f", Arc::new(float_values) as ArrayRef),
1434            ("b", Arc::new(binary_view) as ArrayRef),
1435        ])
1436        .unwrap();
1437
1438        let filter = sparse_filter(1000);
1439
1440        Test::new("coalesce_mixed_inline_binary_view_filtered")
1441            .with_batch(batch.clone())
1442            .with_filter(filter.clone())
1443            .with_batch(batch)
1444            .with_filter(filter)
1445            .with_batch_size(300)
1446            .with_expected_output_sizes(vec![250])
1447            .run();
1448    }
1449
1450    #[test]
1451    fn test_mixed_inline_string_view_filtered() {
1452        let int_values =
1453            Int32Array::from_iter((0..1000).map(|v| if v % 5 == 0 { None } else { Some(v) }));
1454        let float_values = arrow_array::Float64Array::from_iter((0..1000).map(|v| Some(v as f64)));
1455        let string_values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1456        let string_view = StringViewArray::from_iter(
1457            std::iter::repeat(string_values.iter()).flatten().take(1000),
1458        );
1459
1460        let batch = RecordBatch::try_from_iter(vec![
1461            ("i", Arc::new(int_values) as ArrayRef),
1462            ("f", Arc::new(float_values) as ArrayRef),
1463            ("s", Arc::new(string_view) as ArrayRef),
1464        ])
1465        .unwrap();
1466
1467        let filter = sparse_filter(1000);
1468
1469        Test::new("coalesce_mixed_inline_string_view_filtered")
1470            .with_batch(batch.clone())
1471            .with_filter(filter.clone())
1472            .with_batch(batch)
1473            .with_filter(filter)
1474            .with_batch_size(300)
1475            .with_expected_output_sizes(vec![250])
1476            .run();
1477    }
1478
1479    #[test]
1480    fn test_inline_binary_view_sparse() {
1481        // All-inline binary views (<=12 bytes) + 5% selectivity
1482        let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1483        let binary_view =
1484            BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1485        let batch =
1486            RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1487        let filter = very_sparse_filter(1000);
1488
1489        Test::new("inline_binary_view_sparse")
1490            .with_batch(batch.clone())
1491            .with_filter(filter.clone())
1492            .with_batch(batch)
1493            .with_filter(filter)
1494            .with_batch_size(1024)
1495            .with_expected_output_sizes(vec![100])
1496            .run();
1497    }
1498
1499    #[test]
1500    fn test_inline_string_view_sparse() {
1501        let values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1502        let string_view =
1503            StringViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1504        let batch =
1505            RecordBatch::try_from_iter(vec![("c0", Arc::new(string_view) as ArrayRef)]).unwrap();
1506        let filter = very_sparse_filter(1000);
1507
1508        Test::new("inline_string_view_sparse")
1509            .with_batch(batch.clone())
1510            .with_filter(filter.clone())
1511            .with_batch(batch)
1512            .with_filter(filter)
1513            .with_batch_size(1024)
1514            .with_expected_output_sizes(vec![100])
1515            .run();
1516    }
1517
1518    #[test]
1519    fn test_inline_mixed_sparse() {
1520        // Mixing primitives with inline views exercises materialized `Indices`
1521        // selection for both the primitive and the byte-view.
1522        let int_values =
1523            Int32Array::from_iter((0..1000).map(|v| if v % 5 == 0 { None } else { Some(v) }));
1524        let float_values = arrow_array::Float64Array::from_iter((0..1000).map(|v| Some(v as f64)));
1525        let string_values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1526        let string_view = StringViewArray::from_iter(
1527            std::iter::repeat(string_values.iter()).flatten().take(1000),
1528        );
1529        let binary_values: Vec<Option<&[u8]>> = vec![Some(b"x"), None, Some(b"abcdef")];
1530        let binary_view = BinaryViewArray::from_iter(
1531            std::iter::repeat(binary_values.iter()).flatten().take(1000),
1532        );
1533
1534        let batch = RecordBatch::try_from_iter(vec![
1535            ("i", Arc::new(int_values) as ArrayRef),
1536            ("f", Arc::new(float_values) as ArrayRef),
1537            ("s", Arc::new(string_view) as ArrayRef),
1538            ("b", Arc::new(binary_view) as ArrayRef),
1539        ])
1540        .unwrap();
1541        let filter = very_sparse_filter(1000);
1542
1543        Test::new("inline_mixed_sparse")
1544            .with_batch(batch.clone())
1545            .with_filter(filter.clone())
1546            .with_batch(batch)
1547            .with_filter(filter)
1548            .with_batch_size(1024)
1549            .with_expected_output_sizes(vec![100])
1550            .run();
1551    }
1552
1553    #[test]
1554    fn test_inline_crosses_target_batch_size() {
1555        // Each 1000-row batch selects 50 rows; target_batch_size 100
1556        // exercises intermediate batch flushing batch (not just the final
1557        // flush).
1558        let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1559        let make_batch = || {
1560            let binary_view =
1561                BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1562            RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap()
1563        };
1564        let filter = very_sparse_filter(1000);
1565
1566        Test::new("inline_crosses_target_batch_size")
1567            .with_batch(make_batch())
1568            .with_filter(filter.clone())
1569            .with_batch(make_batch())
1570            .with_filter(filter.clone())
1571            .with_batch(make_batch())
1572            .with_filter(filter)
1573            .with_batch_size(100)
1574            .with_expected_output_sizes(vec![100, 50])
1575            .run();
1576    }
1577
1578    #[test]
1579    fn test_inline_filter_rejects_filter_longer_than_batch() {
1580        let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), Some(b"bar")];
1581        let binary_view = BinaryViewArray::from_iter(values);
1582        let batch =
1583            RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1584        let filter = BooleanArray::from(vec![true, false, true]);
1585
1586        let mut coalescer = BatchCoalescer::new(batch.schema(), 100);
1587        let result = coalescer.push_batch_with_filter(batch, &filter);
1588        assert!(result.is_err());
1589        let err = result.unwrap_err().to_string();
1590        assert!(
1591            err.contains("Filter predicate of length 3 is larger than target array of length 2"),
1592            "unexpected error: {err}"
1593        );
1594    }
1595
1596    #[test]
1597    fn test_mixed_boolean_inline_string_view_filtered() {
1598        let bool_values = BooleanArray::from_iter((0..1000).map(|v| Some(v % 3 == 0)));
1599        let string_values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1600        let string_view = StringViewArray::from_iter(
1601            std::iter::repeat(string_values.iter()).flatten().take(1000),
1602        );
1603
1604        let batch = RecordBatch::try_from_iter(vec![
1605            ("b", Arc::new(bool_values) as ArrayRef),
1606            ("s", Arc::new(string_view) as ArrayRef),
1607        ])
1608        .unwrap();
1609
1610        let filter = sparse_filter(1000);
1611
1612        Test::new("coalesce_mixed_boolean_inline_string_view_filtered")
1613            .with_batch(batch.clone())
1614            .with_filter(filter.clone())
1615            .with_batch(batch)
1616            .with_filter(filter)
1617            .with_batch_size(300)
1618            .with_expected_output_sizes(vec![250])
1619            .run();
1620    }
1621
1622    #[test]
1623    fn test_filter_fast_path_schema_capability() {
1624        let supported = Arc::new(Schema::new(vec![
1625            Field::new("primitive", DataType::UInt32, false),
1626            Field::new("utf8_view", DataType::Utf8View, true),
1627            Field::new("binary_view", DataType::BinaryView, true),
1628        ]));
1629        let coalescer = BatchCoalescer::new(supported, 100);
1630        assert!(!coalescer.has_non_specialized_filter_columns);
1631
1632        let utf8 = Arc::new(Schema::new(vec![Field::new("utf8", DataType::Utf8, true)]));
1633        let coalescer = BatchCoalescer::new(utf8, 100);
1634        assert!(coalescer.has_non_specialized_filter_columns);
1635
1636        let boolean = Arc::new(Schema::new(vec![Field::new(
1637            "boolean",
1638            DataType::Boolean,
1639            true,
1640        )]));
1641        let coalescer = BatchCoalescer::new(boolean, 100);
1642        assert!(coalescer.has_non_specialized_filter_columns);
1643    }
1644
1645    #[derive(Debug, Clone, PartialEq)]
1646    struct ExpectedLayout {
1647        len: usize,
1648        capacity: usize,
1649    }
1650
1651    /// Asserts that the buffer layout of the specified StringViewArray matches the expected layout
1652    fn expect_buffer_layout(array: &StringViewArray, expected: Vec<ExpectedLayout>) {
1653        let actual = array
1654            .data_buffers()
1655            .iter()
1656            .map(|b| ExpectedLayout {
1657                len: b.len(),
1658                capacity: b.capacity(),
1659            })
1660            .collect::<Vec<_>>();
1661
1662        assert_eq!(
1663            actual, expected,
1664            "Expected buffer layout {expected:#?} but got {actual:#?}"
1665        );
1666    }
1667
1668    /// Test for [`BatchCoalescer`]
1669    ///
1670    /// Pushes the input batches to the coalescer and verifies that the resulting
1671    /// batches have the
1672    /// 1. expected number of rows
1673    /// 2. The same results when the batches are filtered using the filter kernel
1674    #[derive(Debug, Clone)]
1675    struct Test {
1676        /// A human readable name to assist in debugging
1677        name: String,
1678        /// Batches to feed to the coalescer.
1679        input_batches: Vec<RecordBatch>,
1680        /// Filters to apply to the corresponding input batches.
1681        ///
1682        /// If there are no filters for the input batches, the batch will be
1683        /// pushed as is.
1684        filters: Vec<BooleanArray>,
1685        /// The schema. If not provided, the first batch's schema is used.
1686        schema: Option<SchemaRef>,
1687        /// Expected output sizes of the resulting batches
1688        expected_output_sizes: Vec<usize>,
1689        /// target batch size (default to 1024)
1690        target_batch_size: usize,
1691    }
1692
1693    impl Default for Test {
1694        fn default() -> Self {
1695            Self {
1696                name: "".to_string(),
1697                input_batches: vec![],
1698                filters: vec![],
1699                schema: None,
1700                expected_output_sizes: vec![],
1701                target_batch_size: 1024,
1702            }
1703        }
1704    }
1705
1706    impl Test {
1707        fn new(name: impl Into<String>) -> Self {
1708            Self {
1709                name: name.into(),
1710                ..Self::default()
1711            }
1712        }
1713
1714        /// Append the description to the test name
1715        fn with_description(mut self, description: &str) -> Self {
1716            self.name.push_str(": ");
1717            self.name.push_str(description);
1718            self
1719        }
1720
1721        /// Set the target batch size
1722        fn with_batch_size(mut self, target_batch_size: usize) -> Self {
1723            self.target_batch_size = target_batch_size;
1724            self
1725        }
1726
1727        /// Extend the input batches with `batch`
1728        fn with_batch(mut self, batch: RecordBatch) -> Self {
1729            self.input_batches.push(batch);
1730            self
1731        }
1732
1733        /// Extend the filters with `filter`
1734        fn with_filter(mut self, filter: BooleanArray) -> Self {
1735            self.filters.push(filter);
1736            self
1737        }
1738
1739        /// Replaces the input batches with `batches`
1740        fn with_batches(mut self, batches: impl IntoIterator<Item = RecordBatch>) -> Self {
1741            self.input_batches = batches.into_iter().collect();
1742            self
1743        }
1744
1745        /// Specifies the schema for the test
1746        fn with_schema(mut self, schema: SchemaRef) -> Self {
1747            self.schema = Some(schema);
1748            self
1749        }
1750
1751        /// Extends `sizes` to expected output sizes
1752        fn with_expected_output_sizes(mut self, sizes: impl IntoIterator<Item = usize>) -> Self {
1753            self.expected_output_sizes.extend(sizes);
1754            self
1755        }
1756
1757        /// Runs the test -- see documentation on [`Test`] for details
1758        ///
1759        /// Returns the resulting output batches
1760        fn run(self) -> Vec<RecordBatch> {
1761            // Test several permutations of input batches:
1762            // 1. Removing nulls from some batches (test non-null fast paths)
1763            // 2. Empty batches
1764            // 3. One column (from the batch)
1765            let mut extra_tests = vec![];
1766            extra_tests.push(self.clone().make_half_non_nullable());
1767            extra_tests.push(self.clone().insert_empty_batches());
1768            let single_column_tests = self.make_single_column_tests();
1769            for test in single_column_tests {
1770                extra_tests.push(test.clone().make_half_non_nullable());
1771                extra_tests.push(test);
1772            }
1773
1774            // Run original test case first, so any obvious errors are caught
1775            // by an easier to understand test case
1776            let results = self.run_inner();
1777            // Run the extra cases to expand coverage
1778            for extra in extra_tests {
1779                extra.run_inner();
1780            }
1781
1782            results
1783        }
1784
1785        /// Runs the current test instance
1786        fn run_inner(self) -> Vec<RecordBatch> {
1787            let expected_output = self.expected_output();
1788            let schema = self.schema();
1789
1790            let Self {
1791                name,
1792                input_batches,
1793                filters,
1794                schema: _,
1795                target_batch_size,
1796                expected_output_sizes,
1797            } = self;
1798
1799            println!("Running test '{name}'");
1800
1801            let had_input = input_batches.iter().any(|b| b.num_rows() > 0);
1802
1803            let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), target_batch_size);
1804
1805            // feed input batches and filters to the coalescer
1806            let mut filters = filters.into_iter();
1807            for batch in input_batches {
1808                if let Some(filter) = filters.next() {
1809                    coalescer.push_batch_with_filter(batch, &filter).unwrap();
1810                } else {
1811                    coalescer.push_batch(batch).unwrap();
1812                }
1813            }
1814            assert_eq!(schema, coalescer.schema());
1815
1816            if had_input {
1817                assert!(!coalescer.is_empty(), "Coalescer should not be empty");
1818            } else {
1819                assert!(coalescer.is_empty(), "Coalescer should be empty");
1820            }
1821
1822            coalescer.finish_buffered_batch().unwrap();
1823            if had_input {
1824                assert!(
1825                    coalescer.has_completed_batch(),
1826                    "Coalescer should have completed batches"
1827                );
1828            }
1829
1830            let mut output_batches = vec![];
1831            while let Some(batch) = coalescer.next_completed_batch() {
1832                output_batches.push(batch);
1833            }
1834
1835            // make sure we got the expected number of output batches and content
1836            let mut starting_idx = 0;
1837            let actual_output_sizes: Vec<usize> =
1838                output_batches.iter().map(|b| b.num_rows()).collect();
1839            assert_eq!(
1840                expected_output_sizes, actual_output_sizes,
1841                "Unexpected number of rows in output batches\n\
1842                Expected\n{expected_output_sizes:#?}\nActual:{actual_output_sizes:#?}"
1843            );
1844            let iter = expected_output_sizes
1845                .iter()
1846                .zip(output_batches.iter())
1847                .enumerate();
1848
1849            // Verify that the actual contents of each output batch matches the expected output
1850            for (i, (expected_size, batch)) in iter {
1851                // compare the contents of the batch after normalization (using
1852                // `==` compares the underlying memory layout too)
1853                let expected_batch = expected_output.slice(starting_idx, *expected_size);
1854                let expected_batch = normalize_batch(expected_batch);
1855                let batch = normalize_batch(batch.clone());
1856                assert_eq!(
1857                    expected_batch, batch,
1858                    "Unexpected content in batch {i}:\
1859                    \n\nExpected:\n{expected_batch:#?}\n\nActual:\n{batch:#?}"
1860                );
1861                starting_idx += *expected_size;
1862            }
1863            output_batches
1864        }
1865
1866        /// Return the expected output schema. If not overridden by `with_schema`, it
1867        /// returns the schema of the first input batch.
1868        fn schema(&self) -> SchemaRef {
1869            self.schema
1870                .clone()
1871                .unwrap_or_else(|| Arc::clone(&self.input_batches[0].schema()))
1872        }
1873
1874        /// Returns the expected output as a single `RecordBatch`
1875        fn expected_output(&self) -> RecordBatch {
1876            let schema = self.schema();
1877            if self.filters.is_empty() {
1878                return concat_batches(&schema, &self.input_batches).unwrap();
1879            }
1880
1881            let mut filters = self.filters.iter();
1882            let filtered_batches = self
1883                .input_batches
1884                .iter()
1885                .map(|batch| {
1886                    if let Some(filter) = filters.next() {
1887                        filter_record_batch(batch, filter).unwrap()
1888                    } else {
1889                        batch.clone()
1890                    }
1891                })
1892                .collect::<Vec<_>>();
1893            concat_batches(&schema, &filtered_batches).unwrap()
1894        }
1895
1896        /// Return a copy of self where every other batch has had its nulls removed
1897        /// (there are often fast paths that are used when there are no nulls)
1898        fn make_half_non_nullable(mut self) -> Self {
1899            // remove the nulls from every other batch
1900            self.input_batches = self
1901                .input_batches
1902                .iter()
1903                .enumerate()
1904                .map(|(i, batch)| {
1905                    if i % 2 == 1 {
1906                        batch.clone()
1907                    } else {
1908                        Self::remove_nulls_from_batch(batch)
1909                    }
1910                })
1911                .collect();
1912            self.with_description("non-nullable")
1913        }
1914
1915        /// Insert several empty batches into the input before each existing input
1916        fn insert_empty_batches(mut self) -> Self {
1917            let empty_batch = RecordBatch::new_empty(self.schema());
1918            self.input_batches = self
1919                .input_batches
1920                .into_iter()
1921                .flat_map(|batch| [empty_batch.clone(), batch])
1922                .collect();
1923            let empty_filters = BooleanArray::builder(0).finish();
1924            self.filters = self
1925                .filters
1926                .into_iter()
1927                .flat_map(|filter| [empty_filters.clone(), filter])
1928                .collect();
1929            self.with_description("empty batches inserted")
1930        }
1931
1932        /// Sets one batch to be non-nullable by removing nulls from all columns
1933        fn remove_nulls_from_batch(batch: &RecordBatch) -> RecordBatch {
1934            let new_columns = batch
1935                .columns()
1936                .iter()
1937                .map(Self::remove_nulls_from_array)
1938                .collect::<Vec<_>>();
1939            let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
1940            RecordBatch::try_new_with_options(batch.schema(), new_columns, &options).unwrap()
1941        }
1942
1943        fn remove_nulls_from_array(array: &ArrayRef) -> ArrayRef {
1944            make_array(array.to_data().into_builder().nulls(None).build().unwrap())
1945        }
1946
1947        /// Returns a set of tests where each test that is the sae as self, but
1948        /// has a single column from the original input batch
1949        ///
1950        /// This can be useful to single column optimizations, specifically
1951        /// filter optimization.
1952        fn make_single_column_tests(&self) -> Vec<Self> {
1953            let original_schema = self.schema();
1954            let mut new_tests = vec![];
1955            for column in original_schema.fields() {
1956                let single_column_schema = Arc::new(Schema::new(vec![column.clone()]));
1957
1958                let single_column_batches = self.input_batches.iter().map(|batch| {
1959                    let single_column = batch.column_by_name(column.name()).unwrap();
1960                    RecordBatch::try_new(
1961                        Arc::clone(&single_column_schema),
1962                        vec![single_column.clone()],
1963                    )
1964                    .unwrap()
1965                });
1966
1967                let single_column_test = self
1968                    .clone()
1969                    .with_schema(Arc::clone(&single_column_schema))
1970                    .with_batches(single_column_batches)
1971                    .with_description("single column")
1972                    .with_description(column.name());
1973
1974                new_tests.push(single_column_test);
1975            }
1976            new_tests
1977        }
1978    }
1979
1980    /// Return a RecordBatch with a UInt32Array with the specified range and
1981    /// every third value is null.
1982    fn uint32_batch<T: std::iter::Iterator<Item = u32>>(range: T) -> RecordBatch {
1983        let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, true)]));
1984
1985        let array = UInt32Array::from_iter(range.map(|i| if i % 3 == 0 { None } else { Some(i) }));
1986        RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
1987    }
1988
1989    /// Return a RecordBatch with a UInt32Array with no nulls specified range
1990    fn uint32_batch_non_null<T: std::iter::Iterator<Item = u32>>(range: T) -> RecordBatch {
1991        let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]));
1992
1993        let array = UInt32Array::from_iter_values(range);
1994        RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
1995    }
1996
1997    /// Return a RecordBatch with a UInt64Array with no nulls specified range
1998    fn uint64_batch_non_null<T: std::iter::Iterator<Item = u64>>(range: T) -> RecordBatch {
1999        let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt64, false)]));
2000
2001        let array = UInt64Array::from_iter_values(range);
2002        RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2003    }
2004
2005    /// Return a RecordBatch with a StringArrary with values `value0`, `value1`, ...
2006    /// and every third value is `None`.
2007    fn utf8_batch(range: Range<u32>) -> RecordBatch {
2008        let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Utf8, true)]));
2009
2010        let array = StringArray::from_iter(range.map(|i| {
2011            if i % 3 == 0 {
2012                None
2013            } else {
2014                Some(format!("value{i}"))
2015            }
2016        }));
2017
2018        RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2019    }
2020
2021    /// Return a RecordBatch with a StringViewArray with (only) the specified values
2022    fn stringview_batch<'a>(values: impl IntoIterator<Item = Option<&'a str>>) -> RecordBatch {
2023        let schema = Arc::new(Schema::new(vec![Field::new(
2024            "c0",
2025            DataType::Utf8View,
2026            false,
2027        )]));
2028
2029        let array = StringViewArray::from_iter(values);
2030        RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2031    }
2032
2033    /// Return a RecordBatch with a StringViewArray with num_rows by repeating
2034    /// values over and over.
2035    fn stringview_batch_repeated<'a>(
2036        num_rows: usize,
2037        values: impl IntoIterator<Item = Option<&'a str>>,
2038    ) -> RecordBatch {
2039        let schema = Arc::new(Schema::new(vec![Field::new(
2040            "c0",
2041            DataType::Utf8View,
2042            true,
2043        )]));
2044
2045        // Repeat the values to a total of num_rows
2046        let values: Vec<_> = values.into_iter().collect();
2047        let values_iter = std::iter::repeat(values.iter())
2048            .flatten()
2049            .cloned()
2050            .take(num_rows);
2051
2052        let mut builder = StringViewBuilder::with_capacity(100).with_fixed_block_size(8192);
2053        for val in values_iter {
2054            builder.append_option(val);
2055        }
2056
2057        let array = builder.finish();
2058        RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2059    }
2060
2061    /// Return a RecordBatch of 100 rows
2062    fn multi_column_batch(range: Range<i32>) -> RecordBatch {
2063        let int64_array = Int64Array::from_iter(
2064            range
2065                .clone()
2066                .map(|v| if v % 5 == 0 { None } else { Some(v as i64) }),
2067        );
2068        let string_view_array = StringViewArray::from_iter(range.clone().map(|v| {
2069            if v % 5 == 0 {
2070                None
2071            } else if v % 7 == 0 {
2072                Some(format!("This is a string longer than 12 bytes{v}"))
2073            } else {
2074                Some(format!("Short {v}"))
2075            }
2076        }));
2077        let string_array = StringArray::from_iter(range.clone().map(|v| {
2078            if v % 11 == 0 {
2079                None
2080            } else {
2081                Some(format!("Value {v}"))
2082            }
2083        }));
2084        let timestamp_array = TimestampNanosecondArray::from_iter(range.map(|v| {
2085            if v % 3 == 0 {
2086                None
2087            } else {
2088                Some(v as i64 * 1000) // simulate a timestamp in milliseconds
2089            }
2090        }))
2091        .with_timezone("America/New_York");
2092
2093        RecordBatch::try_from_iter(vec![
2094            ("int64", Arc::new(int64_array) as ArrayRef),
2095            ("stringview", Arc::new(string_view_array) as ArrayRef),
2096            ("string", Arc::new(string_array) as ArrayRef),
2097            ("timestamp", Arc::new(timestamp_array) as ArrayRef),
2098        ])
2099        .unwrap()
2100    }
2101
2102    /// Return a boolean array that filters out randomly selected rows
2103    /// from the input batch with a `selectivity`.
2104    ///
2105    /// For example a `selectivity` of 0.1 will filter out
2106    /// 90% of the rows.
2107    #[derive(Debug)]
2108    struct RandomFilterBuilder {
2109        /// Number of rows to add to each filter
2110        num_rows: usize,
2111        /// selectivity of the filter (between 0.0 and 1.0)
2112        /// 0 selects no rows, 1.0 selects all rows
2113        selectivity: f64,
2114        /// seed for random number generator, increases by one each time
2115        /// `next_filter` is called
2116        seed: u64,
2117    }
2118    impl RandomFilterBuilder {
2119        /// Build the next filter with the current seed and increment the seed
2120        /// by one.
2121        fn next_filter(&mut self) -> BooleanArray {
2122            assert!(self.selectivity >= 0.0 && self.selectivity <= 1.0);
2123            let mut rng = rand::rngs::StdRng::seed_from_u64(self.seed);
2124            self.seed += 1;
2125            BooleanArray::from_iter(
2126                (0..self.num_rows)
2127                    .map(|_| rng.random_bool(self.selectivity))
2128                    .map(Some),
2129            )
2130        }
2131    }
2132
2133    /// Returns the named column as a StringViewArray
2134    fn col_as_string_view<'b>(name: &str, batch: &'b RecordBatch) -> &'b StringViewArray {
2135        batch
2136            .column_by_name(name)
2137            .expect("column not found")
2138            .as_string_view_opt()
2139            .expect("column is not a string view")
2140    }
2141
2142    /// Filter that selects 12.5% of the rows (`idx % 8 == 0`)
2143    fn sparse_filter(len: usize) -> BooleanArray {
2144        BooleanArray::from_iter((0..len).map(|idx| Some(idx % 8 == 0)))
2145    }
2146
2147    /// Filter that selects 5% of the rows (`idx % 20 == 0`)
2148    ///
2149    /// Different code paths are taken for very sparse filters
2150    fn very_sparse_filter(len: usize) -> BooleanArray {
2151        BooleanArray::from_iter((0..len).map(|idx| Some(idx % 20 == 0)))
2152    }
2153
2154    /// Normalize the `RecordBatch` so that the memory layout is consistent
2155    /// (e.g. StringArray is compacted).
2156    fn normalize_batch(batch: RecordBatch) -> RecordBatch {
2157        // Only need to normalize StringViews (as == also tests for memory layout)
2158        let (schema, mut columns, row_count) = batch.into_parts();
2159
2160        for column in columns.iter_mut() {
2161            if let Some(string_view) = column.as_string_view_opt() {
2162                // Re-create the StringViewArray to ensure memory layout is
2163                // consistent
2164                let mut builder = StringViewBuilder::new();
2165                for s in string_view.iter() {
2166                    builder.append_option(s);
2167                }
2168                *column = Arc::new(builder.finish());
2169                continue;
2170            }
2171
2172            if let Some(binary_view) = column.as_binary_view_opt() {
2173                *column = Arc::new(BinaryViewArray::from_iter(binary_view.iter()));
2174            }
2175        }
2176
2177        let options = RecordBatchOptions::new().with_row_count(Some(row_count));
2178        RecordBatch::try_new_with_options(schema, columns, &options).unwrap()
2179    }
2180
2181    /// Helper function to create a test batch with specified number of rows
2182    fn create_test_batch(num_rows: usize) -> RecordBatch {
2183        let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)]));
2184        let array = Int32Array::from_iter_values(0..num_rows as i32);
2185        RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
2186    }
2187    #[test]
2188    fn test_biggest_coalesce_batch_size_none_default() {
2189        // Test that default behavior (None) coalesces all batches
2190        let mut coalescer = BatchCoalescer::new(
2191            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2192            100,
2193        );
2194
2195        // Push a large batch (1000 rows) - should be coalesced normally
2196        let large_batch = create_test_batch(1000);
2197        coalescer.push_batch(large_batch).unwrap();
2198
2199        // Should produce multiple batches of target size (100)
2200        let mut output_batches = vec![];
2201        while let Some(batch) = coalescer.next_completed_batch() {
2202            output_batches.push(batch);
2203        }
2204
2205        coalescer.finish_buffered_batch().unwrap();
2206        while let Some(batch) = coalescer.next_completed_batch() {
2207            output_batches.push(batch);
2208        }
2209
2210        // Should have 10 batches of 100 rows each
2211        assert_eq!(output_batches.len(), 10);
2212        for batch in output_batches {
2213            assert_eq!(batch.num_rows(), 100);
2214        }
2215    }
2216
2217    #[test]
2218    fn test_biggest_coalesce_batch_size_bypass_large_batch() {
2219        // Test that batches larger than biggest_coalesce_batch_size bypass coalescing
2220        let mut coalescer = BatchCoalescer::new(
2221            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2222            100,
2223        );
2224        coalescer.set_biggest_coalesce_batch_size(Some(500));
2225
2226        // Push a large batch (1000 rows) - should bypass coalescing
2227        let large_batch = create_test_batch(1000);
2228        coalescer.push_batch(large_batch.clone()).unwrap();
2229
2230        // Should have one completed batch immediately (the original large batch)
2231        assert!(coalescer.has_completed_batch());
2232        let output_batch = coalescer.next_completed_batch().unwrap();
2233        assert_eq!(output_batch.num_rows(), 1000);
2234
2235        // Should be no more completed batches
2236        assert!(!coalescer.has_completed_batch());
2237        assert_eq!(coalescer.get_buffered_rows(), 0);
2238    }
2239
2240    #[test]
2241    fn test_biggest_coalesce_batch_size_coalesce_small_batch() {
2242        // Test that batches smaller than biggest_coalesce_batch_size are coalesced normally
2243        let mut coalescer = BatchCoalescer::new(
2244            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2245            100,
2246        )
2247        .with_biggest_coalesce_batch_size(Some(500));
2248
2249        // Push small batches that should be coalesced
2250        let small_batch = create_test_batch(50);
2251        coalescer.push_batch(small_batch.clone()).unwrap();
2252
2253        // Should not have completed batch yet (only 50 rows, target is 100)
2254        assert!(!coalescer.has_completed_batch());
2255        assert_eq!(coalescer.get_buffered_rows(), 50);
2256
2257        // Push another small batch
2258        coalescer.push_batch(small_batch).unwrap();
2259
2260        // Now should have a completed batch (100 rows total)
2261        assert!(coalescer.has_completed_batch());
2262        let output_batch = coalescer.next_completed_batch().unwrap();
2263        let size = output_batch
2264            .column(0)
2265            .as_primitive::<Int32Type>()
2266            .get_buffer_memory_size();
2267        assert_eq!(size, 400); // 100 rows * 4 bytes each
2268        assert_eq!(output_batch.num_rows(), 100);
2269
2270        assert_eq!(coalescer.get_buffered_rows(), 0);
2271    }
2272
2273    #[test]
2274    fn test_biggest_coalesce_batch_size_equal_boundary() {
2275        // Test behavior when batch size equals biggest_coalesce_batch_size
2276        let mut coalescer = BatchCoalescer::new(
2277            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2278            100,
2279        );
2280        coalescer.set_biggest_coalesce_batch_size(Some(500));
2281
2282        // Push a batch exactly equal to the limit
2283        let boundary_batch = create_test_batch(500);
2284        coalescer.push_batch(boundary_batch).unwrap();
2285
2286        // Should be coalesced (not bypass) since it's equal, not greater
2287        let mut output_count = 0;
2288        while coalescer.next_completed_batch().is_some() {
2289            output_count += 1;
2290        }
2291
2292        coalescer.finish_buffered_batch().unwrap();
2293        while coalescer.next_completed_batch().is_some() {
2294            output_count += 1;
2295        }
2296
2297        // Should have 5 batches of 100 rows each
2298        assert_eq!(output_count, 5);
2299    }
2300
2301    #[test]
2302    fn test_biggest_coalesce_batch_size_first_large_then_consecutive_bypass() {
2303        // Test the new consecutive large batch bypass behavior
2304        // Pattern: small batches -> first large batch (coalesced) -> consecutive large batches (bypass)
2305        let mut coalescer = BatchCoalescer::new(
2306            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2307            100,
2308        );
2309        coalescer.set_biggest_coalesce_batch_size(Some(200));
2310
2311        let small_batch = create_test_batch(50);
2312
2313        // Push small batch first to create buffered data
2314        coalescer.push_batch(small_batch).unwrap();
2315        assert_eq!(coalescer.get_buffered_rows(), 50);
2316        assert!(!coalescer.has_completed_batch());
2317
2318        // Push first large batch - should go through normal coalescing due to buffered data
2319        let large_batch1 = create_test_batch(250);
2320        coalescer.push_batch(large_batch1).unwrap();
2321
2322        // 50 + 250 = 300 -> 3 complete batches of 100, 0 rows buffered
2323        let mut completed_batches = vec![];
2324        while let Some(batch) = coalescer.next_completed_batch() {
2325            completed_batches.push(batch);
2326        }
2327        assert_eq!(completed_batches.len(), 3);
2328        assert_eq!(coalescer.get_buffered_rows(), 0);
2329
2330        // Now push consecutive large batches - they should bypass
2331        let large_batch2 = create_test_batch(300);
2332        let large_batch3 = create_test_batch(400);
2333
2334        // Push second large batch - should bypass since it's consecutive and buffer is empty
2335        coalescer.push_batch(large_batch2).unwrap();
2336        assert!(coalescer.has_completed_batch());
2337        let output = coalescer.next_completed_batch().unwrap();
2338        assert_eq!(output.num_rows(), 300); // bypassed with original size
2339        assert_eq!(coalescer.get_buffered_rows(), 0);
2340
2341        // Push third large batch - should also bypass
2342        coalescer.push_batch(large_batch3).unwrap();
2343        assert!(coalescer.has_completed_batch());
2344        let output = coalescer.next_completed_batch().unwrap();
2345        assert_eq!(output.num_rows(), 400); // bypassed with original size
2346        assert_eq!(coalescer.get_buffered_rows(), 0);
2347    }
2348
2349    #[test]
2350    fn test_biggest_coalesce_batch_size_empty_batch() {
2351        // Test that empty batches don't trigger the bypass logic
2352        let mut coalescer = BatchCoalescer::new(
2353            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2354            100,
2355        );
2356        coalescer.set_biggest_coalesce_batch_size(Some(50));
2357
2358        let empty_batch = create_test_batch(0);
2359        coalescer.push_batch(empty_batch).unwrap();
2360
2361        // Empty batch should be handled normally (no effect)
2362        assert!(!coalescer.has_completed_batch());
2363        assert_eq!(coalescer.get_buffered_rows(), 0);
2364    }
2365
2366    #[test]
2367    fn test_biggest_coalesce_batch_size_with_buffered_data_no_bypass() {
2368        // Test that when there is buffered data, large batches do NOT bypass (unless consecutive)
2369        let mut coalescer = BatchCoalescer::new(
2370            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2371            100,
2372        );
2373        coalescer.set_biggest_coalesce_batch_size(Some(200));
2374
2375        // Add some buffered data first
2376        let small_batch = create_test_batch(30);
2377        coalescer.push_batch(small_batch.clone()).unwrap();
2378        coalescer.push_batch(small_batch).unwrap();
2379        assert_eq!(coalescer.get_buffered_rows(), 60);
2380
2381        // Push large batch that would normally bypass, but shouldn't because buffered_rows > 0
2382        let large_batch = create_test_batch(250);
2383        coalescer.push_batch(large_batch).unwrap();
2384
2385        // The large batch should be processed through normal coalescing logic
2386        // Total: 60 (buffered) + 250 (new) = 310 rows
2387        // Output: 3 complete batches of 100 rows each, 10 rows remain buffered
2388
2389        let mut completed_batches = vec![];
2390        while let Some(batch) = coalescer.next_completed_batch() {
2391            completed_batches.push(batch);
2392        }
2393
2394        assert_eq!(completed_batches.len(), 3);
2395        for batch in &completed_batches {
2396            assert_eq!(batch.num_rows(), 100);
2397        }
2398        assert_eq!(coalescer.get_buffered_rows(), 10);
2399    }
2400
2401    #[test]
2402    fn test_biggest_coalesce_batch_size_zero_limit() {
2403        // Test edge case where limit is 0 (all batches bypass when no buffered data)
2404        let mut coalescer = BatchCoalescer::new(
2405            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2406            100,
2407        );
2408        coalescer.set_biggest_coalesce_batch_size(Some(0));
2409
2410        // Even a 1-row batch should bypass when there's no buffered data
2411        let tiny_batch = create_test_batch(1);
2412        coalescer.push_batch(tiny_batch).unwrap();
2413
2414        assert!(coalescer.has_completed_batch());
2415        let output = coalescer.next_completed_batch().unwrap();
2416        assert_eq!(output.num_rows(), 1);
2417    }
2418
2419    #[test]
2420    fn test_biggest_coalesce_batch_size_bypass_only_when_no_buffer() {
2421        // Test that bypass only occurs when buffered_rows == 0
2422        let mut coalescer = BatchCoalescer::new(
2423            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2424            100,
2425        );
2426        coalescer.set_biggest_coalesce_batch_size(Some(200));
2427
2428        // First, push a large batch with no buffered data - should bypass
2429        let large_batch = create_test_batch(300);
2430        coalescer.push_batch(large_batch.clone()).unwrap();
2431
2432        assert!(coalescer.has_completed_batch());
2433        let output = coalescer.next_completed_batch().unwrap();
2434        assert_eq!(output.num_rows(), 300); // bypassed
2435        assert_eq!(coalescer.get_buffered_rows(), 0);
2436
2437        // Now add some buffered data
2438        let small_batch = create_test_batch(50);
2439        coalescer.push_batch(small_batch).unwrap();
2440        assert_eq!(coalescer.get_buffered_rows(), 50);
2441
2442        // Push the same large batch again - should NOT bypass this time (not consecutive)
2443        coalescer.push_batch(large_batch).unwrap();
2444
2445        // Should process through normal coalescing: 50 + 300 = 350 rows
2446        // Output: 3 complete batches of 100 rows, 50 rows buffered
2447        let mut completed_batches = vec![];
2448        while let Some(batch) = coalescer.next_completed_batch() {
2449            completed_batches.push(batch);
2450        }
2451
2452        assert_eq!(completed_batches.len(), 3);
2453        for batch in &completed_batches {
2454            assert_eq!(batch.num_rows(), 100);
2455        }
2456        assert_eq!(coalescer.get_buffered_rows(), 50);
2457    }
2458
2459    #[test]
2460    fn test_biggest_coalesce_batch_size_consecutive_large_batches_scenario() {
2461        // Test your exact scenario: 20, 20, 30, 700, 600, 700, 900, 700, 600
2462        let mut coalescer = BatchCoalescer::new(
2463            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2464            1000,
2465        );
2466        coalescer.set_biggest_coalesce_batch_size(Some(500));
2467
2468        // Push small batches first
2469        coalescer.push_batch(create_test_batch(20)).unwrap();
2470        coalescer.push_batch(create_test_batch(20)).unwrap();
2471        coalescer.push_batch(create_test_batch(30)).unwrap();
2472
2473        assert_eq!(coalescer.get_buffered_rows(), 70);
2474        assert!(!coalescer.has_completed_batch());
2475
2476        // Push first large batch (700) - should coalesce due to buffered data
2477        coalescer.push_batch(create_test_batch(700)).unwrap();
2478
2479        // 70 + 700 = 770 rows, not enough for 1000, so all stay buffered
2480        assert_eq!(coalescer.get_buffered_rows(), 770);
2481        assert!(!coalescer.has_completed_batch());
2482
2483        // Push second large batch (600) - should bypass since previous was large
2484        coalescer.push_batch(create_test_batch(600)).unwrap();
2485
2486        // Should flush buffer (770 rows) and bypass the 600
2487        let mut outputs = vec![];
2488        while let Some(batch) = coalescer.next_completed_batch() {
2489            outputs.push(batch);
2490        }
2491        assert_eq!(outputs.len(), 2); // one flushed buffer batch (770) + one bypassed (600)
2492        assert_eq!(outputs[0].num_rows(), 770);
2493        assert_eq!(outputs[1].num_rows(), 600);
2494        assert_eq!(coalescer.get_buffered_rows(), 0);
2495
2496        // Push remaining large batches - should all bypass
2497        let remaining_batches = [700, 900, 700, 600];
2498        for &size in &remaining_batches {
2499            coalescer.push_batch(create_test_batch(size)).unwrap();
2500
2501            assert!(coalescer.has_completed_batch());
2502            let output = coalescer.next_completed_batch().unwrap();
2503            assert_eq!(output.num_rows(), size);
2504            assert_eq!(coalescer.get_buffered_rows(), 0);
2505        }
2506    }
2507
2508    #[test]
2509    fn test_biggest_coalesce_batch_size_truly_consecutive_large_bypass() {
2510        // Test truly consecutive large batches that should all bypass
2511        // This test ensures buffer is completely empty between large batches
2512        let mut coalescer = BatchCoalescer::new(
2513            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2514            100,
2515        );
2516        coalescer.set_biggest_coalesce_batch_size(Some(200));
2517
2518        // Push consecutive large batches with no prior buffered data
2519        let large_batches = vec![
2520            create_test_batch(300),
2521            create_test_batch(400),
2522            create_test_batch(350),
2523            create_test_batch(500),
2524        ];
2525
2526        let mut all_outputs = vec![];
2527
2528        for (i, large_batch) in large_batches.into_iter().enumerate() {
2529            let expected_size = large_batch.num_rows();
2530
2531            // Buffer should be empty before each large batch
2532            assert_eq!(
2533                coalescer.get_buffered_rows(),
2534                0,
2535                "Buffer should be empty before batch {}",
2536                i
2537            );
2538
2539            coalescer.push_batch(large_batch).unwrap();
2540
2541            // Each large batch should bypass and produce exactly one output batch
2542            assert!(
2543                coalescer.has_completed_batch(),
2544                "Should have completed batch after pushing batch {}",
2545                i
2546            );
2547
2548            let output = coalescer.next_completed_batch().unwrap();
2549            assert_eq!(
2550                output.num_rows(),
2551                expected_size,
2552                "Batch {} should have bypassed with original size",
2553                i
2554            );
2555
2556            // Should be no more batches and buffer should be empty
2557            assert!(
2558                !coalescer.has_completed_batch(),
2559                "Should have no more completed batches after batch {}",
2560                i
2561            );
2562            assert_eq!(
2563                coalescer.get_buffered_rows(),
2564                0,
2565                "Buffer should be empty after batch {}",
2566                i
2567            );
2568
2569            all_outputs.push(output);
2570        }
2571
2572        // Verify we got exactly 4 output batches with original sizes
2573        assert_eq!(all_outputs.len(), 4);
2574        assert_eq!(all_outputs[0].num_rows(), 300);
2575        assert_eq!(all_outputs[1].num_rows(), 400);
2576        assert_eq!(all_outputs[2].num_rows(), 350);
2577        assert_eq!(all_outputs[3].num_rows(), 500);
2578    }
2579
2580    #[test]
2581    fn test_biggest_coalesce_batch_size_reset_consecutive_on_small_batch() {
2582        // Test that small batches reset the consecutive large batch tracking
2583        let mut coalescer = BatchCoalescer::new(
2584            Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2585            100,
2586        );
2587        coalescer.set_biggest_coalesce_batch_size(Some(200));
2588
2589        // Push first large batch - should bypass (no buffered data)
2590        coalescer.push_batch(create_test_batch(300)).unwrap();
2591        let output = coalescer.next_completed_batch().unwrap();
2592        assert_eq!(output.num_rows(), 300);
2593
2594        // Push second large batch - should bypass (consecutive)
2595        coalescer.push_batch(create_test_batch(400)).unwrap();
2596        let output = coalescer.next_completed_batch().unwrap();
2597        assert_eq!(output.num_rows(), 400);
2598
2599        // Push small batch - resets consecutive tracking
2600        coalescer.push_batch(create_test_batch(50)).unwrap();
2601        assert_eq!(coalescer.get_buffered_rows(), 50);
2602
2603        // Push large batch again - should NOT bypass due to buffered data
2604        coalescer.push_batch(create_test_batch(350)).unwrap();
2605
2606        // Should coalesce: 50 + 350 = 400 -> 4 complete batches of 100
2607        let mut outputs = vec![];
2608        while let Some(batch) = coalescer.next_completed_batch() {
2609            outputs.push(batch);
2610        }
2611        assert_eq!(outputs.len(), 4);
2612        for batch in outputs {
2613            assert_eq!(batch.num_rows(), 100);
2614        }
2615        assert_eq!(coalescer.get_buffered_rows(), 0);
2616    }
2617
2618    #[test]
2619    fn test_coalasce_push_batch_with_indices() {
2620        const MID_POINT: u32 = 2333;
2621        const TOTAL_ROWS: u32 = 23333;
2622        let batch1 = uint32_batch_non_null(0..MID_POINT);
2623        let batch2 = uint32_batch_non_null((MID_POINT..TOTAL_ROWS).rev());
2624
2625        let mut coalescer = BatchCoalescer::new(
2626            Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])),
2627            TOTAL_ROWS as usize,
2628        );
2629        coalescer.push_batch(batch1).unwrap();
2630
2631        let rev_indices = (0..((TOTAL_ROWS - MID_POINT) as u64)).rev();
2632        let reversed_indices_batch = uint64_batch_non_null(rev_indices);
2633
2634        let reverse_indices = UInt64Array::from(reversed_indices_batch.column(0).to_data());
2635        coalescer
2636            .push_batch_with_indices(batch2, &reverse_indices)
2637            .unwrap();
2638
2639        coalescer.finish_buffered_batch().unwrap();
2640        let actual = coalescer.next_completed_batch().unwrap();
2641
2642        let expected = uint32_batch_non_null(0..TOTAL_ROWS);
2643
2644        assert_eq!(expected, actual);
2645    }
2646
2647    #[test]
2648    fn test_push_batch_schema_mismatch_fewer_columns() {
2649        // Coalescer expects 0 columns, batch has 1
2650        let empty_schema = Arc::new(Schema::empty());
2651        let mut coalescer = BatchCoalescer::new(empty_schema, 100);
2652        let batch = uint32_batch(0..5);
2653        let result = coalescer.push_batch(batch);
2654        assert!(result.is_err());
2655        let err = result.unwrap_err().to_string();
2656        assert!(
2657            err.contains("Batch has 1 columns but BatchCoalescer expects 0"),
2658            "unexpected error: {err}"
2659        );
2660    }
2661
2662    #[test]
2663    fn test_push_batch_schema_mismatch_more_columns() {
2664        // Coalescer expects 2 columns, batch has 1
2665        let schema = Arc::new(Schema::new(vec![
2666            Field::new("c0", DataType::UInt32, false),
2667            Field::new("c1", DataType::UInt32, false),
2668        ]));
2669        let mut coalescer = BatchCoalescer::new(schema, 100);
2670        let batch = uint32_batch(0..5);
2671        let result = coalescer.push_batch(batch);
2672        assert!(result.is_err());
2673        let err = result.unwrap_err().to_string();
2674        assert!(
2675            err.contains("Batch has 1 columns but BatchCoalescer expects 2"),
2676            "unexpected error: {err}"
2677        );
2678    }
2679
2680    #[test]
2681    fn test_push_batch_schema_mismatch_two_vs_zero() {
2682        // Coalescer expects 0 columns, batch has 2
2683        let empty_schema = Arc::new(Schema::empty());
2684        let mut coalescer = BatchCoalescer::new(empty_schema, 100);
2685        let schema = Arc::new(Schema::new(vec![
2686            Field::new("c0", DataType::UInt32, false),
2687            Field::new("c1", DataType::UInt32, false),
2688        ]));
2689        let batch = RecordBatch::try_new(
2690            schema,
2691            vec![
2692                Arc::new(UInt32Array::from(vec![1, 2, 3])),
2693                Arc::new(UInt32Array::from(vec![4, 5, 6])),
2694            ],
2695        )
2696        .unwrap();
2697        let result = coalescer.push_batch(batch);
2698        assert!(result.is_err());
2699        let err = result.unwrap_err().to_string();
2700        assert!(
2701            err.contains("Batch has 2 columns but BatchCoalescer expects 0"),
2702            "unexpected error: {err}"
2703        );
2704    }
2705
2706    #[test]
2707    fn test_size_grows_with_buffering_and_shrinks_when_draining() {
2708        let batch = uint32_batch(0..8);
2709        let mut coalescer = BatchCoalescer::new(batch.schema(), 21);
2710        let baseline = coalescer.size();
2711        assert!(baseline > 0, "size includes container capacities");
2712
2713        // Buffer rows without completing a batch
2714        coalescer.push_batch(batch.clone()).unwrap();
2715        assert!(coalescer.next_completed_batch().is_none());
2716        let buffered = coalescer.size();
2717        assert!(
2718            buffered > baseline,
2719            "buffering rows should grow size ({buffered} > {baseline})"
2720        );
2721
2722        // Push enough to complete several batches
2723        for _ in 0..10 {
2724            coalescer.push_batch(batch.clone()).unwrap();
2725        }
2726        let peak = coalescer.size();
2727        assert!(peak > buffered);
2728
2729        // Draining completed batches must never grow the reported size
2730        let mut prev = peak;
2731        let mut drained_any = false;
2732        while coalescer.next_completed_batch().is_some() {
2733            drained_any = true;
2734            let now = coalescer.size();
2735            assert!(now <= prev, "size grew while draining: {now} > {prev}");
2736            prev = now;
2737        }
2738        assert!(drained_any);
2739        assert!(
2740            prev < peak,
2741            "draining completed batches should release memory"
2742        );
2743    }
2744
2745    #[test]
2746    fn test_size_string_view_buffers_released_after_drain() {
2747        // Long strings spill into external data buffers, exercising the byte-view
2748        // size accounting through the real coalescer path (including compaction).
2749        let batch = stringview_batch_repeated(
2750            1000,
2751            [Some("this string is definitely longer than 12 bytes")],
2752        );
2753        let mut coalescer = BatchCoalescer::new(batch.schema(), 4096);
2754        let baseline = coalescer.size();
2755
2756        for _ in 0..20 {
2757            coalescer.push_batch(batch.clone()).unwrap();
2758        }
2759        let peak = coalescer.size();
2760        assert!(
2761            peak > baseline,
2762            "buffered string view data should grow size ({peak} > {baseline})"
2763        );
2764
2765        coalescer.finish_buffered_batch().unwrap();
2766        while coalescer.next_completed_batch().is_some() {}
2767
2768        // Once fully drained the in-progress byte-view buffers are released.
2769        let drained = coalescer.size();
2770        assert!(
2771            drained < peak,
2772            "draining should release buffered data ({drained} < {peak})"
2773        );
2774    }
2775
2776    /// Every byte added to the accounting must eventually be removed: running the
2777    /// exact same push/finish/drain sequence twice must report identical sizes.
2778    /// This catches accounting leaks and drift without hard-coding magic numbers.
2779    #[test]
2780    fn test_size_accounting_conserved_across_cycles() {
2781        // Primitive column: internal capacities stabilize after the first cycle
2782        // (unlike byte-view, whose buffer sizer keeps growing), so the readings
2783        // are deterministic across cycles.
2784        let batch = uint32_batch(0..8);
2785        let mut coalescer = BatchCoalescer::new(batch.schema(), 4096);
2786
2787        let run_cycle = |coalescer: &mut BatchCoalescer| {
2788            for _ in 0..20 {
2789                coalescer.push_batch(batch.clone()).unwrap();
2790            }
2791            coalescer.finish_buffered_batch().unwrap();
2792            let peak = coalescer.size();
2793            while coalescer.next_completed_batch().is_some() {}
2794            (peak, coalescer.size())
2795        };
2796
2797        let (peak1, drained1) = run_cycle(&mut coalescer);
2798        let (peak2, drained2) = run_cycle(&mut coalescer);
2799
2800        assert_eq!(
2801            peak1, peak2,
2802            "identical work must report identical peak size"
2803        );
2804        assert_eq!(
2805            drained1, drained2,
2806            "fully-drained size must be stable across cycles (no accounting leak)"
2807        );
2808        assert!(
2809            drained1 < peak1,
2810            "draining must release the accounted memory"
2811        );
2812    }
2813}