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

arrow_array/builder/
primitive_run_builder.rs
