arrow_select/coalesce/
primitive.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use crate::coalesce::{InProgressArray, check_filter_len};
6use crate::filter::{FilterIndices, FilterPredicate, FilterSelection, FilterSlices};
7use arrow_array::cast::AsArray;
8use arrow_array::{Array, ArrayRef, ArrowPrimitiveType, PrimitiveArray};
9use arrow_buffer::{NullBuffer, NullBufferBuilder, ScalarBuffer};
10use arrow_schema::{ArrowError, DataType};
11use std::fmt::Debug;
12use std::sync::Arc;
13
14/// InProgressArray for [`PrimitiveArray`]
15#[derive(Debug)]
16pub(crate) struct InProgressPrimitiveArray<T: ArrowPrimitiveType> {
17    /// Data type of the array
18    data_type: DataType,
19    /// The current source, if any
20    source: Option<ArrayRef>,
21    /// the target batch size (and thus size for views allocation)
22    batch_size: usize,
23    /// In progress nulls
24    nulls: NullBufferBuilder,
25    /// The currently in progress array
26    current: Vec<T::Native>,
27}
28
29impl<T: ArrowPrimitiveType> InProgressPrimitiveArray<T> {
30    /// Create a new `InProgressPrimitiveArray`
31    pub(crate) fn new(batch_size: usize, data_type: DataType) -> Self {
32        Self {
33            data_type,
34            batch_size,
35            source: None,
36            nulls: NullBufferBuilder::new(batch_size),
37            current: vec![],
38        }
39    }
40
41    /// Allocate space for output values if necessary.
42    ///
43    /// This is done on write (when we know it is necessary) rather than
44    /// eagerly to avoid allocations that are not used.
45    fn ensure_capacity(&mut self) {
46        if self.current.capacity() == 0 {
47            self.current.reserve(self.batch_size);
48        }
49    }
50
51    /// # Safety
52    ///
53    /// Every index yielded by `indices` must be less than `values.len()`, and
54    /// `indices` must yield exactly `selected_count` indices.
55    unsafe fn append_values_by_indices(
56        current: &mut Vec<T::Native>,
57        values: &[T::Native],
58        indices: FilterIndices<'_>,
59        selected_count: usize,
60    ) {
61        current.reserve(selected_count);
62        let current_len = current.len();
63        let mut written = 0;
64
65        // SAFETY: the caller guarantees that indices are in bounds and that
66        // exactly `selected_count` values are written. Capacity was reserved
67        // above.
68        unsafe {
69            let mut out = current
70                .spare_capacity_mut()
71                .as_mut_ptr()
72                .cast::<T::Native>();
73
74            indices.for_each(|idx| {
75                out.write(*values.get_unchecked(idx));
76                out = out.add(1);
77                written += 1;
78            });
79
80            current.set_len(current_len + written);
81        }
82
83        debug_assert_eq!(written, selected_count);
84    }
85
86    /// # Safety
87    ///
88    /// Every range yielded by `slices` must satisfy
89    /// `start <= end <= values.len()`, and the ranges must contain exactly
90    /// `selected_count` values in total.
91    unsafe fn append_values_by_slices(
92        current: &mut Vec<T::Native>,
93        values: &[T::Native],
94        slices: FilterSlices<'_>,
95        selected_count: usize,
96    ) {
97        current.reserve(selected_count);
98        let current_len = current.len();
99        let mut written = 0;
100
101        // SAFETY: the caller guarantees that slices are in bounds and that
102        // exactly `selected_count` values are written. Capacity was reserved
103        // above.
104        unsafe {
105            let mut out = current
106                .spare_capacity_mut()
107                .as_mut_ptr()
108                .cast::<T::Native>();
109
110            slices.for_each(|(start, end)| {
111                let len = end - start;
112                std::ptr::copy_nonoverlapping(values.as_ptr().add(start), out, len);
113                out = out.add(len);
114                written += len;
115            });
116
117            current.set_len(current_len + written);
118        }
119
120        debug_assert_eq!(written, selected_count);
121    }
122}
123
124#[inline]
125fn primitive_source<T: ArrowPrimitiveType>(
126    source: Option<&ArrayRef>,
127) -> Result<&PrimitiveArray<T>, ArrowError> {
128    Ok(source
129        .ok_or_else(|| {
130            ArrowError::InvalidArgumentError(
131                "Internal Error: InProgressPrimitiveArray: source not set".to_string(),
132            )
133        })?
134        .as_primitive::<T>())
135}
136
137fn append_filtered_nulls(
138    nulls: &mut NullBufferBuilder,
139    source_nulls: Option<&NullBuffer>,
140    filter: &FilterPredicate,
141) {
142    if let Some(filtered_nulls) = filter.filter_nulls(source_nulls) {
143        nulls.append_buffer(&filtered_nulls);
144    } else {
145        nulls.append_n_non_nulls(filter.count());
146    }
147}
148
149impl<T: ArrowPrimitiveType + Debug> InProgressArray for InProgressPrimitiveArray<T> {
150    fn set_source(&mut self, source: Option<ArrayRef>) {
151        self.source = source;
152    }
153
154    fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError> {
155        self.ensure_capacity();
156
157        let s = primitive_source::<T>(self.source.as_ref())?;
158
159        // add nulls if necessary
160        if let Some(nulls) = s.nulls().as_ref() {
161            let nulls = nulls.slice(offset, len);
162            self.nulls.append_buffer(&nulls);
163        } else {
164            self.nulls.append_n_non_nulls(len);
165        }
166
167        // Copy the values
168        let values = s.values();
169        self.current
170            .extend_from_slice(&values[offset..offset + len]);
171
172        Ok(())
173    }
174
175    fn copy_rows_by_filter(&mut self, filter: &FilterPredicate) -> Result<(), ArrowError> {
176        let source_len = primitive_source::<T>(self.source.as_ref())?.len();
177        check_filter_len(filter, source_len)?;
178
179        match filter.selection() {
180            FilterSelection::Indices(indices) => {
181                self.ensure_capacity();
182                let s = primitive_source::<T>(self.source.as_ref())?;
183
184                append_filtered_nulls(&mut self.nulls, s.nulls(), filter);
185                // SAFETY: `check_filter_len` verified the selected indices are
186                // in bounds, and the filter yields `filter.count()` indices.
187                unsafe {
188                    Self::append_values_by_indices(
189                        &mut self.current,
190                        s.values(),
191                        indices,
192                        filter.count(),
193                    )
194                };
195                Ok(())
196            }
197            FilterSelection::Slices(slices) => {
198                self.ensure_capacity();
199                let s = primitive_source::<T>(self.source.as_ref())?;
200
201                append_filtered_nulls(&mut self.nulls, s.nulls(), filter);
202                // SAFETY: `check_filter_len` verified the selected slices are
203                // in bounds, and the filter yields `filter.count()` values.
204                unsafe {
205                    Self::append_values_by_slices(
206                        &mut self.current,
207                        s.values(),
208                        slices,
209                        filter.count(),
210                    )
211                };
212                Ok(())
213            }
214            // Other selection shapes reuse the generic copy_rows path.
215            selection => self.copy_rows_by_selection(selection),
216        }
217    }
218
219    fn finish(&mut self) -> Result<ArrayRef, ArrowError> {
220        // take and reset the current values and nulls
221        let values = std::mem::take(&mut self.current);
222        let nulls = self.nulls.finish();
223        self.nulls = NullBufferBuilder::new(self.batch_size);
224
225        let array = PrimitiveArray::<T>::try_new(ScalarBuffer::from(values), nulls)?
226            // preserve timezone / precision+scale if applicable
227            .with_data_type(self.data_type.clone());
228        Ok(Arc::new(array))
229    }
230
231    fn size(&self) -> usize {
232        self.source
233            .as_ref()
234            .map_or(0, |source| source.get_array_memory_size())
235            + self.current.capacity() * std::mem::size_of::<T::Native>()
236            + self.nulls.allocated_size()
237    }
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243    use crate::filter::FilterBuilder;
244    use arrow_array::types::Int32Type;
245    use arrow_array::{BooleanArray, Int32Array};
246
247    #[test]
248    fn test_copy_rows_by_filter_index_iterator() {
249        let source =
250            Int32Array::from_iter((0..21).map(|idx| if idx % 5 == 0 { None } else { Some(idx) }));
251        let filter = BooleanArray::from_iter(
252            (0..21).map(|idx| Some(matches!(idx, 0 | 1 | 2 | 3 | 5 | 8 | 13))),
253        );
254        let predicate = FilterBuilder::new(&filter).build();
255        let FilterSelection::Indices(indices) = predicate.selection() else {
256            panic!("expected index iterator selection");
257        };
258        let mut selected_indices = Vec::new();
259        indices.for_each(|idx| selected_indices.push(idx));
260        assert_eq!(selected_indices, vec![0, 1, 2, 3, 5, 8, 13]);
261
262        let mut in_progress = InProgressPrimitiveArray::<Int32Type>::new(7, DataType::Int32);
263        in_progress.set_source(Some(Arc::new(source)));
264        in_progress.copy_rows_by_filter(&predicate).unwrap();
265
266        let result = in_progress.finish().unwrap();
267        let result = result.as_primitive::<Int32Type>();
268        let expected = Int32Array::from(vec![
269            None,
270            Some(1),
271            Some(2),
272            Some(3),
273            None,
274            Some(8),
275            Some(13),
276        ]);
277        assert_eq!(result, &expected);
278    }
279
280    #[test]
281    fn test_copy_rows_by_filter_slice_iterator() {
282        let source =
283            Int32Array::from_iter((0..16).map(|idx| if idx % 5 == 0 { None } else { Some(idx) }));
284        let filter = BooleanArray::from_iter((0..16).map(|idx| Some(!matches!(idx, 3 | 9))));
285        let predicate = FilterBuilder::new(&filter).build();
286        let FilterSelection::Slices(slices) = predicate.selection() else {
287            panic!("expected slice iterator selection");
288        };
289        let mut selected_slices = Vec::new();
290        slices.for_each(|slice| selected_slices.push(slice));
291        assert_eq!(selected_slices, vec![(0, 3), (4, 9), (10, 16)]);
292
293        let mut in_progress = InProgressPrimitiveArray::<Int32Type>::new(14, DataType::Int32);
294        in_progress.set_source(Some(Arc::new(source)));
295        in_progress.copy_rows_by_filter(&predicate).unwrap();
296
297        let result = in_progress.finish().unwrap();
298        let result = result.as_primitive::<Int32Type>();
299        let expected = Int32Array::from(vec![
300            None,
301            Some(1),
302            Some(2),
303            Some(4),
304            None,
305            Some(6),
306            Some(7),
307            Some(8),
308            None,
309            Some(11),
310            Some(12),
311            Some(13),
312            Some(14),
313            None,
314        ]);
315        assert_eq!(result, &expected);
316    }
317
318    #[test]
319    fn test_copy_rows_by_filter_rejects_all_filter_longer_than_source() {
320        let source: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3]));
321        let mut in_progress = InProgressPrimitiveArray::<Int32Type>::new(4, DataType::Int32);
322        in_progress.set_source(Some(source));
323
324        let filter = BooleanArray::from(vec![true; 4]);
325        let predicate = FilterBuilder::new(&filter).build();
326        assert!(matches!(predicate.selection(), FilterSelection::All { .. }));
327
328        let err = in_progress.copy_rows_by_filter(&predicate).unwrap_err();
329        assert!(matches!(
330            err,
331            ArrowError::InvalidArgumentError(message)
332                if message == "Filter predicate of length 4 is larger than source array of length 3"
333        ));
334    }
335
336    #[test]
337    fn test_size_empty() {
338        // A fresh in-progress array has allocated nothing yet
339        let in_progress = InProgressPrimitiveArray::<Int32Type>::new(64, DataType::Int32);
340        assert_eq!(in_progress.size(), 0);
341    }
342
343    #[test]
344    fn test_size_counts_source() {
345        let mut in_progress = InProgressPrimitiveArray::<Int32Type>::new(64, DataType::Int32);
346        let source: ArrayRef = Arc::new(Int32Array::from_iter_values(0..100));
347        in_progress.set_source(Some(Arc::clone(&source)));
348        // Nothing copied yet, so size is exactly the source's memory
349        assert_eq!(in_progress.size(), source.get_array_memory_size());
350    }
351
352    #[test]
353    fn test_size_counts_values_buffer_and_resets_on_finish() {
354        const BATCH_SIZE: usize = 64;
355        let mut in_progress =
356            InProgressPrimitiveArray::<Int32Type>::new(BATCH_SIZE, DataType::Int32);
357        // Non-null source: the nulls builder stays empty (allocated_size == 0),
358        // so the only growth is the values buffer.
359        let source: ArrayRef = Arc::new(Int32Array::from_iter_values(0..100));
360        let source_size = source.get_array_memory_size();
361        in_progress.set_source(Some(Arc::clone(&source)));
362
363        in_progress.copy_rows(0, 50).unwrap();
364        assert!(
365            in_progress.size() >= source_size + 50 * size_of::<i32>(),
366            "values buffer under-counted: {} < {} + {} * {}",
367            in_progress.size(),
368            source_size,
369            50,
370            size_of::<i32>(),
371        );
372
373        // finish() takes the buffered values/nulls but keeps the source, so the
374        // reported size drops back to exactly the source.
375        in_progress.finish().unwrap();
376        assert_eq!(in_progress.size(), source_size);
377    }
378
379    #[test]
380    fn test_size_counts_null_buffer() {
381        const BATCH_SIZE: usize = 64;
382
383        let in_progress_bytes = |source: ArrayRef| {
384            let mut in_progress =
385                InProgressPrimitiveArray::<Int32Type>::new(BATCH_SIZE, DataType::Int32);
386            let source_len = source.len();
387            in_progress.set_source(Some(source));
388            in_progress.copy_rows(0, source_len / 2).unwrap();
389            in_progress.size()
390        };
391
392        // All values valid: the nulls builder never allocates.
393        let all_valid = in_progress_bytes(Arc::new(Int32Array::from_iter_values(0..100)));
394        // Some values null: copying materializes a null buffer that must count.
395        let with_nulls = in_progress_bytes(Arc::new(Int32Array::from_iter(
396            (0..100).map(|i| (i % 2 == 0).then_some(i)),
397        )));
398
399        assert!(
400            with_nulls > all_valid,
401            "null buffer must be included in size(): with_nulls={with_nulls} all_valid={all_valid}"
402        );
403    }
404}
arrow_select/coalesce/primitive.rs

arrow_select/coalesce/
primitive.rs
