arrow_select/coalesce/
byte.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::filter::{FilterIndices, FilterPredicate, FilterSelection};
7use arrow_array::builder::GenericByteBuilder;
8use arrow_array::cast::AsArray;
9use arrow_array::types::ByteArrayType;
10use arrow_array::{Array, ArrayRef, GenericByteArray, OffsetSizeTrait};
11use arrow_buffer::ArrowNativeType;
12use arrow_schema::ArrowError;
13use std::sync::Arc;
14
15pub(crate) struct InProgressByteArray<T: ByteArrayType> {
16    source: Option<ArrayRef>,
17    batch_size: usize,
18    direct: GenericByteBuilder<T>,
19}
20
21impl<T: ByteArrayType> std::fmt::Debug for InProgressByteArray<T> {
22    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
23        f.debug_struct("InProgressByteArray")
24            .field(
25                "direct_rows",
26                &self.direct.offsets_slice().len().saturating_sub(1),
27            )
28            .finish()
29    }
30}
31
32impl<T: ByteArrayType> InProgressByteArray<T>
33where
34    T::Offset: OffsetSizeTrait,
35{
36    pub(crate) fn new(batch_size: usize) -> Self {
37        Self {
38            source: None,
39            batch_size,
40            direct: GenericByteBuilder::with_capacity(0, 0),
41        }
42    }
43
44    fn checked_values_len(current: usize, additional: usize) -> Result<usize, ArrowError> {
45        let new_len = current
46            .checked_add(additional)
47            .ok_or(ArrowError::OffsetOverflowError(usize::MAX))?;
48        if new_len > T::Offset::MAX_OFFSET {
49            return Err(ArrowError::OffsetOverflowError(new_len));
50        }
51        Ok(new_len)
52    }
53
54    fn ensure_capacity(&mut self, source: &GenericByteArray<T>, rows: usize) {
55        if self.direct.offsets_slice().len() == 1 && !source.is_empty() {
56            let offsets = source.value_offsets();
57            let bytes = offsets[source.len()].as_usize() - offsets[0].as_usize();
58            let rows = rows.min(self.batch_size);
59            let estimated = bytes.saturating_mul(rows) / source.len();
60            // Later slices may contain slightly larger values than this first sample.
61            let capacity = estimated.saturating_add(estimated / 8);
62            self.direct = GenericByteBuilder::with_capacity(rows, capacity);
63        }
64    }
65
66    fn append_filter_indices(
67        &mut self,
68        source: &GenericByteArray<T>,
69        indices: FilterIndices<'_>,
70        count: usize,
71    ) -> Result<(), ArrowError> {
72        self.ensure_capacity(source, count);
73        indices.try_for_each(|idx| {
74            if source.is_null(idx) {
75                self.direct.append_null();
76            } else {
77                let value = source.value(idx);
78                Self::checked_values_len(
79                    self.direct.values_slice().len(),
80                    source.value_length(idx).as_usize(),
81                )?;
82                self.direct.append_value(value);
83            }
84            Ok::<(), ArrowError>(())
85        })?;
86        Ok(())
87    }
88}
89
90impl<T: ByteArrayType + 'static> InProgressArray for InProgressByteArray<T>
91where
92    T::Offset: OffsetSizeTrait,
93{
94    fn set_source(&mut self, source: Option<ArrayRef>) {
95        self.source = source;
96    }
97
98    fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError> {
99        let source = self.source.as_ref().ok_or_else(|| {
100            ArrowError::InvalidArgumentError(
101                "Internal Error: InProgressByteArray: source not set".to_string(),
102            )
103        })?;
104        let source = source.as_bytes::<T>().slice(offset, len);
105        self.ensure_capacity(&source, len);
106        self.direct.append_array(&source)
107    }
108
109    fn copy_rows_by_filter_from(
110        &mut self,
111        source: ArrayRef,
112        filter: &FilterPredicate,
113    ) -> Result<(), ArrowError> {
114        match filter.selection() {
115            FilterSelection::None => Ok(()),
116            // Index selections are sparse enough that copying values directly avoids the
117            // extra allocation of a filtered byte array. Dense contiguous selections use
118            // the filter kernel below, which copies their ranges in bulk.
119            FilterSelection::Indices(indices) => {
120                self.append_filter_indices(source.as_bytes::<T>(), indices, filter.count())
121            }
122            FilterSelection::Slices(_) | FilterSelection::All { .. } => {
123                let filtered = filter.filter(source.as_ref())?;
124                self.direct.append_array(filtered.as_bytes::<T>())
125            }
126        }
127    }
128
129    fn finish(&mut self) -> Result<ArrayRef, ArrowError> {
130        Ok(Arc::new(self.direct.finish()))
131    }
132
133    fn size(&self) -> usize {
134        self.source
135            .as_ref()
136            .map_or(0, |source| source.get_array_memory_size())
137            + self.direct.offsets_capacity() * size_of::<T::Offset>()
138            + self.direct.values_capacity()
139            + self.direct.validity_capacity()
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146    use crate::filter::FilterBuilder;
147    use arrow_array::types::GenericStringType;
148    use arrow_array::{BooleanArray, StringArray};
149
150    #[test]
151    fn copy_rows_and_filter_preserve_order() {
152        let mut values = vec![Some("padding"); 4098];
153        values[1..4].copy_from_slice(&[Some("zero"), None, Some("two")]);
154        values[2048] = Some("far");
155        let source: ArrayRef = Arc::new(StringArray::from(values).slice(1, 4096));
156        let mask =
157            BooleanArray::from_iter((0..source.len()).map(|idx| Some(matches!(idx, 0..=2 | 2047))));
158        let predicate = FilterBuilder::new(&mask).optimize().build();
159        let mut in_progress = InProgressByteArray::<GenericStringType<i32>>::new(source.len());
160        in_progress.set_source(Some(Arc::clone(&source)));
161        in_progress.copy_rows(0, 1).unwrap();
162        in_progress
163            .copy_rows_by_filter_from(Arc::clone(&source), &predicate)
164            .unwrap();
165        in_progress.copy_rows(2, 1).unwrap();
166
167        let actual = in_progress.finish().unwrap();
168        let expected = StringArray::from(vec![
169            Some("zero"),
170            Some("zero"),
171            None,
172            Some("two"),
173            Some("far"),
174            Some("two"),
175        ]);
176        assert_eq!(actual.as_string::<i32>(), &expected);
177    }
178
179    #[test]
180    fn checked_values_len_boundaries() {
181        type Bytes = InProgressByteArray<GenericStringType<i32>>;
182        let max = i32::MAX as usize;
183
184        assert_eq!(Bytes::checked_values_len(max - 1, 1).unwrap(), max);
185        assert!(Bytes::checked_values_len(max, 1).is_err());
186        assert!(Bytes::checked_values_len(usize::MAX, 1).is_err());
187    }
188}
arrow_select/coalesce/byte.rs

arrow_select/coalesce/
byte.rs
