arrow_select/coalesce/
fixed_size_binary.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::builder::FixedSizeBinaryBuilder;
21use arrow_array::cast::AsArray;
22use arrow_array::{Array, ArrayRef};
23use arrow_schema::ArrowError;
24use std::sync::Arc;
25
26/// Specialized [`InProgressArray`] for `FixedSizeBinary` columns.
27#[derive(Debug)]
28pub(crate) struct InProgressFixedSizeBinaryArray {
29    source: Option<ArrayRef>,
30    value_length: i32,
31    batch_size: usize,
32    builder: Option<FixedSizeBinaryBuilder>,
33}
34
35impl InProgressFixedSizeBinaryArray {
36    pub(crate) fn new(value_length: i32, batch_size: usize) -> Self {
37        Self {
38            source: None,
39            value_length,
40            batch_size,
41            builder: None,
42        }
43    }
44
45    fn ensure_builder(&mut self) -> &mut FixedSizeBinaryBuilder {
46        self.builder.get_or_insert_with(|| {
47            FixedSizeBinaryBuilder::with_capacity(self.batch_size, self.value_length)
48        })
49    }
50}
51
52impl InProgressArray for InProgressFixedSizeBinaryArray {
53    fn set_source(&mut self, source: Option<ArrayRef>) {
54        self.source = source;
55    }
56
57    fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError> {
58        let source = self.source.as_ref().ok_or_else(|| {
59            ArrowError::InvalidArgumentError(
60                "Internal Error: InProgressFixedSizeBinaryArray: source not set".to_string(),
61            )
62        })?;
63        let sliced = source.as_fixed_size_binary().slice(offset, len);
64        self.ensure_builder().append_array(&sliced)?;
65        Ok(())
66    }
67
68    fn copy_rows_by_filter_from(
69        &mut self,
70        source: ArrayRef,
71        filter: &FilterPredicate,
72    ) -> Result<(), ArrowError> {
73        let filtered = filter.filter(source.as_ref())?;
74        if !filtered.is_empty() {
75            self.ensure_builder()
76                .append_array(filtered.as_fixed_size_binary())?;
77        }
78        Ok(())
79    }
80
81    fn finish(&mut self) -> Result<ArrayRef, ArrowError> {
82        let mut b = self
83            .builder
84            .take()
85            .unwrap_or_else(|| FixedSizeBinaryBuilder::new(self.value_length));
86        Ok(Arc::new(b.finish()))
87    }
88
89    fn size(&self) -> usize {
90        self.builder.as_ref().map_or(0, |b| b.capacity())
91            + self
92                .source
93                .as_ref()
94                .map_or(0, |a| a.get_array_memory_size())
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101    use crate::filter::FilterBuilder;
102    use arrow_array::{BooleanArray, FixedSizeBinaryArray};
103
104    fn make_fsb(value_length: i32, data: &[Option<&[u8]>]) -> FixedSizeBinaryArray {
105        let mut b = FixedSizeBinaryBuilder::with_capacity(data.len(), value_length);
106        for v in data {
107            match v {
108                Some(bytes) => b.append_value(bytes).unwrap(),
109                None => b.append_null(),
110            }
111        }
112        b.finish()
113    }
114
115    #[test]
116    fn test_roundtrip_with_nulls() {
117        let source = Arc::new(make_fsb(4, &[Some(b"abcd"), None, Some(b"ijkl")])) as ArrayRef;
118        let mut coalescer = InProgressFixedSizeBinaryArray::new(4, 8);
119        coalescer.set_source(Some(Arc::clone(&source)));
120        coalescer.copy_rows(0, 3).unwrap();
121        let output = coalescer.finish().unwrap();
122        let output = output.as_fixed_size_binary();
123        assert_eq!(output.len(), 3);
124        assert_eq!(output.value(0), b"abcd");
125        assert!(output.is_null(1));
126        assert_eq!(output.value(2), b"ijkl");
127    }
128
129    #[test]
130    fn test_offset_copy() {
131        let source =
132            Arc::new(make_fsb(4, &[Some(b"aaaa"), Some(b"bbbb"), Some(b"cccc")])) as ArrayRef;
133        let mut coalescer = InProgressFixedSizeBinaryArray::new(4, 8);
134        coalescer.set_source(Some(Arc::clone(&source)));
135        coalescer.copy_rows(1, 2).unwrap();
136        let output = coalescer.finish().unwrap();
137        let output = output.as_fixed_size_binary();
138        assert_eq!(output.len(), 2);
139        assert_eq!(output.value(0), b"bbbb");
140        assert_eq!(output.value(1), b"cccc");
141    }
142
143    #[test]
144    fn test_finish_preserves_value_length() {
145        // finish() with no rows written should still produce an array with the correct value_length
146        let mut coalescer = InProgressFixedSizeBinaryArray::new(16, 8);
147        let output = coalescer.finish().unwrap();
148        let output = output.as_fixed_size_binary();
149        assert_eq!(output.len(), 0);
150        assert_eq!(output.value_length(), 16);
151    }
152
153    #[test]
154    fn test_filter_path() {
155        let source = Arc::new(make_fsb(
156            4,
157            &[Some(b"aaaa"), Some(b"bbbb"), Some(b"cccc"), Some(b"dddd")],
158        )) as ArrayRef;
159        let filter_mask = BooleanArray::from(vec![true, false, true, false]);
160        let predicate = FilterBuilder::new(&filter_mask).build();
161        let mut coalescer = InProgressFixedSizeBinaryArray::new(4, 8);
162        coalescer
163            .copy_rows_by_filter_from(source, &predicate)
164            .unwrap();
165        let output = coalescer.finish().unwrap();
166        let output = output.as_fixed_size_binary();
167        assert_eq!(output.len(), 2);
168        assert_eq!(output.value(0), b"aaaa");
169        assert_eq!(output.value(1), b"cccc");
170    }
171}
arrow_select/coalesce/fixed_size_binary.rs

arrow_select/coalesce/
fixed_size_binary.rs
