arrow_array/builder/
map_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::{Array, ArrayRef, MapArray, StructArray};
7use arrow_buffer::Buffer;
8use arrow_buffer::{NullBuffer, NullBufferBuilder};
9use arrow_data::ArrayData;
10use arrow_schema::{ArrowError, DataType, Field, FieldRef};
11use std::any::Any;
12use std::sync::Arc;
13
14/// Builder for [`MapArray`]
15///
16/// ```
17/// # use arrow_array::builder::{Int32Builder, MapBuilder, StringBuilder};
18/// # use arrow_array::{Int32Array, StringArray};
19///
20/// let string_builder = StringBuilder::new();
21/// let int_builder = Int32Builder::with_capacity(4);
22///
23/// // Construct `[{"joe": 1}, {"blogs": 2, "foo": 4}, {}, null]`
24/// let mut builder = MapBuilder::new(None, string_builder, int_builder);
25///
26/// builder.keys().append_value("joe");
27/// builder.values().append_value(1);
28/// builder.append(true).unwrap();
29///
30/// builder.keys().append_value("blogs");
31/// builder.values().append_value(2);
32/// builder.keys().append_value("foo");
33/// builder.values().append_value(4);
34/// builder.append(true).unwrap();
35/// builder.append(true).unwrap();
36/// builder.append(false).unwrap();
37///
38/// let array = builder.finish();
39/// assert_eq!(array.value_offsets(), &[0, 1, 3, 3, 3]);
40/// assert_eq!(array.values().as_ref(), &Int32Array::from(vec![1, 2, 4]));
41/// assert_eq!(array.keys().as_ref(), &StringArray::from(vec!["joe", "blogs", "foo"]));
42///
43/// ```
44#[derive(Debug)]
45pub struct MapBuilder<K: ArrayBuilder, V: ArrayBuilder> {
46    offsets_builder: Vec<i32>,
47    null_buffer_builder: NullBufferBuilder,
48    field_names: MapFieldNames,
49    key_builder: K,
50    value_builder: V,
51    key_field: Option<FieldRef>,
52    value_field: Option<FieldRef>,
53}
54
55/// The [`Field`] names for a [`MapArray`]
56#[derive(Debug, Clone)]
57pub struct MapFieldNames {
58    /// [`Field`] name for map entries
59    pub entry: String,
60    /// [`Field`] name for map key
61    pub key: String,
62    /// [`Field`] name for map value
63    pub value: String,
64}
65
66impl Default for MapFieldNames {
67    fn default() -> Self {
68        Self {
69            entry: Field::MAP_ENTRIES_FIELD_DEFAULT_NAME.to_string(),
70            key: Field::MAP_KEY_FIELD_DEFAULT_NAME.to_string(),
71            value: Field::MAP_VALUE_FIELD_DEFAULT_NAME.to_string(),
72        }
73    }
74}
75
76impl<K: ArrayBuilder, V: ArrayBuilder> MapBuilder<K, V> {
77    /// Creates a new `MapBuilder`
78    pub fn new(field_names: Option<MapFieldNames>, key_builder: K, value_builder: V) -> Self {
79        let capacity = key_builder.len();
80        Self::with_capacity(field_names, key_builder, value_builder, capacity)
81    }
82
83    /// Creates a new `MapBuilder` with capacity
84    pub fn with_capacity(
85        field_names: Option<MapFieldNames>,
86        key_builder: K,
87        value_builder: V,
88        capacity: usize,
89    ) -> Self {
90        let mut offsets_builder = Vec::with_capacity(capacity + 1);
91        offsets_builder.push(0);
92        Self {
93            offsets_builder,
94            null_buffer_builder: NullBufferBuilder::new(capacity),
95            field_names: field_names.unwrap_or_default(),
96            key_builder,
97            value_builder,
98            key_field: None,
99            value_field: None,
100        }
101    }
102
103    /// Override the field passed to [`MapBuilder::new`]
104    ///
105    /// By default, a non-nullable field is created with the name `keys`
106    ///
107    /// Note: [`Self::finish`] and [`Self::finish_cloned`] will panic if the
108    /// field's data type does not match that of `K` or the field is nullable
109    pub fn with_keys_field(self, field: impl Into<FieldRef>) -> Self {
110        Self {
111            key_field: Some(field.into()),
112            ..self
113        }
114    }
115
116    /// Override the field passed to [`MapBuilder::new`]
117    ///
118    /// By default, a nullable field is created with the name `values`
119    ///
120    /// Note: [`Self::finish`] and [`Self::finish_cloned`] will panic if the
121    /// field's data type does not match that of `V`
122    pub fn with_values_field(self, field: impl Into<FieldRef>) -> Self {
123        Self {
124            value_field: Some(field.into()),
125            ..self
126        }
127    }
128
129    /// Returns the key array builder of the map
130    pub fn keys(&mut self) -> &mut K {
131        &mut self.key_builder
132    }
133
134    /// Returns the value array builder of the map
135    pub fn values(&mut self) -> &mut V {
136        &mut self.value_builder
137    }
138
139    /// Returns both the key and value array builders of the map
140    pub fn entries(&mut self) -> (&mut K, &mut V) {
141        (&mut self.key_builder, &mut self.value_builder)
142    }
143
144    /// Validates that key and value builders have equal lengths.
145    #[inline]
146    fn validate_equal_lengths(&self) -> Result<(), ArrowError> {
147        if self.key_builder.len() != self.value_builder.len() {
148            return Err(ArrowError::InvalidArgumentError(format!(
149                "Cannot append to a map builder when its keys and values have unequal lengths of {} and {}",
150                self.key_builder.len(),
151                self.value_builder.len()
152            )));
153        }
154        Ok(())
155    }
156
157    /// Returns the offset of the next map slot, i.e. the current number of entries.
158    ///
159    /// Returns an error if the number of entries does not fit in an `i32` offset.
160    #[inline]
161    fn next_offset(&self) -> Result<i32, ArrowError> {
162        let len = self.key_builder.len();
163        i32::try_from(len).map_err(|_| {
164            ArrowError::InvalidArgumentError(format!(
165                "Cannot append to a map builder with {len} entries: offset exceeds i32::MAX"
166            ))
167        })
168    }
169
170    /// Finish the current map array slot
171    ///
172    /// Returns an error if the key and values builders are in an inconsistent state,
173    /// or if the number of entries exceeds `i32::MAX`.
174    #[inline]
175    pub fn append(&mut self, is_valid: bool) -> Result<(), ArrowError> {
176        self.validate_equal_lengths()?;
177        let offset = self.next_offset()?;
178        self.offsets_builder.push(offset);
179        self.null_buffer_builder.append(is_valid);
180        Ok(())
181    }
182
183    /// Append `n` nulls to this [`MapBuilder`]
184    ///
185    /// Returns an error if the key and values builders are in an inconsistent state,
186    /// or if the number of entries exceeds `i32::MAX`.
187    #[inline]
188    pub fn append_nulls(&mut self, n: usize) -> Result<(), ArrowError> {
189        self.validate_equal_lengths()?;
190        let offset = self.next_offset()?;
191        self.offsets_builder.extend(std::iter::repeat_n(offset, n));
192        self.null_buffer_builder.append_n_nulls(n);
193        Ok(())
194    }
195
196    /// Builds the [`MapArray`]
197    ///
198    /// # Panics
199    ///
200    /// Panics if the fields set with [`Self::with_keys_field`] or
201    /// [`Self::with_values_field`] do not match the data types of the key and value
202    /// builders, or if the keys contain nulls
203    pub fn finish(&mut self) -> MapArray {
204        let len = self.len();
205        // Build the keys
206        let keys_arr = self.key_builder.finish();
207        let values_arr = self.value_builder.finish();
208        let offset_buffer = Buffer::from_vec(std::mem::take(&mut self.offsets_builder));
209        self.offsets_builder.push(0);
210        let null_bit_buffer = self.null_buffer_builder.finish();
211
212        self.finish_helper(keys_arr, values_arr, offset_buffer, null_bit_buffer, len)
213    }
214
215    /// Builds the [`MapArray`] without resetting the builder.
216    ///
217    /// # Panics
218    ///
219    /// Panics if the fields set with [`Self::with_keys_field`] or
220    /// [`Self::with_values_field`] do not match the data types of the key and value
221    /// builders, or if the keys contain nulls
222    pub fn finish_cloned(&self) -> MapArray {
223        let len = self.len();
224        // Build the keys
225        let keys_arr = self.key_builder.finish_cloned();
226        let values_arr = self.value_builder.finish_cloned();
227        let offset_buffer = Buffer::from_slice_ref(self.offsets_builder.as_slice());
228        let nulls = self.null_buffer_builder.finish_cloned();
229        self.finish_helper(keys_arr, values_arr, offset_buffer, nulls, len)
230    }
231
232    fn finish_preserve_values(&mut self) -> MapArray {
233        let len = self.len();
234        // Build the keys
235        let keys_arr = self.key_builder.finish_preserve_values();
236        let values_arr = self.value_builder.finish_preserve_values();
237        let offset_buffer = Buffer::from_vec(std::mem::take(&mut self.offsets_builder));
238        self.offsets_builder.push(0);
239        let null_bit_buffer = self.null_buffer_builder.finish();
240
241        self.finish_helper(keys_arr, values_arr, offset_buffer, null_bit_buffer, len)
242    }
243
244    fn finish_helper(
245        &self,
246        keys_arr: Arc<dyn Array>,
247        values_arr: Arc<dyn Array>,
248        offset_buffer: Buffer,
249        nulls: Option<NullBuffer>,
250        len: usize,
251    ) -> MapArray {
252        assert!(
253            keys_arr.null_count() == 0,
254            "Keys array must have no null values, found {} null value(s)",
255            keys_arr.null_count()
256        );
257
258        let keys_field = match &self.key_field {
259            Some(f) => {
260                assert!(!f.is_nullable(), "Keys field must not be nullable");
261                f.clone()
262            }
263            None => Arc::new(Field::new(
264                self.field_names.key.as_str(),
265                keys_arr.data_type().clone(),
266                false, // always non-nullable
267            )),
268        };
269        let values_field = match &self.value_field {
270            Some(f) => f.clone(),
271            None => Arc::new(Field::new(
272                self.field_names.value.as_str(),
273                values_arr.data_type().clone(),
274                true,
275            )),
276        };
277
278        let struct_array =
279            StructArray::from(vec![(keys_field, keys_arr), (values_field, values_arr)]);
280
281        let map_field = Arc::new(Field::new(
282            self.field_names.entry.as_str(),
283            struct_array.data_type().clone(),
284            false, // always non-nullable
285        ));
286        let array_data = ArrayData::builder(DataType::Map(map_field, false)) // TODO: support sorted keys
287            .len(len)
288            .add_buffer(offset_buffer)
289            .add_child_data(struct_array.into_data())
290            .nulls(nulls);
291
292        // SAFETY: builder is constructed from valid offset buffer and struct child array maintained by the builder
293        let array_data = unsafe { array_data.build_unchecked() };
294
295        MapArray::from(array_data)
296    }
297
298    /// Returns the current null buffer as a slice
299    pub fn validity_slice(&self) -> Option<&[u8]> {
300        self.null_buffer_builder.as_slice()
301    }
302}
303
304impl<K: ArrayBuilder, V: ArrayBuilder> ArrayBuilder for MapBuilder<K, V> {
305    fn len(&self) -> usize {
306        self.null_buffer_builder.len()
307    }
308
309    fn finish(&mut self) -> ArrayRef {
310        Arc::new(self.finish())
311    }
312
313    /// Builds the array without resetting the builder.
314    fn finish_cloned(&self) -> ArrayRef {
315        Arc::new(self.finish_cloned())
316    }
317
318    fn finish_preserve_values(&mut self) -> ArrayRef {
319        Arc::new(self.finish_preserve_values())
320    }
321
322    fn as_any(&self) -> &dyn Any {
323        self
324    }
325
326    fn as_any_mut(&mut self) -> &mut dyn Any {
327        self
328    }
329
330    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
331        self
332    }
333}
334
335#[cfg(test)]
336mod tests {
337    use super::*;
338    use crate::builder::{Int32Builder, StringBuilder, make_builder, tests::PreserveValuesMock};
339    use crate::{Int32Array, StringArray};
340    use std::collections::HashMap;
341
342    #[test]
343    #[should_panic(expected = "Keys array must have no null values, found 1 null value(s)")]
344    fn test_map_builder_with_null_keys_panics() {
345        let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
346        builder.keys().append_null();
347        builder.values().append_value(42);
348        builder.append(true).unwrap();
349
350        builder.finish();
351    }
352
353    #[test]
354    fn test_boxed_map_builder() {
355        let keys_builder = make_builder(&DataType::Utf8, 5);
356        let values_builder = make_builder(&DataType::Int32, 5);
357
358        let mut builder = MapBuilder::new(None, keys_builder, values_builder);
359        builder
360            .keys()
361            .as_any_mut()
362            .downcast_mut::<StringBuilder>()
363            .expect("should be an StringBuilder")
364            .append_value("1");
365        builder
366            .values()
367            .as_any_mut()
368            .downcast_mut::<Int32Builder>()
369            .expect("should be an Int32Builder")
370            .append_value(42);
371        builder.append(true).unwrap();
372
373        let map_array = builder.finish();
374
375        assert_eq!(
376            map_array
377                .keys()
378                .as_any()
379                .downcast_ref::<StringArray>()
380                .expect("should be an StringArray")
381                .value(0),
382            "1"
383        );
384        assert_eq!(
385            map_array
386                .values()
387                .as_any()
388                .downcast_ref::<Int32Array>()
389                .expect("should be an Int32Array")
390                .value(0),
391            42
392        );
393    }
394
395    #[test]
396    fn test_with_values_field() {
397        let value_field = Arc::new(Field::new("bars", DataType::Int32, false));
398        let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new())
399            .with_values_field(value_field.clone());
400        builder.keys().append_value(1);
401        builder.values().append_value(2);
402        builder.append(true).unwrap();
403        builder.append(false).unwrap(); // This is fine as nullability refers to nullability of values
404        builder.keys().append_value(3);
405        builder.values().append_value(4);
406        builder.append(true).unwrap();
407        let map = builder.finish();
408
409        assert_eq!(map.len(), 3);
410        assert_eq!(
411            map.data_type(),
412            &DataType::Map(
413                Arc::new(Field::new(
414                    Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
415                    DataType::Struct(
416                        vec![
417                            Arc::new(Field::new(
418                                Field::MAP_KEY_FIELD_DEFAULT_NAME,
419                                DataType::Int32,
420                                false
421                            )),
422                            value_field.clone()
423                        ]
424                        .into()
425                    ),
426                    false,
427                )),
428                false
429            )
430        );
431
432        builder.keys().append_value(5);
433        builder.values().append_value(6);
434        builder.append(true).unwrap();
435        let map = builder.finish();
436
437        assert_eq!(map.len(), 1);
438        assert_eq!(
439            map.data_type(),
440            &DataType::Map(
441                Arc::new(Field::new(
442                    Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
443                    DataType::Struct(
444                        vec![
445                            Arc::new(Field::new(
446                                Field::MAP_KEY_FIELD_DEFAULT_NAME,
447                                DataType::Int32,
448                                false
449                            )),
450                            value_field
451                        ]
452                        .into()
453                    ),
454                    false,
455                )),
456                false
457            )
458        );
459    }
460
461    #[test]
462    fn test_with_keys_field() {
463        let mut key_metadata = HashMap::new();
464        key_metadata.insert("foo".to_string(), "bar".to_string());
465        let key_field = Arc::new(
466            Field::new("other_key", DataType::Int32, false).with_metadata(key_metadata.clone()),
467        );
468        let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new())
469            .with_keys_field(key_field.clone());
470        builder.keys().append_value(1);
471        builder.values().append_value(2);
472        builder.append(true).unwrap();
473        let map = builder.finish();
474
475        assert_eq!(map.len(), 1);
476        assert_eq!(
477            map.data_type(),
478            &DataType::Map(
479                Arc::new(Field::new(
480                    Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
481                    DataType::Struct(
482                        vec![
483                            Arc::new(
484                                Field::new("other_key", DataType::Int32, false)
485                                    .with_metadata(key_metadata)
486                            ),
487                            Arc::new(Field::new(
488                                Field::MAP_VALUE_FIELD_DEFAULT_NAME,
489                                DataType::Int32,
490                                true
491                            ))
492                        ]
493                        .into()
494                    ),
495                    false,
496                )),
497                false
498            )
499        );
500    }
501
502    #[test]
503    fn test_append_nulls() {
504        let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new());
505
506        builder.keys().append_value(1);
507        builder.values().append_value(100);
508        builder.append(true).unwrap();
509
510        builder.append_nulls(3).unwrap();
511
512        builder.keys().append_value(2);
513        builder.values().append_value(200);
514        builder.append(true).unwrap();
515
516        let map = builder.finish();
517        assert_eq!(map.len(), 5);
518        assert_eq!(map.null_count(), 3);
519        assert!(map.is_valid(0));
520        assert!(map.is_null(1));
521        assert!(map.is_null(2));
522        assert!(map.is_null(3));
523        assert!(map.is_valid(4));
524        assert_eq!(map.value_offsets(), &[0, 1, 1, 1, 1, 2]);
525    }
526
527    #[test]
528    fn test_append_nulls_inconsistent_state() {
529        let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new());
530        // Add a key without a matching value
531        builder.keys().append_value(1);
532
533        let result = builder.append_nulls(2);
534        assert!(result.is_err());
535        assert!(result.unwrap_err().to_string().contains("unequal lengths"));
536    }
537
538    #[test]
539    #[should_panic(expected = "Keys field must not be nullable")]
540    fn test_with_nullable_keys_field() {
541        let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new())
542            .with_keys_field(Arc::new(Field::new(
543                Field::MAP_KEY_FIELD_DEFAULT_NAME,
544                DataType::Int32,
545                true,
546            )));
547
548        builder.keys().append_value(1);
549        builder.values().append_value(2);
550        builder.append(true).unwrap();
551
552        builder.finish();
553    }
554
555    #[test]
556    #[should_panic(expected = "Incorrect datatype")]
557    fn test_keys_field_type_mismatch() {
558        let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new())
559            .with_keys_field(Arc::new(Field::new(
560                Field::MAP_KEY_FIELD_DEFAULT_NAME,
561                DataType::Utf8,
562                false,
563            )));
564
565        builder.keys().append_value(1);
566        builder.values().append_value(2);
567        builder.append(true).unwrap();
568
569        builder.finish();
570    }
571
572    #[test]
573    fn test_finish_preserve_values() {
574        let mut builder = MapBuilder::new(
575            None,
576            PreserveValuesMock::default(),
577            PreserveValuesMock::default(),
578        );
579
580        builder.keys().inner.append_value(1);
581        builder.values().inner.append_value(2);
582        builder.append(true).unwrap();
583
584        let map = builder.finish_preserve_values();
585
586        assert_eq!(1, map.len());
587        assert_eq!(1, builder.keys().called);
588        assert_eq!(1, builder.values().called);
589    }
590
591    /// Mock builder that reports an arbitrary length without allocating anything,
592    /// so offset overflow can be exercised without building `i32::MAX` entries.
593    struct FixedLenBuilder(usize);
594
595    impl ArrayBuilder for FixedLenBuilder {
596        fn as_any(&self) -> &dyn Any {
597            self
598        }
599
600        fn as_any_mut(&mut self) -> &mut dyn Any {
601            self
602        }
603
604        fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
605            self
606        }
607
608        fn len(&self) -> usize {
609            self.0
610        }
611
612        fn finish(&mut self) -> ArrayRef {
613            panic!("finish should never be called on FixedLenBuilder")
614        }
615
616        fn finish_cloned(&self) -> ArrayRef {
617            panic!("finish_cloned should never be called on FixedLenBuilder")
618        }
619    }
620
621    #[test]
622    fn test_append_offset_overflow() {
623        let len = i32::MAX as usize + 1;
624        let mut builder = MapBuilder::new(None, FixedLenBuilder(len), FixedLenBuilder(len));
625
626        let result = builder.append(true);
627        assert!(result.is_err());
628        assert!(result.unwrap_err().to_string().contains("exceeds i32::MAX"));
629
630        let result = builder.append_nulls(1);
631        assert!(result.is_err());
632        assert!(result.unwrap_err().to_string().contains("exceeds i32::MAX"));
633
634        // The builder must be left untouched after a rejected append
635        assert_eq!(builder.offsets_builder, vec![0]);
636        assert_eq!(builder.len(), 0);
637    }
638}
arrow_array/builder/map_builder.rs

arrow_array/builder/
map_builder.rs
