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

arrow_array/builder/
fixed_size_list_builder.rs
