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

arrow_array/builder/
primitive_run_builder.rs
