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

arrow_array/builder/
map_builder.rs
