arrow_array/builder/
fixed_size_list_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::builder::ArrayBuilder;
6use crate::{ArrayRef, FixedSizeListArray};
7use arrow_buffer::NullBufferBuilder;
8use arrow_schema::{Field, FieldRef};
9use std::any::Any;
10use std::sync::Arc;
11
12///  Builder for [`FixedSizeListArray`]
13/// ```
14/// use arrow_array::{builder::{Int32Builder, FixedSizeListBuilder}, Array, Int32Array};
15/// let values_builder = Int32Builder::new();
16/// let mut builder = FixedSizeListBuilder::new(values_builder, 3);
17///
18/// //  [[0, 1, 2], null, [3, null, 5], [6, 7, null]]
19/// builder.values().append_value(0);
20/// builder.values().append_value(1);
21/// builder.values().append_value(2);
22/// builder.append(true);
23/// builder.values().append_null();
24/// builder.values().append_null();
25/// builder.values().append_null();
26/// builder.append(false);
27/// builder.values().append_value(3);
28/// builder.values().append_null();
29/// builder.values().append_value(5);
30/// builder.append(true);
31/// builder.values().append_value(6);
32/// builder.values().append_value(7);
33/// builder.values().append_null();
34/// builder.append(true);
35/// let list_array = builder.finish();
36/// assert_eq!(
37///     *list_array.value(0),
38///     Int32Array::from(vec![Some(0), Some(1), Some(2)])
39/// );
40/// assert!(list_array.is_null(1));
41/// assert_eq!(
42///     *list_array.value(2),
43///     Int32Array::from(vec![Some(3), None, Some(5)])
44/// );
45/// assert_eq!(
46///     *list_array.value(3),
47///     Int32Array::from(vec![Some(6), Some(7), None])
48/// )
49/// ```
50///
51#[derive(Debug)]
52pub struct FixedSizeListBuilder<T: ArrayBuilder> {
53    null_buffer_builder: NullBufferBuilder,
54    values_builder: T,
55    list_len: i32,
56    field: Option<FieldRef>,
57}
58
59impl<T: ArrayBuilder> FixedSizeListBuilder<T> {
60    /// Creates a new [`FixedSizeListBuilder`] from a given values array builder
61    /// `value_length` is the number of values within each array
62    pub fn new(values_builder: T, value_length: i32) -> Self {
63        let capacity = values_builder
64            .len()
65            .checked_div(value_length as _)
66            .unwrap_or_default();
67
68        Self::with_capacity(values_builder, value_length, capacity)
69    }
70
71    /// Creates a new [`FixedSizeListBuilder`] from a given values array builder
72    /// `value_length` is the number of values within each array
73    /// `capacity` is the number of items to pre-allocate space for in this builder
74    pub fn with_capacity(values_builder: T, value_length: i32, capacity: usize) -> Self {
75        Self {
76            null_buffer_builder: NullBufferBuilder::new(capacity),
77            values_builder,
78            list_len: value_length,
79            field: None,
80        }
81    }
82
83    /// Override the field passed to [`FixedSizeListArray::new`]
84    ///
85    /// By default, a nullable field is created with the name `item`
86    ///
87    /// Note: [`Self::finish`] and [`Self::finish_cloned`] will panic if the
88    /// field's data type does not match that of `T`
89    pub fn with_field(self, field: impl Into<FieldRef>) -> Self {
90        Self {
91            field: Some(field.into()),
92            ..self
93        }
94    }
95}
96
97impl<T: ArrayBuilder> ArrayBuilder for FixedSizeListBuilder<T>
98where
99    T: 'static,
100{
101    /// Returns the builder as a non-mutable `Any` reference.
102    fn as_any(&self) -> &dyn Any {
103        self
104    }
105
106    /// Returns the builder as a mutable `Any` reference.
107    fn as_any_mut(&mut self) -> &mut dyn Any {
108        self
109    }
110
111    /// Returns the boxed builder as a box of `Any`.
112    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
113        self
114    }
115
116    /// Returns the number of array slots in the builder
117    fn len(&self) -> usize {
118        self.null_buffer_builder.len()
119    }
120
121    /// Builds the array and reset this builder.
122    fn finish(&mut self) -> ArrayRef {
123        Arc::new(self.finish())
124    }
125
126    /// Builds the array without resetting the builder.
127    fn finish_cloned(&self) -> ArrayRef {
128        Arc::new(self.finish_cloned())
129    }
130
131    fn finish_preserve_values(&mut self) -> ArrayRef {
132        Arc::new(self.finish_preserve_values())
133    }
134}
135
136impl<T: ArrayBuilder> FixedSizeListBuilder<T>
137where
138    T: 'static,
139{
140    /// Returns the child array builder as a mutable reference.
141    ///
142    /// This mutable reference can be used to append values into the child array builder,
143    /// but you must call [`append`](#method.append) to delimit each distinct list value.
144    pub fn values(&mut self) -> &mut T {
145        &mut self.values_builder
146    }
147
148    /// Returns the length of the list
149    pub fn value_length(&self) -> i32 {
150        self.list_len
151    }
152
153    /// Finish the current fixed-length list array slot
154    #[inline]
155    pub fn append(&mut self, is_valid: bool) {
156        self.null_buffer_builder.append(is_valid);
157    }
158
159    /// Builds the [`FixedSizeListBuilder`] and reset this builder.
160    ///
161    /// # Panics
162    ///
163    /// Panics if the length of the child array is not `self.len() * value_length`
164    pub fn finish(&mut self) -> FixedSizeListArray {
165        let len = self.len();
166        let values = self.values_builder.finish();
167        let nulls = self.null_buffer_builder.finish();
168
169        assert_eq!(
170            values.len(),
171            len * self.list_len as usize,
172            "Length of the child array ({}) must be the multiple of the value length ({}) and the array length ({}).",
173            values.len(),
174            self.list_len,
175            len,
176        );
177
178        let field = self
179            .field
180            .clone()
181            .unwrap_or_else(|| Arc::new(Field::new_list_field(values.data_type().clone(), true)));
182
183        FixedSizeListArray::new(field, self.list_len, values, nulls)
184    }
185
186    /// Builds the [`FixedSizeListBuilder`] without resetting the builder.
187    ///
188    /// # Panics
189    ///
190    /// Panics if the length of the child array is not `self.len() * value_length`
191    pub fn finish_cloned(&self) -> FixedSizeListArray {
192        let len = self.len();
193        let values = self.values_builder.finish_cloned();
194        let nulls = self.null_buffer_builder.finish_cloned();
195
196        assert_eq!(
197            values.len(),
198            len * self.list_len as usize,
199            "Length of the child array ({}) must be the multiple of the value length ({}) and the array length ({}).",
200            values.len(),
201            self.list_len,
202            len,
203        );
204
205        let field = self
206            .field
207            .clone()
208            .unwrap_or_else(|| Arc::new(Field::new_list_field(values.data_type().clone(), true)));
209
210        FixedSizeListArray::new(field, self.list_len, values, nulls)
211    }
212
213    fn finish_preserve_values(&mut self) -> FixedSizeListArray {
214        let len = self.len();
215        let values = self.values_builder.finish_preserve_values();
216        let nulls = self.null_buffer_builder.finish();
217
218        assert_eq!(
219            values.len(),
220            len * self.list_len as usize,
221            "Length of the child array ({}) must be the multiple of the value length ({}) and the array length ({}).",
222            values.len(),
223            self.list_len,
224            len,
225        );
226
227        let field = self
228            .field
229            .clone()
230            .unwrap_or_else(|| Arc::new(Field::new_list_field(values.data_type().clone(), true)));
231
232        FixedSizeListArray::new(field, self.list_len, values, nulls)
233    }
234
235    /// Returns the current null buffer as a slice
236    pub fn validity_slice(&self) -> Option<&[u8]> {
237        self.null_buffer_builder.as_slice()
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244    use arrow_schema::DataType;
245
246    use crate::Array;
247    use crate::Int32Array;
248    use crate::builder::{Int32Builder, tests::PreserveValuesMock};
249
250    fn make_list_builder(
251        include_null_element: bool,
252        include_null_in_values: bool,
253    ) -> FixedSizeListBuilder<crate::builder::PrimitiveBuilder<crate::types::Int32Type>> {
254        let values_builder = Int32Builder::new();
255        let mut builder = FixedSizeListBuilder::new(values_builder, 3);
256
257        builder.values().append_value(0);
258        builder.values().append_value(1);
259        builder.values().append_value(2);
260        builder.append(true);
261
262        builder.values().append_value(2);
263        builder.values().append_value(3);
264        builder.values().append_value(4);
265        builder.append(true);
266
267        if include_null_element {
268            builder.values().append_null();
269            builder.values().append_null();
270            builder.values().append_null();
271            builder.append(false);
272        } else {
273            builder.values().append_value(2);
274            builder.values().append_value(3);
275            builder.values().append_value(4);
276            builder.append(true);
277        }
278
279        builder.values().append_value(3);
280        if include_null_in_values {
281            builder.values().append_null();
282        } else {
283            builder.values().append_value(4);
284        }
285        builder.values().append_value(5);
286        builder.append(true);
287
288        builder
289    }
290
291    #[test]
292    fn test_fixed_size_list_array_builder() {
293        let mut builder = make_list_builder(true, true);
294
295        let list_array = builder.finish();
296
297        assert_eq!(DataType::Int32, list_array.value_type());
298        assert_eq!(4, list_array.len());
299        assert_eq!(1, list_array.null_count());
300        assert_eq!(6, list_array.value_offset_at(2));
301        assert_eq!(3, list_array.value_length());
302    }
303
304    #[test]
305    fn test_fixed_size_list_array_builder_with_field() {
306        let builder = make_list_builder(false, false);
307        let mut builder = builder.with_field(Field::new("list_element", DataType::Int32, false));
308        let list_array = builder.finish();
309
310        assert_eq!(DataType::Int32, list_array.value_type());
311        assert_eq!(4, list_array.len());
312        assert_eq!(0, list_array.null_count());
313        assert_eq!(6, list_array.value_offset_at(2));
314        assert_eq!(3, list_array.value_length());
315    }
316
317    #[test]
318    fn test_fixed_size_list_array_builder_with_field_and_null() {
319        let builder = make_list_builder(true, false);
320        let mut builder = builder.with_field(Field::new("list_element", DataType::Int32, false));
321        let list_array = builder.finish();
322
323        assert_eq!(DataType::Int32, list_array.value_type());
324        assert_eq!(4, list_array.len());
325        assert_eq!(1, list_array.null_count());
326        assert_eq!(6, list_array.value_offset_at(2));
327        assert_eq!(3, list_array.value_length());
328    }
329
330    #[test]
331    #[should_panic(expected = "Found unmasked nulls for non-nullable FixedSizeListArray")]
332    fn test_fixed_size_list_array_builder_with_field_null_panic() {
333        let builder = make_list_builder(true, true);
334        let mut builder = builder.with_field(Field::new("list_item", DataType::Int32, false));
335
336        builder.finish();
337    }
338
339    #[test]
340    #[should_panic(expected = "FixedSizeListArray expected data type Int64 got Int32")]
341    fn test_fixed_size_list_array_builder_with_field_type_panic() {
342        let values_builder = Int32Builder::new();
343        let builder = FixedSizeListBuilder::new(values_builder, 3);
344        let mut builder = builder.with_field(Field::new("list_item", DataType::Int64, true));
345
346        //  [[0, 1, 2], null, [3, null, 5], [6, 7, null]]
347        builder.values().append_value(0);
348        builder.values().append_value(1);
349        builder.values().append_value(2);
350        builder.append(true);
351        builder.values().append_null();
352        builder.values().append_null();
353        builder.values().append_null();
354        builder.append(false);
355        builder.values().append_value(3);
356        builder.values().append_value(4);
357        builder.values().append_value(5);
358        builder.append(true);
359
360        builder.finish();
361    }
362
363    #[test]
364    fn test_fixed_size_list_array_builder_cloned_with_field() {
365        let builder = make_list_builder(true, true);
366        let builder = builder.with_field(Field::new("list_element", DataType::Int32, true));
367
368        let list_array = builder.finish_cloned();
369
370        assert_eq!(DataType::Int32, list_array.value_type());
371        assert_eq!(4, list_array.len());
372        assert_eq!(1, list_array.null_count());
373        assert_eq!(6, list_array.value_offset_at(2));
374        assert_eq!(3, list_array.value_length());
375    }
376
377    #[test]
378    #[should_panic(expected = "Found unmasked nulls for non-nullable FixedSizeListArray")]
379    fn test_fixed_size_list_array_builder_cloned_with_field_null_panic() {
380        let builder = make_list_builder(true, true);
381        let builder = builder.with_field(Field::new("list_item", DataType::Int32, false));
382
383        builder.finish_cloned();
384    }
385
386    #[test]
387    fn test_fixed_size_list_array_builder_cloned_with_field_and_null() {
388        let builder = make_list_builder(true, false);
389        let mut builder = builder.with_field(Field::new("list_element", DataType::Int32, false));
390        let list_array = builder.finish();
391
392        assert_eq!(DataType::Int32, list_array.value_type());
393        assert_eq!(4, list_array.len());
394        assert_eq!(1, list_array.null_count());
395        assert_eq!(6, list_array.value_offset_at(2));
396        assert_eq!(3, list_array.value_length());
397    }
398
399    #[test]
400    #[should_panic(expected = "FixedSizeListArray expected data type Int64 got Int32")]
401    fn test_fixed_size_list_array_builder_cloned_with_field_type_panic() {
402        let builder = make_list_builder(false, false);
403        let builder = builder.with_field(Field::new("list_item", DataType::Int64, true));
404
405        builder.finish_cloned();
406    }
407
408    #[test]
409    fn test_fixed_size_list_array_builder_finish_cloned() {
410        let mut builder = make_list_builder(true, true);
411
412        let mut list_array = builder.finish_cloned();
413
414        assert_eq!(DataType::Int32, list_array.value_type());
415        assert_eq!(4, list_array.len());
416        assert_eq!(1, list_array.null_count());
417        assert_eq!(3, list_array.value_length());
418
419        builder.values().append_value(6);
420        builder.values().append_value(7);
421        builder.values().append_null();
422        builder.append(true);
423        builder.values().append_null();
424        builder.values().append_null();
425        builder.values().append_null();
426        builder.append(false);
427        list_array = builder.finish();
428
429        assert_eq!(DataType::Int32, list_array.value_type());
430        assert_eq!(6, list_array.len());
431        assert_eq!(2, list_array.null_count());
432        assert_eq!(6, list_array.value_offset_at(2));
433        assert_eq!(3, list_array.value_length());
434    }
435
436    #[test]
437    fn test_fixed_size_list_array_builder_with_field_empty() {
438        let values_builder = Int32Array::builder(0);
439        let mut builder = FixedSizeListBuilder::new(values_builder, 3).with_field(Field::new(
440            "list_item",
441            DataType::Int32,
442            false,
443        ));
444        assert!(builder.is_empty());
445        let arr = builder.finish();
446        assert_eq!(0, arr.len());
447        assert_eq!(0, builder.len());
448    }
449
450    #[test]
451    fn test_fixed_size_list_array_builder_cloned_with_field_empty() {
452        let values_builder = Int32Array::builder(0);
453        let builder = FixedSizeListBuilder::new(values_builder, 3).with_field(Field::new(
454            "list_item",
455            DataType::Int32,
456            false,
457        ));
458        assert!(builder.is_empty());
459        let arr = builder.finish_cloned();
460        assert_eq!(0, arr.len());
461        assert_eq!(0, builder.len());
462    }
463
464    #[test]
465    fn test_fixed_size_list_array_builder_empty() {
466        let values_builder = Int32Array::builder(5);
467        let mut builder = FixedSizeListBuilder::new(values_builder, 3);
468        assert!(builder.is_empty());
469        let arr = builder.finish();
470        assert_eq!(0, arr.len());
471        assert_eq!(0, builder.len());
472    }
473
474    #[test]
475    fn test_fixed_size_list_array_builder_finish() {
476        let values_builder = Int32Array::builder(5);
477        let mut builder = FixedSizeListBuilder::new(values_builder, 3);
478
479        builder.values().append_slice(&[1, 2, 3]);
480        builder.append(true);
481        builder.values().append_slice(&[4, 5, 6]);
482        builder.append(true);
483
484        let mut arr = builder.finish();
485        assert_eq!(2, arr.len());
486        assert_eq!(0, builder.len());
487
488        builder.values().append_slice(&[7, 8, 9]);
489        builder.append(true);
490        arr = builder.finish();
491        assert_eq!(1, arr.len());
492        assert_eq!(0, builder.len());
493    }
494
495    #[test]
496    #[should_panic(
497        expected = "Length of the child array (10) must be the multiple of the value length (3) and the array length (3)."
498    )]
499    fn test_fixed_size_list_array_builder_fail() {
500        let values_builder = Int32Array::builder(5);
501        let mut builder = FixedSizeListBuilder::new(values_builder, 3);
502
503        builder.values().append_slice(&[1, 2, 3]);
504        builder.append(true);
505        builder.values().append_slice(&[4, 5, 6]);
506        builder.append(true);
507        builder.values().append_slice(&[7, 8, 9, 10]);
508        builder.append(true);
509
510        builder.finish();
511    }
512
513    #[test]
514    fn test_finish_preserve_values() {
515        let mut builder = FixedSizeListBuilder::new(PreserveValuesMock::default(), 2);
516
517        builder.values().inner.append_value(0);
518        builder.values().inner.append_value(1);
519        builder.append(true);
520
521        let arr = builder.finish_preserve_values();
522
523        assert_eq!(1, arr.len());
524        assert_eq!(1, builder.values().called);
525    }
526}
arrow_array/builder/fixed_size_list_builder.rs

arrow_array/builder/
fixed_size_list_builder.rs
