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 indexes to zero.
204        self.current_value = None;
205        self.current_run_end_index = 0;
206        self.prev_run_end_index = 0;
207
208        // build the run encoded array by adding run_ends and values array as its children.
209        let run_ends_array = self.run_ends_builder.finish();
210        let values_array = self.values_builder.finish();
211        RunArray::<R>::try_new(&run_ends_array, &values_array).unwrap()
212    }
213
214    /// Creates the RunArray and without resetting the builder.
215    ///
216    /// # Panics
217    ///
218    /// Panics if RunArray cannot be built.
219    pub fn finish_cloned(&self) -> RunArray<R> {
220        let mut run_ends_array = self.run_ends_builder.finish_cloned();
221        let mut values_array = self.values_builder.finish_cloned();
222
223        // Add current run if one exists
224        if self.prev_run_end_index != self.current_run_end_index {
225            let mut run_end_builder = run_ends_array.into_builder().unwrap();
226            let mut values_builder = values_array.into_builder().unwrap();
227            self.append_run_end_with_builders(&mut run_end_builder, &mut values_builder);
228            run_ends_array = run_end_builder.finish();
229            values_array = values_builder.finish();
230        }
231
232        RunArray::try_new(&run_ends_array, &values_array).unwrap()
233    }
234
235    // Appends the current run to the array.
236    fn append_run_end(&mut self) {
237        // empty array or the function called without appending any value.
238        if self.prev_run_end_index == self.current_run_end_index {
239            return;
240        }
241        let run_end_index = self.run_end_index_as_native();
242        self.run_ends_builder.append_value(run_end_index);
243        self.values_builder.append_option(self.current_value);
244        self.prev_run_end_index = self.current_run_end_index;
245    }
246
247    // Similar to `append_run_end` but on custom builders.
248    // Used in `finish_cloned` which is not suppose to mutate `self`.
249    fn append_run_end_with_builders(
250        &self,
251        run_ends_builder: &mut PrimitiveBuilder<R>,
252        values_builder: &mut PrimitiveBuilder<V>,
253    ) {
254        let run_end_index = self.run_end_index_as_native();
255        run_ends_builder.append_value(run_end_index);
256        values_builder.append_option(self.current_value);
257    }
258
259    fn run_end_index_as_native(&self) -> R::Native {
260        R::Native::from_usize(self.current_run_end_index)
261        .unwrap_or_else(|| panic!(
262                "Cannot convert `current_run_end_index` {} from `usize` to native form of arrow datatype {}",
263                self.current_run_end_index,
264                R::DATA_TYPE
265        ))
266    }
267}
268
269impl<R, V> Extend<Option<V::Native>> for PrimitiveRunBuilder<R, V>
270where
271    R: RunEndIndexType,
272    V: ArrowPrimitiveType,
273{
274    fn extend<T: IntoIterator<Item = Option<V::Native>>>(&mut self, iter: T) {
275        for elem in iter {
276            self.append_option(elem);
277        }
278    }
279}
280
281#[cfg(test)]
282mod tests {
283    use arrow_schema::DataType;
284
285    use crate::builder::PrimitiveRunBuilder;
286    use crate::cast::AsArray;
287    use crate::types::{Decimal128Type, Int16Type, TimestampMicrosecondType, UInt32Type};
288    use crate::{Array, Decimal128Array, TimestampMicrosecondArray, UInt32Array};
289
290    #[test]
291    fn test_primitive_ree_array_builder() {
292        let mut builder = PrimitiveRunBuilder::<Int16Type, UInt32Type>::new();
293        builder.append_value(1234);
294        builder.append_value(1234);
295        builder.append_value(1234);
296        builder.append_null();
297        builder.append_value(5678);
298        builder.append_value(5678);
299
300        let array = builder.finish();
301
302        assert_eq!(array.null_count(), 0);
303        assert_eq!(array.logical_null_count(), 1);
304        assert_eq!(array.len(), 6);
305
306        assert_eq!(array.run_ends().values(), &[3, 4, 6]);
307
308        let av = array.values();
309
310        assert!(!av.is_null(0));
311        assert!(av.is_null(1));
312        assert!(!av.is_null(2));
313
314        // Values are polymorphic and so require a downcast.
315        let ava: &UInt32Array = av.as_primitive::<UInt32Type>();
316
317        assert_eq!(ava, &UInt32Array::from(vec![Some(1234), None, Some(5678)]));
318    }
319
320    #[test]
321    fn test_extend() {
322        let mut builder = PrimitiveRunBuilder::<Int16Type, Int16Type>::new();
323        builder.extend([1, 2, 2, 5, 5, 4, 4].into_iter().map(Some));
324        builder.extend([4, 4, 6, 2].into_iter().map(Some));
325        let array = builder.finish();
326
327        assert_eq!(array.len(), 11);
328        assert_eq!(array.null_count(), 0);
329        assert_eq!(array.logical_null_count(), 0);
330        assert_eq!(array.run_ends().values(), &[1, 3, 5, 9, 10, 11]);
331        assert_eq!(
332            array.values().as_primitive::<Int16Type>().values(),
333            &[1, 2, 5, 4, 6, 2]
334        );
335    }
336
337    #[test]
338    fn test_reuse_after_finish() {
339        let mut builder = PrimitiveRunBuilder::<Int16Type, Int16Type>::new();
340        builder.extend([1, 2, 3].into_iter().map(Some));
341        builder.finish();
342
343        // A single run with the same length as the previous batch
344        builder.extend([7, 7, 7].into_iter().map(Some));
345        let array = builder.finish();
346        assert_eq!(array.len(), 3);
347        assert_eq!(array.run_ends().values(), &[3]);
348        assert_eq!(array.values().as_primitive::<Int16Type>().values(), &[7]);
349
350        // Nothing appended since the last finish
351        assert_eq!(builder.finish().len(), 0);
352    }
353
354    #[test]
355    #[should_panic(expected = "incompatible data type for builder")]
356    fn test_override_data_type_invalid() {
357        PrimitiveRunBuilder::<Int16Type, UInt32Type>::new().with_data_type(DataType::UInt64);
358    }
359
360    #[test]
361    fn test_override_data_type() {
362        // Noop.
363        PrimitiveRunBuilder::<Int16Type, UInt32Type>::new().with_data_type(DataType::UInt32);
364
365        // Setting scale & precision.
366        let mut builder = PrimitiveRunBuilder::<Int16Type, Decimal128Type>::new()
367            .with_data_type(DataType::Decimal128(1, 2));
368        builder.append_value(123);
369        let array = builder.finish();
370        let array = array.downcast::<Decimal128Array>().unwrap();
371        let values = array.values();
372        assert_eq!(values.precision(), 1);
373        assert_eq!(values.scale(), 2);
374
375        // Setting timezone.
376        let mut builder = PrimitiveRunBuilder::<Int16Type, TimestampMicrosecondType>::new()
377            .with_data_type(DataType::Timestamp(
378                arrow_schema::TimeUnit::Microsecond,
379                Some("Europe/Paris".into()),
380            ));
381        builder.append_value(1);
382        let array = builder.finish();
383        let array = array.downcast::<TimestampMicrosecondArray>().unwrap();
384        let values = array.values();
385        assert_eq!(values.timezone(), Some("Europe/Paris"));
386    }
387
388    #[test]
389    fn test_finish_cloned_keeps_data_type() {
390        let data_type = DataType::Decimal128(5, 2);
391        let mut builder = PrimitiveRunBuilder::<Int16Type, Decimal128Type>::new()
392            .with_data_type(data_type.clone());
393        // Leave the run open, so `finish_cloned` has to add it to a copy of the values
394        builder.append_value(123);
395        let array = builder.finish_cloned();
396        assert_eq!(array.values().data_type(), &data_type);
397    }
398}
arrow_array/builder/primitive_run_builder.rs

arrow_array/builder/
primitive_run_builder.rs
