arrow_select/coalesce/
boolean.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::filter::FilterPredicate;
20use arrow_array::cast::AsArray;
21use arrow_array::{Array, ArrayRef, BooleanArray};
22use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBufferBuilder};
23use arrow_schema::ArrowError;
24use std::sync::Arc;
25
26#[derive(Debug)]
27pub(crate) struct InProgressBooleanArray {
28    source: Option<ArrayRef>,
29    batch_size: usize,
30    values: Option<BooleanBufferBuilder>,
31    nulls: NullBufferBuilder,
32}
33
34impl InProgressBooleanArray {
35    pub(crate) fn new(batch_size: usize) -> Self {
36        Self {
37            source: None,
38            batch_size,
39            values: None,
40            nulls: NullBufferBuilder::new(batch_size),
41        }
42    }
43
44    fn values(&mut self) -> &mut BooleanBufferBuilder {
45        self.values
46            .get_or_insert_with(|| BooleanBufferBuilder::new(self.batch_size))
47    }
48
49    fn append_array(&mut self, array: &BooleanArray) {
50        self.values().append_buffer(array.values());
51        if let Some(nulls) = array.nulls() {
52            self.nulls.append_buffer(nulls);
53        } else {
54            self.nulls.append_n_non_nulls(array.len());
55        }
56    }
57}
58
59impl InProgressArray for InProgressBooleanArray {
60    fn set_source(&mut self, source: Option<ArrayRef>) {
61        self.source = source;
62    }
63
64    fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError> {
65        let source = self.source.as_ref().ok_or_else(|| {
66            ArrowError::InvalidArgumentError(
67                "Internal Error: InProgressBooleanArray: source not set".to_string(),
68            )
69        })?;
70        let array = source.as_boolean().slice(offset, len);
71        self.append_array(&array);
72        Ok(())
73    }
74
75    fn copy_rows_by_filter_from(
76        &mut self,
77        source: ArrayRef,
78        filter: &FilterPredicate,
79    ) -> Result<(), ArrowError> {
80        let filtered = filter.filter(source.as_ref())?;
81        self.append_array(filtered.as_boolean());
82        Ok(())
83    }
84
85    fn finish(&mut self) -> Result<ArrayRef, ArrowError> {
86        let values = self
87            .values
88            .take()
89            .map(|mut values| values.finish())
90            .unwrap_or_else(|| BooleanBuffer::new_unset(0));
91        let nulls = self.nulls.finish();
92        self.nulls = NullBufferBuilder::new(self.batch_size);
93        Ok(Arc::new(BooleanArray::new(values, nulls)))
94    }
95
96    fn size(&self) -> usize {
97        self.source
98            .as_ref()
99            .map_or(0, |source| source.get_array_memory_size())
100            + self
101                .values
102                .as_ref()
103                .map_or(0, |values| values.capacity() / 8)
104            + self.nulls.allocated_size()
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111    use crate::filter::FilterBuilder;
112
113    #[test]
114    fn copies_slices_and_sparse_filters() {
115        let source = Arc::new(BooleanArray::from(vec![
116            Some(true),
117            None,
118            Some(false),
119            Some(true),
120            Some(false),
121        ])) as ArrayRef;
122        let mut in_progress = InProgressBooleanArray::new(8);
123        in_progress.set_source(Some(Arc::clone(&source)));
124        in_progress.copy_rows(1, 2).unwrap();
125        in_progress.set_source(None);
126
127        let filter = BooleanArray::from(vec![true, false, false, true, false]);
128        let predicate = FilterBuilder::new(&filter).build();
129        in_progress
130            .copy_rows_by_filter_from(source, &predicate)
131            .unwrap();
132
133        let output = in_progress.finish().unwrap();
134        let expected = BooleanArray::from(vec![None, Some(false), Some(true), Some(true)]);
135        assert_eq!(output.as_boolean(), &expected);
136    }
137
138    #[test]
139    fn finish_resets_buffers() {
140        let mut in_progress = InProgressBooleanArray::new(8);
141        assert_eq!(in_progress.finish().unwrap().len(), 0);
142
143        let source = Arc::new(BooleanArray::from(vec![
144            Some(true),
145            None,
146            Some(false),
147            None,
148        ]));
149        in_progress.set_source(Some(Arc::clone(&source) as ArrayRef));
150        in_progress.copy_rows(0, 2).unwrap();
151        let first = in_progress.finish().unwrap();
152        assert_eq!(first.as_boolean(), &source.slice(0, 2));
153
154        in_progress.copy_rows(2, 2).unwrap();
155        let second = in_progress.finish().unwrap();
156        assert_eq!(second.as_boolean(), &source.slice(2, 2));
157        assert_eq!(in_progress.finish().unwrap().len(), 0);
158    }
159}
arrow_select/coalesce/boolean.rs

arrow_select/coalesce/
boolean.rs
