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

arrow_select/coalesce/
generic.rs
