arrow_array/builder/
primitive_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 std::{any::Any, sync::Arc};
6
7use crate::{ArrayRef, ArrowPrimitiveType, RunArray, types::RunEndIndexType};
8
9use super::{ArrayBuilder, PrimitiveBuilder};
10
11use arrow_buffer::ArrowNativeType;
12
13/// Builder for [`RunArray`] of [`PrimitiveArray`](crate::array::PrimitiveArray)
14///
15/// # Example:
16///
17/// ```
18///
19/// # use arrow_array::builder::PrimitiveRunBuilder;
20/// # use arrow_array::cast::AsArray;
21/// # use arrow_array::types::{UInt32Type, Int16Type};
22/// # use arrow_array::{Array, UInt32Array, Int16Array};
23///
24/// let mut builder =
25/// PrimitiveRunBuilder::<Int16Type, UInt32Type>::new();
26/// builder.append_value(1234);
27/// builder.append_value(1234);
28/// builder.append_value(1234);
29/// builder.append_null();
30/// builder.append_value(5678);
31/// builder.append_value(5678);
32/// let array = builder.finish();
33///
34/// assert_eq!(array.run_ends().values(), &[3, 4, 6]);
35///
36/// let av = array.values();
37///
38/// assert!(!av.is_null(0));
39/// assert!(av.is_null(1));
40/// assert!(!av.is_null(2));
41///
42/// // Values are polymorphic and so require a downcast.
43/// let ava: &UInt32Array = av.as_primitive::<UInt32Type>();
44///
45/// assert_eq!(ava, &UInt32Array::from(vec![Some(1234), None, Some(5678)]));
46/// ```
47#[derive(Debug)]
48pub struct PrimitiveRunBuilder<R, V>
49where
50    R: RunEndIndexType,
51    V: ArrowPrimitiveType,
52{
53    run_ends_builder: PrimitiveBuilder<R>,
54    values_builder: PrimitiveBuilder<V>,
55    current_value: Option<V::Native>,
56    current_run_end_index: usize,
57    prev_run_end_index: usize,
58}
59
60impl<R, V> Default for PrimitiveRunBuilder<R, V>
61where
62    R: RunEndIndexType,
63    V: ArrowPrimitiveType,
64{
65    fn default() -> Self {
66        Self::new()
67    }
68}
69
70impl<R, V> PrimitiveRunBuilder<R, V>
71where
72    R: RunEndIndexType,
73    V: ArrowPrimitiveType,
74{
75    /// Creates a new `PrimitiveRunBuilder`
76    pub fn new() -> Self {
77        Self {
78            run_ends_builder: PrimitiveBuilder::new(),
79            values_builder: PrimitiveBuilder::new(),
80            current_value: None,
81            current_run_end_index: 0,
82            prev_run_end_index: 0,
83        }
84    }
85
86    /// Creates a new `PrimitiveRunBuilder` with the provided capacity
87    ///
88    /// `capacity`: the expected number of run-end encoded values.
89    pub fn with_capacity(capacity: usize) -> Self {
90        Self {
91            run_ends_builder: PrimitiveBuilder::with_capacity(capacity),
92            values_builder: PrimitiveBuilder::with_capacity(capacity),
93            current_value: None,
94            current_run_end_index: 0,
95            prev_run_end_index: 0,
96        }
97    }
98
99    /// Overrides the data type of the values child array.
100    ///
101    /// By default, `V::DATA_TYPE` is used (via [`PrimitiveBuilder`]). This
102    /// allows setting the timezone of a Timestamp, the precision & scale of a
103    /// Decimal, etc.
104    ///
105    /// # Panics
106    ///
107    /// This method panics if `values_builder` rejects `data_type`.
108    pub fn with_data_type(mut self, data_type: arrow_schema::DataType) -> Self {
109        self.values_builder = self.values_builder.with_data_type(data_type);
110        self
111    }
112}
113
114impl<R, V> ArrayBuilder for PrimitiveRunBuilder<R, V>
115where
116    R: RunEndIndexType,
117    V: ArrowPrimitiveType,
118{
119    /// Returns the builder as a non-mutable `Any` reference.
120    fn as_any(&self) -> &dyn Any {
121        self
122    }
123
124    /// Returns the builder as a mutable `Any` reference.
125    fn as_any_mut(&mut self) -> &mut dyn Any {
126        self
127    }
128
129    /// Returns the boxed builder as a box of `Any`.
130    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
131        self
132    }
133
134    /// Returns the length of logical array encoded by
135    /// the eventual runs array.
136    fn len(&self) -> usize {
137        self.current_run_end_index
138    }
139
140    /// Builds the array and reset this builder.
141    fn finish(&mut self) -> ArrayRef {
142        Arc::new(self.finish())
143    }
144
145    /// Builds the array without resetting the builder.
146    fn finish_cloned(&self) -> ArrayRef {
147        Arc::new(self.finish_cloned())
148    }
149}
150
151impl<R, V> PrimitiveRunBuilder<R, V>
152where
153    R: RunEndIndexType,
154    V: ArrowPrimitiveType,
155{
156    /// Appends optional value to the logical array encoded by the RunArray.
157    pub fn append_option(&mut self, value: Option<V::Native>) {
158        if self.current_run_end_index == 0 {
159            self.current_run_end_index = 1;
160            self.current_value = value;
161            return;
162        }
163        if self.current_value != value {
164            self.append_run_end();
165            self.current_value = value;
166        }
167
168        self.current_run_end_index += 1;
169    }
170
171    /// Appends value to the logical array encoded by the run-ends array.
172    pub fn append_value(&mut self, value: V::Native) {
173        self.append_option(Some(value))
174    }
175
176    /// Appends null to the logical array encoded by the run-ends array.
177    pub fn append_null(&mut self) {
178        self.append_option(None)
179    }
180
181    /// Creates the RunArray and resets the builder.
182    ///
183    /// # Panics
184    ///
185    /// Panics if RunArray cannot be built.
186    pub fn finish(&mut self) -> RunArray<R> {
187        // write the last run end to the array.
188        self.append_run_end();
189
190        // reset the run indexes to zero.
191        self.current_value = None;
192        self.current_run_end_index = 0;
193        self.prev_run_end_index = 0;
194
195        // build the run encoded array by adding run_ends and values array as its children.
196        let run_ends_array = self.run_ends_builder.finish();
197        let values_array = self.values_builder.finish();
198        RunArray::<R>::try_new(&run_ends_array, &values_array).unwrap()
199    }
200
201    /// Creates the RunArray and without resetting the builder.
202    ///
203    /// # Panics
204    ///
205    /// Panics if RunArray cannot be built.
206    pub fn finish_cloned(&self) -> RunArray<R> {
207        let mut run_ends_array = self.run_ends_builder.finish_cloned();
208        let mut values_array = self.values_builder.finish_cloned();
209
210        // Add current run if one exists
211        if self.prev_run_end_index != self.current_run_end_index {
212            let mut run_end_builder = run_ends_array.into_builder().unwrap();
213            let mut values_builder = values_array.into_builder().unwrap();
214            self.append_run_end_with_builders(&mut run_end_builder, &mut values_builder);
215            run_ends_array = run_end_builder.finish();
216            values_array = values_builder.finish();
217        }
218
219        RunArray::try_new(&run_ends_array, &values_array).unwrap()
220    }
221
222    // Appends the current run to the array.
223    fn append_run_end(&mut self) {
224        // empty array or the function called without appending any value.
225        if self.prev_run_end_index == self.current_run_end_index {
226            return;
227        }
228        let run_end_index = self.run_end_index_as_native();
229        self.run_ends_builder.append_value(run_end_index);
230        self.values_builder.append_option(self.current_value);
231        self.prev_run_end_index = self.current_run_end_index;
232    }
233
234    // Similar to `append_run_end` but on custom builders.
235    // Used in `finish_cloned` which is not suppose to mutate `self`.
236    fn append_run_end_with_builders(
237        &self,
238        run_ends_builder: &mut PrimitiveBuilder<R>,
239        values_builder: &mut PrimitiveBuilder<V>,
240    ) {
241        let run_end_index = self.run_end_index_as_native();
242        run_ends_builder.append_value(run_end_index);
243        values_builder.append_option(self.current_value);
244    }
245
246    fn run_end_index_as_native(&self) -> R::Native {
247        R::Native::from_usize(self.current_run_end_index)
248        .unwrap_or_else(|| panic!(
249                "Cannot convert `current_run_end_index` {} from `usize` to native form of arrow datatype {}",
250                self.current_run_end_index,
251                R::DATA_TYPE
252        ))
253    }
254}
255
256impl<R, V> Extend<Option<V::Native>> for PrimitiveRunBuilder<R, V>
257where
258    R: RunEndIndexType,
259    V: ArrowPrimitiveType,
260{
261    fn extend<T: IntoIterator<Item = Option<V::Native>>>(&mut self, iter: T) {
262        for elem in iter {
263            self.append_option(elem);
264        }
265    }
266}
267
268#[cfg(test)]
269mod tests {
270    use arrow_schema::DataType;
271
272    use crate::builder::PrimitiveRunBuilder;
273    use crate::cast::AsArray;
274    use crate::types::{Decimal128Type, Int16Type, TimestampMicrosecondType, UInt32Type};
275    use crate::{Array, Decimal128Array, TimestampMicrosecondArray, UInt32Array};
276
277    #[test]
278    fn test_primitive_ree_array_builder() {
279        let mut builder = PrimitiveRunBuilder::<Int16Type, UInt32Type>::new();
280        builder.append_value(1234);
281        builder.append_value(1234);
282        builder.append_value(1234);
283        builder.append_null();
284        builder.append_value(5678);
285        builder.append_value(5678);
286
287        let array = builder.finish();
288
289        assert_eq!(array.null_count(), 0);
290        assert_eq!(array.logical_null_count(), 1);
291        assert_eq!(array.len(), 6);
292
293        assert_eq!(array.run_ends().values(), &[3, 4, 6]);
294
295        let av = array.values();
296
297        assert!(!av.is_null(0));
298        assert!(av.is_null(1));
299        assert!(!av.is_null(2));
300
301        // Values are polymorphic and so require a downcast.
302        let ava: &UInt32Array = av.as_primitive::<UInt32Type>();
303
304        assert_eq!(ava, &UInt32Array::from(vec![Some(1234), None, Some(5678)]));
305    }
306
307    #[test]
308    fn test_extend() {
309        let mut builder = PrimitiveRunBuilder::<Int16Type, Int16Type>::new();
310        builder.extend([1, 2, 2, 5, 5, 4, 4].into_iter().map(Some));
311        builder.extend([4, 4, 6, 2].into_iter().map(Some));
312        let array = builder.finish();
313
314        assert_eq!(array.len(), 11);
315        assert_eq!(array.null_count(), 0);
316        assert_eq!(array.logical_null_count(), 0);
317        assert_eq!(array.run_ends().values(), &[1, 3, 5, 9, 10, 11]);
318        assert_eq!(
319            array.values().as_primitive::<Int16Type>().values(),
320            &[1, 2, 5, 4, 6, 2]
321        );
322    }
323
324    #[test]
325    fn test_reuse_after_finish() {
326        let mut builder = PrimitiveRunBuilder::<Int16Type, Int16Type>::new();
327        builder.extend([1, 2, 3].into_iter().map(Some));
328        builder.finish();
329
330        // A single run with the same length as the previous batch
331        builder.extend([7, 7, 7].into_iter().map(Some));
332        let array = builder.finish();
333        assert_eq!(array.len(), 3);
334        assert_eq!(array.run_ends().values(), &[3]);
335        assert_eq!(array.values().as_primitive::<Int16Type>().values(), &[7]);
336
337        // Nothing appended since the last finish
338        assert_eq!(builder.finish().len(), 0);
339    }
340
341    #[test]
342    #[should_panic(expected = "incompatible data type for builder")]
343    fn test_override_data_type_invalid() {
344        PrimitiveRunBuilder::<Int16Type, UInt32Type>::new().with_data_type(DataType::UInt64);
345    }
346
347    #[test]
348    fn test_override_data_type() {
349        // Noop.
350        PrimitiveRunBuilder::<Int16Type, UInt32Type>::new().with_data_type(DataType::UInt32);
351
352        // Setting scale & precision.
353        let mut builder = PrimitiveRunBuilder::<Int16Type, Decimal128Type>::new()
354            .with_data_type(DataType::Decimal128(1, 2));
355        builder.append_value(123);
356        let array = builder.finish();
357        let array = array.downcast::<Decimal128Array>().unwrap();
358        let values = array.values();
359        assert_eq!(values.precision(), 1);
360        assert_eq!(values.scale(), 2);
361
362        // Setting timezone.
363        let mut builder = PrimitiveRunBuilder::<Int16Type, TimestampMicrosecondType>::new()
364            .with_data_type(DataType::Timestamp(
365                arrow_schema::TimeUnit::Microsecond,
366                Some("Europe/Paris".into()),
367            ));
368        builder.append_value(1);
369        let array = builder.finish();
370        let array = array.downcast::<TimestampMicrosecondArray>().unwrap();
371        let values = array.values();
372        assert_eq!(values.timezone(), Some("Europe/Paris"));
373    }
374
375    #[test]
376    fn test_finish_cloned_keeps_data_type() {
377        let data_type = DataType::Decimal128(5, 2);
378        let mut builder = PrimitiveRunBuilder::<Int16Type, Decimal128Type>::new()
379            .with_data_type(data_type.clone());
380        // Leave the run open, so `finish_cloned` has to add it to a copy of the values
381        builder.append_value(123);
382        let array = builder.finish_cloned();
383        assert_eq!(array.values().data_type(), &data_type);
384    }
385}
arrow_array/builder/primitive_run_builder.rs

arrow_array/builder/
primitive_run_builder.rs
