arrow_select/coalesce/
generic.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;
8use arrow_array::{Array, ArrayRef};
9use arrow_schema::ArrowError;
10
11/// Generic implementation for [`InProgressArray`] that works with any type of
12/// array.
13///
14/// Internally, this buffers arrays and then calls other kernels such as
15/// [`concat`] to produce the final array.
16///
17/// [`concat`]: crate::concat::concat
18#[derive(Debug)]
19pub(crate) struct GenericInProgressArray {
20    /// The current source
21    source: Option<ArrayRef>,
22
23    /// Is [`Self::source`] referenced in [`Self::buffered_arrays`]
24    source_data_referenced_in_buffers: bool,
25    /// The buffered array slices
26    buffered_arrays: Vec<ArrayRef>,
27
28    /// The number of bytes the arrays in [`Self::buffered_arrays`] takes
29    total_size_of_non_shared_buffers: usize,
30}
31
32impl GenericInProgressArray {
33    /// Create a new `GenericInProgressArray`
34    pub(crate) fn new() -> Self {
35        Self {
36            source: None,
37            buffered_arrays: vec![],
38            total_size_of_non_shared_buffers: 0,
39            source_data_referenced_in_buffers: false,
40        }
41    }
42}
43impl InProgressArray for GenericInProgressArray {
44    fn set_source(&mut self, source: Option<ArrayRef>) {
45        if let Some(old_source) = self.source.take() {
46            //  If the source is still referenced in buffered_arrays,
47            // then count it now
48            if self.source_data_referenced_in_buffers {
49                self.total_size_of_non_shared_buffers += old_source.get_array_memory_size();
50            }
51        }
52        self.source_data_referenced_in_buffers = false;
53        self.source = source;
54    }
55
56    fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError> {
57        let source = self.source.as_ref().ok_or_else(|| {
58            ArrowError::InvalidArgumentError(
59                "Internal Error: GenericInProgressArray: source not set".to_string(),
60            )
61        })?;
62        // No need to count the size of the array that was pushed to buffered_arrays
63        // since the data is shared with the `source` memory that is already counted
64        self.source_data_referenced_in_buffers = true;
65        let array = source.slice(offset, len);
66        self.buffered_arrays.push(array);
67        Ok(())
68    }
69
70    fn copy_rows_by_filter_from(
71        &mut self,
72        source: ArrayRef,
73        filter: &FilterPredicate,
74    ) -> Result<(), ArrowError> {
75        let array = filter.filter(source.as_ref())?;
76        self.total_size_of_non_shared_buffers += array.get_array_memory_size();
77        self.buffered_arrays.push(array);
78        Ok(())
79    }
80
81    fn finish(&mut self) -> Result<ArrayRef, ArrowError> {
82        // Concatenate all buffered arrays into a single array, which uses 2x
83        // peak memory
84        let array = concat(
85            &self
86                .buffered_arrays
87                .iter()
88                .map(|array| array.as_ref())
89                .collect::<Vec<_>>(),
90        )?;
91        self.buffered_arrays.clear();
92        self.total_size_of_non_shared_buffers = 0;
93        self.source_data_referenced_in_buffers = false;
94        Ok(array)
95    }
96
97    fn size(&self) -> usize {
98        self.total_size_of_non_shared_buffers
99            + self.buffered_arrays.capacity() * size_of::<ArrayRef>()
100            + self
101                .source
102                .as_ref()
103                .map_or(0, |a| a.get_array_memory_size())
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110    use arrow_array::Int32Array;
111    use std::sync::Arc;
112
113    fn arr(range: std::ops::Range<i32>) -> ArrayRef {
114        Arc::new(Int32Array::from_iter_values(range))
115    }
116
117    #[test]
118    fn test_size_empty() {
119        let in_progress = GenericInProgressArray::new();
120        assert_eq!(in_progress.size(), 0);
121    }
122
123    #[test]
124    fn test_source_is_counted_in_memory_and_released_when_not_used() {
125        let mut in_progress = GenericInProgressArray::new();
126
127        // Starting with no used memory
128        assert_eq!(in_progress.size(), 0);
129        {
130            let source1 = arr(0..100);
131            in_progress.set_source(Some(Arc::clone(&source1)));
132
133            // We are holding on source so this is the size
134            assert_eq!(in_progress.size(), source1.get_array_memory_size());
135        }
136
137        {
138            // Replacing an unused source drops the old size and counts only the new one
139            let source2 = arr(0..40);
140            in_progress.set_source(Some(Arc::clone(&source2)));
141            assert_eq!(in_progress.size(), source2.get_array_memory_size());
142        }
143
144        // Drop the used source
145        in_progress.set_source(None);
146
147        // The source is no longer being held, so it should reduce the size
148        assert_eq!(in_progress.size(), 0);
149    }
150
151    #[test]
152    fn test_double_copy_on_same_source_should_not_double_count() {
153        let mut in_progress = GenericInProgressArray::new();
154
155        // Starting with no used memory
156        assert_eq!(in_progress.size(), 0);
157
158        let source = arr(0..100);
159        in_progress.set_source(Some(Arc::clone(&source)));
160
161        // We are holding on source so this is the size
162        let size_before_copy = in_progress.size();
163        assert_eq!(size_before_copy, source.get_array_memory_size());
164
165        for _ in 0..2 {
166            // Only copy a subset
167            in_progress.copy_rows(0, 98).unwrap();
168
169            // The size should now account for the buffered but not the actual array data since we are still holding on it
170            assert!(
171                in_progress.size() > size_before_copy,
172                "size after copy {} should be greater than before copy {size_before_copy}",
173                in_progress.size()
174            );
175            {
176                let in_progress_size = in_progress.size() as f64;
177                let source_size = source.get_array_memory_size();
178                let size_if_source_and_sliced_would_be_counted = (source_size as f64) * 1.8;
179                assert!(
180                    in_progress_size < size_if_source_and_sliced_would_be_counted,
181                    "size after copy {in_progress_size} should not include the source and sliced array (should be greater than {size_if_source_and_sliced_would_be_counted}), source size is {source_size}"
182                );
183            }
184        }
185
186        let size_before_clear_source = in_progress.size();
187
188        // Drop the used source
189        in_progress.set_source(None);
190
191        // The source is still being held in the buffered
192        assert_eq!(in_progress.size(), size_before_clear_source);
193
194        {
195            let source2 = arr(0..40);
196            in_progress.set_source(Some(Arc::clone(&source2)));
197            assert_eq!(
198                in_progress.size(),
199                size_before_clear_source + source2.get_array_memory_size()
200            );
201            in_progress.set_source(None);
202        }
203
204        in_progress.finish().unwrap();
205
206        // There is still some memory being held by some leftover capacity but not arrays
207        assert!(in_progress.size() < source.get_array_memory_size());
208    }
209}
arrow_select/coalesce/generic.rs

arrow_select/coalesce/
generic.rs
