arrow_array/builder/
generic_byte_run_builder.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 crate::types::bytes::ByteArrayNativeType;
19use std::{any::Any, sync::Arc};
20
21use crate::{
22    ArrayRef, ArrowPrimitiveType, RunArray,
23    types::{BinaryType, ByteArrayType, LargeBinaryType, LargeUtf8Type, RunEndIndexType, Utf8Type},
24};
25
26use super::{ArrayBuilder, GenericByteBuilder, PrimitiveBuilder};
27
28use arrow_buffer::ArrowNativeType;
29
30/// Builder for [`RunArray`] of [`GenericByteArray`](crate::array::GenericByteArray)
31///
32/// # Example:
33///
34/// ```
35///
36/// # use arrow_array::builder::GenericByteRunBuilder;
37/// # use arrow_array::{GenericByteArray, BinaryArray};
38/// # use arrow_array::types::{BinaryType, Int16Type};
39/// # use arrow_array::{Array, Int16Array};
40/// # use arrow_array::cast::AsArray;
41///
42/// let mut builder =
43/// GenericByteRunBuilder::<Int16Type, BinaryType>::new();
44/// builder.extend([Some(b"abc"), Some(b"abc"), None, Some(b"def")].into_iter());
45/// builder.append_value(b"def");
46/// builder.append_null();
47/// let array = builder.finish();
48///
49/// assert_eq!(array.run_ends().values(), &[2, 3, 5, 6]);
50///
51/// let av = array.values();
52///
53/// assert!(!av.is_null(0));
54/// assert!(av.is_null(1));
55/// assert!(!av.is_null(2));
56/// assert!(av.is_null(3));
57///
58/// // Values are polymorphic and so require a downcast.
59/// let ava: &BinaryArray = av.as_binary();
60///
61/// assert_eq!(ava.value(0), b"abc");
62/// assert_eq!(ava.value(2), b"def");
63/// ```
64#[derive(Debug)]
65pub struct GenericByteRunBuilder<R, V>
66where
67    R: ArrowPrimitiveType,
68    V: ByteArrayType,
69{
70    run_ends_builder: PrimitiveBuilder<R>,
71    values_builder: GenericByteBuilder<V>,
72    current_value: Vec<u8>,
73    has_current_value: bool,
74    current_run_end_index: usize,
75    prev_run_end_index: usize,
76}
77
78impl<R, V> Default for GenericByteRunBuilder<R, V>
79where
80    R: ArrowPrimitiveType,
81    V: ByteArrayType,
82{
83    fn default() -> Self {
84        Self::new()
85    }
86}
87
88impl<R, V> GenericByteRunBuilder<R, V>
89where
90    R: ArrowPrimitiveType,
91    V: ByteArrayType,
92{
93    /// Creates a new `GenericByteRunBuilder`
94    pub fn new() -> Self {
95        Self {
96            run_ends_builder: PrimitiveBuilder::new(),
97            values_builder: GenericByteBuilder::<V>::new(),
98            current_value: Vec::new(),
99            has_current_value: false,
100            current_run_end_index: 0,
101            prev_run_end_index: 0,
102        }
103    }
104
105    /// Creates a new `GenericByteRunBuilder` with the provided capacity
106    ///
107    /// `capacity`: the expected number of run-end encoded values.
108    /// `data_capacity`: the expected number of bytes of run end encoded values
109    pub fn with_capacity(capacity: usize, data_capacity: usize) -> Self {
110        Self {
111            run_ends_builder: PrimitiveBuilder::with_capacity(capacity),
112            values_builder: GenericByteBuilder::<V>::with_capacity(capacity, data_capacity),
113            current_value: Vec::new(),
114            has_current_value: false,
115            current_run_end_index: 0,
116            prev_run_end_index: 0,
117        }
118    }
119}
120
121impl<R, V> ArrayBuilder for GenericByteRunBuilder<R, V>
122where
123    R: RunEndIndexType,
124    V: ByteArrayType,
125{
126    /// Returns the builder as a non-mutable `Any` reference.
127    fn as_any(&self) -> &dyn Any {
128        self
129    }
130
131    /// Returns the builder as a mutable `Any` reference.
132    fn as_any_mut(&mut self) -> &mut dyn Any {
133        self
134    }
135
136    /// Returns the boxed builder as a box of `Any`.
137    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
138        self
139    }
140
141    /// Returns the length of logical array encoded by
142    /// the eventual runs array.
143    fn len(&self) -> usize {
144        self.current_run_end_index
145    }
146
147    /// Builds the array and reset this builder.
148    fn finish(&mut self) -> ArrayRef {
149        Arc::new(self.finish())
150    }
151
152    /// Builds the array without resetting the builder.
153    fn finish_cloned(&self) -> ArrayRef {
154        Arc::new(self.finish_cloned())
155    }
156}
157
158impl<R, V> GenericByteRunBuilder<R, V>
159where
160    R: RunEndIndexType,
161    V: ByteArrayType,
162{
163    /// Appends optional value to the logical array encoded by the RunArray.
164    pub fn append_option(&mut self, input_value: Option<impl AsRef<V::Native>>) {
165        match input_value {
166            Some(value) => self.append_value(value),
167            None => self.append_null(),
168        }
169    }
170
171    /// Appends value to the logical array encoded by the RunArray.
172    pub fn append_value(&mut self, input_value: impl AsRef<V::Native>) {
173        let value: &[u8] = input_value.as_ref().as_ref();
174        if !self.has_current_value {
175            self.append_run_end();
176            self.current_value.extend_from_slice(value);
177            self.has_current_value = true;
178        } else if self.current_value.as_slice() != value {
179            self.append_run_end();
180            self.current_value.clear();
181            self.current_value.extend_from_slice(value);
182        }
183        self.current_run_end_index += 1;
184    }
185
186    /// Appends null to the logical array encoded by the RunArray.
187    pub fn append_null(&mut self) {
188        if self.has_current_value {
189            self.append_run_end();
190            self.current_value.clear();
191            self.has_current_value = false;
192        }
193        self.current_run_end_index += 1;
194    }
195
196    /// Creates the RunArray and resets the builder.
197    ///
198    /// # Panics
199    ///
200    /// Panics if RunArray cannot be built.
201    pub fn finish(&mut self) -> RunArray<R> {
202        // write the last run end to the array.
203        self.append_run_end();
204
205        // reset the run end index to zero.
206        self.current_value.clear();
207        self.has_current_value = false;
208        self.current_run_end_index = 0;
209        self.prev_run_end_index = 0;
210
211        // build the run encoded array by adding run_ends and values array as its children.
212        let run_ends_array = self.run_ends_builder.finish();
213        let values_array = self.values_builder.finish();
214        RunArray::<R>::try_new(&run_ends_array, &values_array).unwrap()
215    }
216
217    /// Creates the RunArray and without resetting the builder.
218    ///
219    /// # Panics
220    ///
221    /// Panics if RunArray cannot be built.
222    pub fn finish_cloned(&self) -> RunArray<R> {
223        let mut run_ends_array = self.run_ends_builder.finish_cloned();
224        let mut values_array = self.values_builder.finish_cloned();
225
226        // Add current run if one exists
227        if self.prev_run_end_index != self.current_run_end_index {
228            let mut run_end_builder = run_ends_array.into_builder().unwrap();
229            let mut values_builder = values_array.into_builder().unwrap();
230            self.append_run_end_with_builders(&mut run_end_builder, &mut values_builder);
231            run_ends_array = run_end_builder.finish();
232            values_array = values_builder.finish();
233        }
234
235        RunArray::<R>::try_new(&run_ends_array, &values_array).unwrap()
236    }
237
238    // Appends the current run to the array.
239    fn append_run_end(&mut self) {
240        // empty array or the function called without appending any value.
241        if self.prev_run_end_index == self.current_run_end_index {
242            return;
243        }
244        let run_end_index = self.run_end_index_as_native();
245        self.run_ends_builder.append_value(run_end_index);
246        if self.has_current_value {
247            let slice = self.current_value.as_slice();
248            let native = unsafe {
249                // Safety:
250                // As self.current_value is created from V::Native. The value V::Native can be
251                // built back from the bytes without validations
252                V::Native::from_bytes_unchecked(slice)
253            };
254            self.values_builder.append_value(native);
255        } else {
256            self.values_builder.append_null();
257        }
258        self.prev_run_end_index = self.current_run_end_index;
259    }
260
261    // Similar to `append_run_end` but on custom builders.
262    // Used in `finish_cloned` which is not suppose to mutate `self`.
263    fn append_run_end_with_builders(
264        &self,
265        run_ends_builder: &mut PrimitiveBuilder<R>,
266        values_builder: &mut GenericByteBuilder<V>,
267    ) {
268        let run_end_index = self.run_end_index_as_native();
269        run_ends_builder.append_value(run_end_index);
270        if self.has_current_value {
271            let slice = self.current_value.as_slice();
272            let native = unsafe {
273                // Safety:
274                // As self.current_value is created from V::Native. The value V::Native can be
275                // built back from the bytes without validations
276                V::Native::from_bytes_unchecked(slice)
277            };
278            values_builder.append_value(native);
279        } else {
280            values_builder.append_null();
281        }
282    }
283
284    fn run_end_index_as_native(&self) -> R::Native {
285        R::Native::from_usize(self.current_run_end_index).unwrap_or_else(|| {
286            panic!(
287                "Cannot convert the value {} from `usize` to native form of arrow datatype {}",
288                self.current_run_end_index,
289                R::DATA_TYPE
290            )
291        })
292    }
293}
294
295impl<R, V, S> Extend<Option<S>> for GenericByteRunBuilder<R, V>
296where
297    R: RunEndIndexType,
298    V: ByteArrayType,
299    S: AsRef<V::Native>,
300{
301    fn extend<T: IntoIterator<Item = Option<S>>>(&mut self, iter: T) {
302        for elem in iter {
303            self.append_option(elem);
304        }
305    }
306}
307
308/// Builder for [`RunArray`] of [`StringArray`](crate::array::StringArray)
309///
310/// ```
311/// // Create a run-end encoded array with run-end indexes data type as `i16`.
312/// // The encoded values are Strings.
313///
314/// # use arrow_array::builder::StringRunBuilder;
315/// # use arrow_array::{Int16Array, StringArray};
316/// # use arrow_array::types::Int16Type;
317/// # use arrow_array::cast::AsArray;
318/// #
319/// let mut builder = StringRunBuilder::<Int16Type>::new();
320///
321/// // The builder builds the dictionary value by value
322/// builder.append_value("abc");
323/// builder.append_null();
324/// builder.extend([Some("def"), Some("def"), Some("abc")]);
325/// let array = builder.finish();
326///
327/// assert_eq!(array.run_ends().values(), &[1, 2, 4, 5]);
328///
329/// // Values are polymorphic and so require a downcast.
330/// let av = array.values();
331/// let ava: &StringArray = av.as_string::<i32>();
332///
333/// assert_eq!(ava.value(0), "abc");
334/// assert!(av.is_null(1));
335/// assert_eq!(ava.value(2), "def");
336/// assert_eq!(ava.value(3), "abc");
337///
338/// ```
339pub type StringRunBuilder<K> = GenericByteRunBuilder<K, Utf8Type>;
340
341/// Builder for [`RunArray`] of [`LargeStringArray`](crate::array::LargeStringArray)
342pub type LargeStringRunBuilder<K> = GenericByteRunBuilder<K, LargeUtf8Type>;
343
344/// Builder for [`RunArray`] of [`BinaryArray`](crate::array::BinaryArray)
345///
346/// ```
347/// // Create a run-end encoded array with run-end indexes data type as `i16`.
348/// // The encoded data is binary values.
349///
350/// # use arrow_array::builder::BinaryRunBuilder;
351/// # use arrow_array::{BinaryArray, Int16Array};
352/// # use arrow_array::cast::AsArray;
353/// # use arrow_array::types::Int16Type;
354///
355/// let mut builder = BinaryRunBuilder::<Int16Type>::new();
356///
357/// // The builder builds the dictionary value by value
358/// builder.append_value(b"abc");
359/// builder.append_null();
360/// builder.extend([Some(b"def"), Some(b"def"), Some(b"abc")]);
361/// let array = builder.finish();
362///
363/// assert_eq!(array.run_ends().values(), &[1, 2, 4, 5]);
364///
365/// // Values are polymorphic and so require a downcast.
366/// let av = array.values();
367/// let ava: &BinaryArray = av.as_binary();
368///
369/// assert_eq!(ava.value(0), b"abc");
370/// assert!(av.is_null(1));
371/// assert_eq!(ava.value(2), b"def");
372/// assert_eq!(ava.value(3), b"abc");
373///
374/// ```
375pub type BinaryRunBuilder<K> = GenericByteRunBuilder<K, BinaryType>;
376
377/// Builder for [`RunArray`] of [`LargeBinaryArray`](crate::array::LargeBinaryArray)
378pub type LargeBinaryRunBuilder<K> = GenericByteRunBuilder<K, LargeBinaryType>;
379
380#[cfg(test)]
381mod tests {
382    use super::*;
383
384    use crate::GenericByteArray;
385    use crate::Int16RunArray;
386    use crate::array::Array;
387    use crate::cast::AsArray;
388    use crate::types::{Int16Type, Int32Type};
389
390    fn test_bytes_run_builder<T>(values: Vec<&T::Native>)
391    where
392        T: ByteArrayType,
393        <T as ByteArrayType>::Native: PartialEq,
394        <T as ByteArrayType>::Native: AsRef<<T as ByteArrayType>::Native>,
395    {
396        let mut builder = GenericByteRunBuilder::<Int16Type, T>::new();
397        builder.append_value(values[0]);
398        builder.append_value(values[0]);
399        builder.append_value(values[0]);
400        builder.append_null();
401        builder.append_null();
402        builder.append_value(values[1]);
403        builder.append_value(values[1]);
404        builder.append_value(values[2]);
405        builder.append_value(values[2]);
406        builder.append_value(values[2]);
407        builder.append_value(values[2]);
408        let array = builder.finish();
409
410        assert_eq!(array.len(), 11);
411        assert_eq!(array.null_count(), 0);
412        assert_eq!(array.logical_null_count(), 2);
413
414        assert_eq!(array.run_ends().values(), &[3, 5, 7, 11]);
415
416        // Values are polymorphic and so require a downcast.
417        let av = array.values();
418        let ava: &GenericByteArray<T> = av.as_any().downcast_ref::<GenericByteArray<T>>().unwrap();
419
420        assert_eq!(*ava.value(0), *values[0]);
421        assert!(ava.is_null(1));
422        assert_eq!(*ava.value(2), *values[1]);
423        assert_eq!(*ava.value(3), *values[2]);
424    }
425
426    #[test]
427    fn test_string_run_builder() {
428        test_bytes_run_builder::<Utf8Type>(vec!["abc", "def", "ghi"]);
429    }
430
431    #[test]
432    fn test_string_run_builder_with_empty_strings() {
433        test_bytes_run_builder::<Utf8Type>(vec!["abc", "", "ghi"]);
434    }
435
436    #[test]
437    fn test_binary_run_builder() {
438        test_bytes_run_builder::<BinaryType>(vec![b"abc", b"def", b"ghi"]);
439    }
440
441    fn test_bytes_run_builder_finish_cloned<T>(values: Vec<&T::Native>)
442    where
443        T: ByteArrayType,
444        <T as ByteArrayType>::Native: PartialEq,
445        <T as ByteArrayType>::Native: AsRef<<T as ByteArrayType>::Native>,
446    {
447        let mut builder = GenericByteRunBuilder::<Int16Type, T>::new();
448
449        builder.append_value(values[0]);
450        builder.append_null();
451        builder.append_value(values[1]);
452        builder.append_value(values[1]);
453        builder.append_value(values[0]);
454        let mut array: Int16RunArray = builder.finish_cloned();
455
456        assert_eq!(array.len(), 5);
457        assert_eq!(array.null_count(), 0);
458        assert_eq!(array.logical_null_count(), 1);
459
460        assert_eq!(array.run_ends().values(), &[1, 2, 4, 5]);
461
462        // Values are polymorphic and so require a downcast.
463        let av = array.values();
464        let ava: &GenericByteArray<T> = av.as_any().downcast_ref::<GenericByteArray<T>>().unwrap();
465
466        assert_eq!(ava.value(0), values[0]);
467        assert!(ava.is_null(1));
468        assert_eq!(ava.value(2), values[1]);
469        assert_eq!(ava.value(3), values[0]);
470
471        // Append last value before `finish_cloned` (`value[0]`) again and ensure it has only
472        // one entry in final output.
473        builder.append_value(values[0]);
474        builder.append_value(values[0]);
475        builder.append_value(values[1]);
476        array = builder.finish();
477
478        assert_eq!(array.len(), 8);
479        assert_eq!(array.null_count(), 0);
480        assert_eq!(array.logical_null_count(), 1);
481
482        assert_eq!(array.run_ends().values(), &[1, 2, 4, 7, 8]);
483
484        // Values are polymorphic and so require a downcast.
485        let av2 = array.values();
486        let ava2: &GenericByteArray<T> =
487            av2.as_any().downcast_ref::<GenericByteArray<T>>().unwrap();
488
489        assert_eq!(ava2.value(0), values[0]);
490        assert!(ava2.is_null(1));
491        assert_eq!(ava2.value(2), values[1]);
492        // The value appended before and after `finish_cloned` has only one entry.
493        assert_eq!(ava2.value(3), values[0]);
494        assert_eq!(ava2.value(4), values[1]);
495    }
496
497    #[test]
498    fn test_string_run_builder_finish_cloned() {
499        test_bytes_run_builder_finish_cloned::<Utf8Type>(vec!["abc", "def", "ghi"]);
500    }
501
502    #[test]
503    fn test_binary_run_builder_finish_cloned() {
504        test_bytes_run_builder_finish_cloned::<BinaryType>(vec![b"abc", b"def", b"ghi"]);
505    }
506
507    #[test]
508    fn test_extend() {
509        let mut builder = StringRunBuilder::<Int32Type>::new();
510        builder.extend(["a", "a", "a", "", "", "b", "b"].into_iter().map(Some));
511        builder.extend(["b", "cupcakes", "cupcakes"].into_iter().map(Some));
512        let array = builder.finish();
513
514        assert_eq!(array.len(), 10);
515        assert_eq!(array.run_ends().values(), &[3, 5, 8, 10]);
516
517        let str_array = array.values().as_string::<i32>();
518        assert_eq!(str_array.value(0), "a");
519        assert_eq!(str_array.value(1), "");
520        assert_eq!(str_array.value(2), "b");
521        assert_eq!(str_array.value(3), "cupcakes");
522    }
523}
arrow_array/builder/generic_byte_run_builder.rs

arrow_array/builder/
generic_byte_run_builder.rs
