arrow_select/coalesce/
list.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 super::InProgressArray;
6use crate::concat::concat;
7use crate::filter::{FilterPredicate, FilterSelection};
8use crate::interleave::interleave;
9use arrow_array::{Array, ArrayRef, new_empty_array};
10use arrow_schema::{ArrowError, DataType};
11
12/// Bound retained input data while batching enough to amortize interleave.
13const MAX_PENDING_SOURCES: usize = 64;
14const MAX_PENDING_SOURCE_BYTES: usize = 64 * 1024 * 1024;
15
16#[derive(Debug)]
17pub(crate) struct InProgressListArray {
18    /// The current source for contiguous row copies.
19    source: Option<ArrayRef>,
20    /// The output array type, either List or LargeList.
21    data_type: DataType,
22    /// Materialized sparse groups and contiguous source slices in output order.
23    buffered_arrays: Vec<ArrayRef>,
24    /// Source arrays retained until the next sparse interleave.
25    sparse_sources: Vec<ArrayRef>,
26    /// Selected `(source index, row index)` pairs for `sparse_sources`.
27    sparse_indices: Vec<(usize, usize)>,
28    /// Total memory size reported by the retained sparse sources.
29    sparse_source_bytes: usize,
30    /// Whether buffered slices reference the current source's data.
31    source_data_referenced_in_buffers: bool,
32    /// Memory retained by buffered arrays that is not counted by `source`.
33    total_size_of_non_shared_buffers: usize,
34}
35
36impl InProgressListArray {
37    pub(crate) fn new(data_type: DataType) -> Self {
38        Self {
39            source: None,
40            data_type,
41            buffered_arrays: Vec::new(),
42            sparse_sources: Vec::new(),
43            sparse_indices: Vec::new(),
44            sparse_source_bytes: 0,
45            source_data_referenced_in_buffers: false,
46            total_size_of_non_shared_buffers: 0,
47        }
48    }
49
50    fn sparse_limit_reached(&self) -> bool {
51        self.sparse_sources.len() >= MAX_PENDING_SOURCES
52            || self.sparse_source_bytes >= MAX_PENDING_SOURCE_BYTES
53    }
54
55    fn flush_sparse(&mut self) -> Result<(), ArrowError> {
56        if self.sparse_indices.is_empty() {
57            self.sparse_sources.clear();
58            self.sparse_source_bytes = 0;
59            return Ok(());
60        }
61
62        let sources = self
63            .sparse_sources
64            .iter()
65            .map(|array| array.as_ref())
66            .collect::<Vec<_>>();
67        let array = interleave(&sources, &self.sparse_indices)?;
68        self.total_size_of_non_shared_buffers += array.get_array_memory_size();
69        self.buffered_arrays.push(array);
70        self.sparse_sources.clear();
71        self.sparse_indices.clear();
72        self.sparse_source_bytes = 0;
73        Ok(())
74    }
75}
76
77impl InProgressArray for InProgressListArray {
78    fn set_source(&mut self, source: Option<ArrayRef>) {
79        if let Some(old_source) = self.source.take()
80            && self.source_data_referenced_in_buffers
81        {
82            self.total_size_of_non_shared_buffers += old_source.get_array_memory_size();
83        }
84        self.source_data_referenced_in_buffers = false;
85        self.source = source;
86    }
87
88    fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError> {
89        self.flush_sparse()?;
90        let source = self.source.as_ref().ok_or_else(|| {
91            ArrowError::InvalidArgumentError("InProgressListArray source not set".into())
92        })?;
93        self.source_data_referenced_in_buffers = true;
94        self.buffered_arrays.push(source.slice(offset, len));
95        Ok(())
96    }
97
98    fn copy_rows_by_filter_from(
99        &mut self,
100        source: ArrayRef,
101        filter: &FilterPredicate,
102    ) -> Result<(), ArrowError> {
103        let source_size = source.get_array_memory_size();
104        if !self.sparse_sources.is_empty()
105            && (self.sparse_limit_reached()
106                || self.sparse_source_bytes.saturating_add(source_size) > MAX_PENDING_SOURCE_BYTES)
107        {
108            self.flush_sparse()?;
109        }
110
111        let source_index = self.sparse_sources.len();
112        match filter.selection() {
113            FilterSelection::None => return Ok(()),
114            FilterSelection::All { len } => self
115                .sparse_indices
116                .extend((0..len).map(|index| (source_index, index))),
117            FilterSelection::Slices(slices) => slices.for_each(|(start, end)| {
118                self.sparse_indices
119                    .extend((start..end).map(|index| (source_index, index)));
120            }),
121            FilterSelection::Indices(indices) => indices.for_each(|index| {
122                self.sparse_indices.push((source_index, index));
123            }),
124        }
125        self.sparse_sources.push(source);
126        self.sparse_source_bytes += source_size;
127        if self.sparse_limit_reached() {
128            self.flush_sparse()?;
129        }
130        Ok(())
131    }
132
133    fn finish(&mut self) -> Result<ArrayRef, ArrowError> {
134        self.flush_sparse()?;
135        let result = match self.buffered_arrays.len() {
136            0 => new_empty_array(&self.data_type),
137            1 => self.buffered_arrays.pop().unwrap(),
138            _ => concat(
139                &self
140                    .buffered_arrays
141                    .iter()
142                    .map(|array| array.as_ref())
143                    .collect::<Vec<_>>(),
144            )?,
145        };
146        self.buffered_arrays.clear();
147        self.total_size_of_non_shared_buffers = 0;
148        self.source_data_referenced_in_buffers = false;
149        Ok(result)
150    }
151
152    fn size(&self) -> usize {
153        self.total_size_of_non_shared_buffers
154            + self.buffered_arrays.capacity() * size_of::<ArrayRef>()
155            + self.sparse_sources.capacity() * size_of::<ArrayRef>()
156            + self.sparse_indices.capacity() * size_of::<(usize, usize)>()
157            + self
158                .source
159                .as_ref()
160                .map_or(0, |source| source.get_array_memory_size())
161            + self.sparse_source_bytes
162    }
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168    use crate::filter::FilterBuilder;
169    use arrow_array::{BooleanArray, Int32Array, ListArray, StructArray, UInt32Array};
170    use arrow_buffer::{NullBuffer, OffsetBuffer};
171    use arrow_schema::Field;
172    use std::sync::Arc;
173
174    fn list_array() -> ArrayRef {
175        let field = Arc::new(Field::new("item", DataType::Int32, true));
176        let offsets = OffsetBuffer::from_lengths([2, 1, 0, 2, 0]);
177        let values = Arc::new(Int32Array::from(vec![
178            Some(1),
179            None,
180            Some(2),
181            Some(3),
182            Some(4),
183        ]));
184        let nulls = NullBuffer::from(vec![true, true, false, true, true]);
185        Arc::new(ListArray::new(field, offsets, values, Some(nulls)))
186    }
187
188    #[test]
189    fn sparse_copy_preserves_nulls_and_empty_lists() {
190        let source = list_array();
191        let filter = BooleanArray::from(vec![true, false, true, true, true]);
192        let predicate = FilterBuilder::new(&filter).optimize().build();
193        let expected = crate::filter::filter(source.as_ref(), &filter).unwrap();
194        let mut in_progress = InProgressListArray::new(source.data_type().clone());
195        in_progress
196            .copy_rows_by_filter_from(source, &predicate)
197            .unwrap();
198        assert_eq!(in_progress.sparse_sources.len(), 1);
199        assert!(in_progress.buffered_arrays.is_empty());
200        assert_eq!(in_progress.finish().unwrap().as_ref(), expected.as_ref());
201    }
202
203    #[test]
204    fn sparse_copy_nested_children() {
205        let child = StructArray::from(vec![(
206            Arc::new(Field::new("id", DataType::Int32, false)),
207            Arc::new(Int32Array::from_iter_values(0..6)) as ArrayRef,
208        )]);
209        let field = Arc::new(Field::new("item", child.data_type().clone(), true));
210        let source = Arc::new(ListArray::new(
211            field,
212            OffsetBuffer::from_lengths([2, 1, 3]),
213            Arc::new(child),
214            None,
215        )) as ArrayRef;
216        let filter = BooleanArray::from(vec![true, false, true]);
217        let predicate = FilterBuilder::new(&filter).optimize().build();
218        let expected = crate::take::take(
219            source.as_ref(),
220            &UInt32Array::from_iter_values([0, 2]),
221            None,
222        )
223        .unwrap();
224        let mut in_progress = InProgressListArray::new(source.data_type().clone());
225        in_progress
226            .copy_rows_by_filter_from(source, &predicate)
227            .unwrap();
228        let output = in_progress.finish().unwrap();
229        assert_eq!(output.as_ref(), expected.as_ref());
230    }
231
232    #[test]
233    fn preserves_order_when_sparse_and_contiguous_inputs_mix() {
234        let first = list_array();
235        let second = list_array();
236        let filter = BooleanArray::from(vec![true, false, false, true, false]);
237        let predicate = FilterBuilder::new(&filter).optimize().build();
238        let filtered = crate::filter::filter(first.as_ref(), &filter).unwrap();
239        let second_slice = second.slice(1, 2);
240        let expected = concat(&[filtered.as_ref(), second_slice.as_ref()]).unwrap();
241        let mut in_progress = InProgressListArray::new(first.data_type().clone());
242        in_progress
243            .copy_rows_by_filter_from(first, &predicate)
244            .unwrap();
245        in_progress.set_source(Some(second));
246        in_progress.copy_rows(1, 2).unwrap();
247
248        assert_eq!(in_progress.finish().unwrap().as_ref(), expected.as_ref());
249    }
250
251    #[test]
252    fn bounds_retained_sparse_sources() {
253        let source = list_array();
254        let filter = BooleanArray::from(vec![true, false, false, false, false]);
255        let predicate = FilterBuilder::new(&filter).optimize().build();
256        let mut in_progress = InProgressListArray::new(source.data_type().clone());
257
258        for _ in 0..MAX_PENDING_SOURCES {
259            in_progress
260                .copy_rows_by_filter_from(Arc::clone(&source), &predicate)
261                .unwrap();
262        }
263
264        assert!(in_progress.sparse_sources.is_empty());
265        assert_eq!(in_progress.sparse_source_bytes, 0);
266        assert_eq!(in_progress.buffered_arrays.len(), 1);
267        assert_eq!(in_progress.finish().unwrap().len(), MAX_PENDING_SOURCES);
268    }
269}
arrow_select/coalesce/list.rs

arrow_select/coalesce/
list.rs
