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

arrow_select/coalesce/
fixed_size_binary.rs
