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

arrow_array/builder/
generic_byte_run_builder.rs
