arrow_array/array/
map_array.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::array::{get_offsets_from_buffer, print_long_array};
19use crate::builder::MapFieldNames;
20use crate::iterator::MapArrayIter;
21use crate::{Array, ArrayAccessor, ArrayRef, ListArray, StringArray, StructArray, make_array};
22use arrow_buffer::{ArrowNativeType, Buffer, NullBuffer, OffsetBuffer, ToByteSlice};
23use arrow_data::{ArrayData, ArrayDataBuilder};
24use arrow_schema::{ArrowError, DataType, Field, FieldRef, Fields};
25use std::any::Any;
26use std::sync::Arc;
27
28/// An array of key-value maps
29///
30/// Keys should always be non-null, but values can be null.
31///
32/// [`MapArray`] is physically a [`ListArray`] of key values pairs stored as an `entries`
33/// [`StructArray`] with 2 child fields.
34///
35/// # See also
36/// * [`MapBuilder`](crate::builder::MapBuilder) for how to construct a [`MapArray`]
37/// * [`Self::from_vec_of_maps`] for ergonomically creating maps for testing
38#[derive(Clone)]
39pub struct MapArray {
40    data_type: DataType,
41    nulls: Option<NullBuffer>,
42    /// The [`StructArray`] that is the direct child of this array
43    entries: StructArray,
44    /// The start and end offsets of each entry
45    value_offsets: OffsetBuffer<i32>,
46}
47
48impl MapArray {
49    /// Create a new [`MapArray`] from the provided parts
50    ///
51    /// See [`MapBuilder`](crate::builder::MapBuilder) for a higher-level interface
52    /// to construct a [`MapArray`]
53    ///
54    /// # Errors
55    ///
56    /// Errors if
57    ///
58    /// * `offsets.len() - 1 != nulls.len()`
59    /// * `offsets.last() > entries.len()`
60    /// * `field.is_nullable()`
61    /// * `entries.null_count() != 0`
62    /// * `entries.columns().len() != 2`
63    /// * `field.data_type() != entries.data_type()`
64    /// * the keys field is nullable
65    pub fn try_new(
66        field: FieldRef,
67        offsets: OffsetBuffer<i32>,
68        entries: StructArray,
69        nulls: Option<NullBuffer>,
70        ordered: bool,
71    ) -> Result<Self, ArrowError> {
72        let len = offsets.len() - 1; // Offsets guaranteed to not be empty
73        let end_offset = offsets.last().unwrap().as_usize();
74        // don't need to check other values of `offsets` because they are checked
75        // during construction of `OffsetBuffer`
76        if end_offset > entries.len() {
77            return Err(ArrowError::InvalidArgumentError(format!(
78                "Max offset of {end_offset} exceeds length of entries {}",
79                entries.len()
80            )));
81        }
82
83        if let Some(n) = nulls.as_ref() {
84            if n.len() != len {
85                return Err(ArrowError::InvalidArgumentError(format!(
86                    "Incorrect length of null buffer for MapArray, expected {len} got {}",
87                    n.len(),
88                )));
89            }
90        }
91        if field.is_nullable() || entries.null_count() != 0 {
92            return Err(ArrowError::InvalidArgumentError(
93                "MapArray entries cannot contain nulls".to_string(),
94            ));
95        }
96
97        if field.data_type() != entries.data_type() {
98            return Err(ArrowError::InvalidArgumentError(format!(
99                "MapArray expected data type {} got {} for {:?}",
100                field.data_type(),
101                entries.data_type(),
102                field.name()
103            )));
104        }
105
106        if entries.columns().len() != 2 {
107            return Err(ArrowError::InvalidArgumentError(format!(
108                "MapArray entries must contain two children, got {}",
109                entries.columns().len()
110            )));
111        }
112
113        // The Arrow spec requires the "key" field to be non-nullable
114        // <https://github.com/apache/arrow/blob/98347d233f03bcf4d116d77f1769d498902b1fc8/format/Schema.fbs#L138>
115        if entries.fields()[0].is_nullable() {
116            return Err(ArrowError::InvalidArgumentError(
117                "MapArray keys field cannot be nullable".to_string(),
118            ));
119        }
120        // No need to verify if key contain nulls since `StructArray`
121        // already disallow nulls for non nullable fields
122
123        Ok(Self {
124            data_type: DataType::Map(field, ordered),
125            nulls,
126            entries,
127            value_offsets: offsets,
128        })
129    }
130
131    /// Create a new [`MapArray`] from the provided parts
132    ///
133    /// See [`MapBuilder`](crate::builder::MapBuilder) for a higher-level interface
134    /// to construct a [`MapArray`]
135    ///
136    /// # Panics
137    ///
138    /// Panics if [`Self::try_new`] returns an error
139    pub fn new(
140        field: FieldRef,
141        offsets: OffsetBuffer<i32>,
142        entries: StructArray,
143        nulls: Option<NullBuffer>,
144        ordered: bool,
145    ) -> Self {
146        Self::try_new(field, offsets, entries, nulls, ordered).unwrap()
147    }
148
149    /// Deconstruct this array into its constituent parts
150    pub fn into_parts(
151        self,
152    ) -> (
153        FieldRef,
154        OffsetBuffer<i32>,
155        StructArray,
156        Option<NullBuffer>,
157        bool,
158    ) {
159        let (f, ordered) = match self.data_type {
160            DataType::Map(f, ordered) => (f, ordered),
161            _ => unreachable!(),
162        };
163        (f, self.value_offsets, self.entries, self.nulls, ordered)
164    }
165
166    /// Returns a reference to the offsets of this map
167    ///
168    /// Unlike [`Self::value_offsets`] this returns the [`OffsetBuffer`]
169    /// allowing for zero-copy cloning
170    #[inline]
171    pub fn offsets(&self) -> &OffsetBuffer<i32> {
172        &self.value_offsets
173    }
174
175    /// Returns a reference to the keys of this map
176    pub fn keys(&self) -> &ArrayRef {
177        self.entries.column(0)
178    }
179
180    /// Returns a reference to the values of this map
181    pub fn values(&self) -> &ArrayRef {
182        self.entries.column(1)
183    }
184
185    /// Returns a reference to the [`StructArray`] entries of this map
186    pub fn entries(&self) -> &StructArray {
187        &self.entries
188    }
189
190    /// Returns a reference to the fields of the [`StructArray`] that backs this map.
191    pub fn entries_fields(&self) -> (&Field, &Field) {
192        (
193            self.entries.field(0).as_ref(),
194            self.entries.field(1).as_ref(),
195        )
196    }
197
198    /// Returns the data type of the map's keys.
199    pub fn key_type(&self) -> &DataType {
200        self.keys().data_type()
201    }
202
203    /// Returns the data type of the map's values.
204    pub fn value_type(&self) -> &DataType {
205        self.values().data_type()
206    }
207
208    /// Returns ith value of this map array.
209    ///
210    /// Note: This method does not check for nulls and the value is arbitrary
211    /// if [`is_null`](Self::is_null) returns true for the index.
212    ///
213    /// # Safety
214    /// Caller must ensure that the index is within the array bounds
215    pub unsafe fn value_unchecked(&self, i: usize) -> StructArray {
216        let end = *unsafe { self.value_offsets().get_unchecked(i + 1) };
217        let start = *unsafe { self.value_offsets().get_unchecked(i) };
218        self.entries
219            .slice(start.to_usize().unwrap(), (end - start).to_usize().unwrap())
220    }
221
222    /// Returns ith value of this map array.
223    ///
224    /// This is a [`StructArray`] containing two fields
225    ///
226    /// Note: This method does not check for nulls and the value is arbitrary
227    /// (but still well-defined) if [`is_null`](Self::is_null) returns true for the index.
228    ///
229    /// # Panics
230    /// Panics if index `i` is out of bounds
231    pub fn value(&self, i: usize) -> StructArray {
232        let end = self.value_offsets()[i + 1] as usize;
233        let start = self.value_offsets()[i] as usize;
234        self.entries.slice(start, end - start)
235    }
236
237    /// Returns the offset values in the offsets buffer
238    #[inline]
239    pub fn value_offsets(&self) -> &[i32] {
240        &self.value_offsets
241    }
242
243    /// Returns the length for value at index `i`.
244    #[inline]
245    pub fn value_length(&self, i: usize) -> i32 {
246        let offsets = self.value_offsets();
247        offsets[i + 1] - offsets[i]
248    }
249
250    /// Returns a zero-copy slice of this array with the indicated offset and length.
251    pub fn slice(&self, offset: usize, length: usize) -> Self {
252        Self {
253            data_type: self.data_type.clone(),
254            nulls: self.nulls.as_ref().map(|n| n.slice(offset, length)),
255            entries: self.entries.clone(),
256            value_offsets: self.value_offsets.slice(offset, length),
257        }
258    }
259
260    /// constructs a new iterator
261    pub fn iter(&self) -> MapArrayIter<'_> {
262        MapArrayIter::new(self)
263    }
264}
265
266impl From<ArrayData> for MapArray {
267    fn from(data: ArrayData) -> Self {
268        Self::try_new_from_array_data(data)
269            .expect("Expected infallible creation of MapArray from ArrayData failed")
270    }
271}
272
273impl From<MapArray> for ArrayData {
274    fn from(array: MapArray) -> Self {
275        let len = array.len();
276        let builder = ArrayDataBuilder::new(array.data_type)
277            .len(len)
278            .nulls(array.nulls)
279            .buffers(vec![array.value_offsets.into_inner().into_inner()])
280            .child_data(vec![array.entries.to_data()]);
281
282        unsafe { builder.build_unchecked() }
283    }
284}
285
286type Entries<Key, Value> = Vec<(Key, Value)>;
287
288impl MapArray {
289    fn try_new_from_array_data(data: ArrayData) -> Result<Self, ArrowError> {
290        let (data_type, len, nulls, offset, mut buffers, mut child_data) = data.into_parts();
291
292        if !matches!(data_type, DataType::Map(_, _)) {
293            return Err(ArrowError::InvalidArgumentError(format!(
294                "MapArray expected ArrayData with DataType::Map got {data_type}",
295            )));
296        }
297
298        if buffers.len() != 1 {
299            return Err(ArrowError::InvalidArgumentError(format!(
300                "MapArray data should contain a single buffer only (value offsets), had {}",
301                buffers.len(),
302            )));
303        }
304        let buffer = buffers.pop().expect("checked above");
305
306        if child_data.len() != 1 {
307            return Err(ArrowError::InvalidArgumentError(format!(
308                "MapArray should contain a single child array (values array), had {}",
309                child_data.len()
310            )));
311        }
312        let entries = child_data.pop().expect("checked above");
313
314        if let DataType::Struct(fields) = entries.data_type() {
315            if fields.len() != 2 {
316                return Err(ArrowError::InvalidArgumentError(format!(
317                    "MapArray should contain a struct array with 2 fields, have {} fields",
318                    fields.len()
319                )));
320            }
321        } else {
322            return Err(ArrowError::InvalidArgumentError(format!(
323                "MapArray should contain a struct array child, found {:?}",
324                entries.data_type()
325            )));
326        }
327        let entries = entries.into();
328
329        // SAFETY:
330        // ArrayData is valid, and verified type above
331        let value_offsets = unsafe { get_offsets_from_buffer(buffer, offset, len) };
332
333        Ok(Self {
334            data_type,
335            nulls,
336            entries,
337            value_offsets,
338        })
339    }
340
341    /// Creates map array from provided keys, values and entry_offsets.
342    pub fn new_from_strings<'a>(
343        keys: impl Iterator<Item = &'a str>,
344        values: &dyn Array,
345        entry_offsets: &[u32],
346    ) -> Result<Self, ArrowError> {
347        let entry_offsets_buffer = Buffer::from(entry_offsets.to_byte_slice());
348        let keys_data = StringArray::from_iter_values(keys);
349
350        let keys_field = Arc::new(Field::new(
351            Field::MAP_KEY_FIELD_DEFAULT_NAME,
352            DataType::Utf8,
353            false,
354        ));
355        let values_field = Arc::new(Field::new(
356            Field::MAP_VALUE_FIELD_DEFAULT_NAME,
357            values.data_type().clone(),
358            values.null_count() > 0,
359        ));
360
361        let entry_struct = StructArray::from(vec![
362            (keys_field, Arc::new(keys_data) as ArrayRef),
363            (values_field, make_array(values.to_data())),
364        ]);
365
366        let map_data_type = DataType::Map(
367            Arc::new(Field::new(
368                Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
369                entry_struct.data_type().clone(),
370                false,
371            )),
372            false,
373        );
374        let map_data = ArrayData::builder(map_data_type)
375            .len(entry_offsets.len() - 1)
376            .add_buffer(entry_offsets_buffer)
377            .add_child_data(entry_struct.into_data())
378            .build()?;
379
380        Ok(MapArray::from(map_data))
381    }
382
383    /// Helper to create [`MapArray`] from [`Vec`]s of entries so the code will look clean and straightforward
384    ///
385    /// the input is: `Vec<Option<Map>>` where each `Map` is `Vec<(Key, Option<Value>)>`
386    ///
387    /// Useful for tests, this should not be used for performance sensitive operations
388    ///
389    /// ```
390    /// use std::collections::HashMap;
391    /// # use arrow_array::{MapArray, Int32Array, StringArray};
392    ///
393    /// let map = vec![
394    ///    // {}
395    ///    Some(vec![]),
396    ///    // null
397    ///    None,
398    ///    // { "a": 1, "b": null, "cd": 4 }
399    ///    Some(vec![
400    ///        ("a", Some(1)),
401    ///        ("b", None),
402    ///        ("cd", Some(4)),
403    ///    ]),
404    ///    // { "e": 0 }
405    ///    Some(vec![("e", Some(0))]),
406    /// ];
407    /// let ordered = true;
408    ///
409    /// // created map: [{}, null, {"a": 1, "b": null, "cd": 4}, {"e": 0}]
410    /// let map_array = MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(map, ordered);
411    /// // Or you could fill the last 2 generics manually for the key array item and value array item
412    /// // let map_array = MapArray::from_vec_of_maps::<StringArray, Int32Array, &str, i32>(map, ordered);
413    ///```
414    #[allow(clippy::type_complexity)]
415    pub fn from_vec_of_maps<KeyArray, ValueArray, K, V>(
416        input: Vec<Option<Entries<K, Option<V>>>>,
417        ordered: bool,
418    ) -> Self
419    where
420        KeyArray: Array + 'static,
421        ValueArray: Array + 'static,
422        Vec<K>: Into<KeyArray>,
423        Vec<Option<V>>: Into<ValueArray>,
424    {
425        let offsets = OffsetBuffer::<i32>::from_lengths(
426            input.iter().map(|v| v.as_ref().map_or(0, |m| m.len())),
427        );
428        let nulls = NullBuffer::from_iter(input.iter().map(|v| v.is_some()));
429        let nulls = Some(nulls).filter(|b| b.null_count() > 0);
430
431        let (keys, values): (Vec<K>, Vec<Option<V>>) = input
432            .into_iter()
433            .flatten()
434            .flat_map(|m| m.into_iter())
435            .unzip();
436
437        let keys_array: ArrayRef = Arc::new(<Vec<K> as Into<KeyArray>>::into(keys));
438        let values_array: ArrayRef = Arc::new(<Vec<Option<V>> as Into<ValueArray>>::into(values));
439
440        let field_names = MapFieldNames::default();
441
442        let entries = StructArray::new(
443            Fields::from(vec![
444                Field::new(field_names.key, keys_array.data_type().clone(), false),
445                Field::new(
446                    field_names.value,
447                    values_array.data_type().clone(),
448                    values_array.is_nullable(),
449                ),
450            ]),
451            vec![keys_array, values_array],
452            None,
453        );
454
455        MapArray::new(
456            Arc::new(Field::new(
457                field_names.entry,
458                entries.data_type().clone(),
459                false,
460            )),
461            offsets,
462            entries,
463            nulls,
464            ordered,
465        )
466    }
467}
468
469/// SAFETY: Correctly implements the contract of Arrow Arrays
470unsafe impl Array for MapArray {
471    fn as_any(&self) -> &dyn Any {
472        self
473    }
474
475    fn to_data(&self) -> ArrayData {
476        self.clone().into_data()
477    }
478
479    fn into_data(self) -> ArrayData {
480        self.into()
481    }
482
483    fn data_type(&self) -> &DataType {
484        &self.data_type
485    }
486
487    fn slice(&self, offset: usize, length: usize) -> ArrayRef {
488        Arc::new(self.slice(offset, length))
489    }
490
491    fn len(&self) -> usize {
492        self.value_offsets.len() - 1
493    }
494
495    fn is_empty(&self) -> bool {
496        self.value_offsets.len() <= 1
497    }
498
499    fn shrink_to_fit(&mut self) {
500        if let Some(nulls) = &mut self.nulls {
501            nulls.shrink_to_fit();
502        }
503        self.entries.shrink_to_fit();
504        self.value_offsets.shrink_to_fit();
505    }
506
507    fn offset(&self) -> usize {
508        0
509    }
510
511    fn nulls(&self) -> Option<&NullBuffer> {
512        self.nulls.as_ref()
513    }
514
515    fn logical_null_count(&self) -> usize {
516        // More efficient that the default implementation
517        self.null_count()
518    }
519
520    fn get_buffer_memory_size(&self) -> usize {
521        let mut size = self.entries.get_buffer_memory_size();
522        size += self.value_offsets.inner().inner().capacity();
523        if let Some(n) = self.nulls.as_ref() {
524            size += n.buffer().capacity();
525        }
526        size
527    }
528
529    fn get_array_memory_size(&self) -> usize {
530        let mut size = std::mem::size_of::<Self>() + self.entries.get_array_memory_size();
531        size += self.value_offsets.inner().inner().capacity();
532        if let Some(n) = self.nulls.as_ref() {
533            size += n.buffer().capacity();
534        }
535        size
536    }
537
538    #[cfg(feature = "pool")]
539    fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
540        self.value_offsets.claim(pool);
541        self.entries.claim(pool);
542        if let Some(nulls) = &self.nulls {
543            nulls.claim(pool);
544        }
545    }
546}
547
548impl ArrayAccessor for &MapArray {
549    type Item = StructArray;
550
551    fn value(&self, index: usize) -> Self::Item {
552        MapArray::value(self, index)
553    }
554
555    unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
556        MapArray::value(self, index)
557    }
558}
559
560impl std::fmt::Debug for MapArray {
561    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
562        write!(f, "MapArray\n[\n")?;
563        print_long_array(self, f, |array, index, f| {
564            std::fmt::Debug::fmt(&array.value(index), f)
565        })?;
566        write!(f, "]")
567    }
568}
569
570impl From<MapArray> for ListArray {
571    fn from(value: MapArray) -> Self {
572        let field = match value.data_type() {
573            DataType::Map(field, _) => field,
574            _ => unreachable!("This should be a map type."),
575        };
576        let data_type = DataType::List(field.clone());
577        let builder = value.into_data().into_builder().data_type(data_type);
578        let array_data = unsafe { builder.build_unchecked() };
579
580        ListArray::from(array_data)
581    }
582}
583
584#[cfg(test)]
585mod tests {
586    use crate::builder::{Int32Builder, MapBuilder, StringBuilder};
587    use crate::cast::AsArray;
588    use crate::types::UInt32Type;
589    use crate::{Int32Array, UInt32Array};
590    use arrow_schema::Fields;
591
592    use super::*;
593
594    fn create_from_buffers() -> MapArray {
595        // Construct key and values
596        let keys_data = ArrayData::builder(DataType::Int32)
597            .len(8)
598            .add_buffer(Buffer::from([0, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
599            .build()
600            .unwrap();
601        let values_data = ArrayData::builder(DataType::UInt32)
602            .len(8)
603            .add_buffer(Buffer::from(
604                [0u32, 10, 20, 30, 40, 50, 60, 70].to_byte_slice(),
605            ))
606            .build()
607            .unwrap();
608
609        // Construct a buffer for value offsets, for the nested array:
610        //  [[0, 1, 2], [3, 4, 5], [6, 7]]
611        let entry_offsets = Buffer::from([0, 3, 6, 8].to_byte_slice());
612
613        let keys = Arc::new(Field::new(
614            Field::MAP_KEY_FIELD_DEFAULT_NAME,
615            DataType::Int32,
616            false,
617        ));
618        let values = Arc::new(Field::new(
619            Field::MAP_VALUE_FIELD_DEFAULT_NAME,
620            DataType::UInt32,
621            false,
622        ));
623        let entry_struct = StructArray::from(vec![
624            (keys, make_array(keys_data)),
625            (values, make_array(values_data)),
626        ]);
627
628        // Construct a map array from the above two
629        let map_data_type = DataType::Map(
630            Arc::new(Field::new(
631                Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
632                entry_struct.data_type().clone(),
633                false,
634            )),
635            false,
636        );
637        let map_data = ArrayData::builder(map_data_type)
638            .len(3)
639            .add_buffer(entry_offsets)
640            .add_child_data(entry_struct.into_data())
641            .build()
642            .unwrap();
643        MapArray::from(map_data)
644    }
645
646    #[test]
647    fn test_map_array() {
648        // Construct key and values
649        let key_data = ArrayData::builder(DataType::Int32)
650            .len(8)
651            .add_buffer(Buffer::from([0, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
652            .build()
653            .unwrap();
654        let value_data = ArrayData::builder(DataType::UInt32)
655            .len(8)
656            .add_buffer(Buffer::from(
657                [0u32, 10, 20, 0, 40, 0, 60, 70].to_byte_slice(),
658            ))
659            .null_bit_buffer(Some(Buffer::from(&[0b11010110])))
660            .build()
661            .unwrap();
662
663        // Construct a buffer for value offsets, for the nested array:
664        //  [[0, 1, 2], [3, 4, 5], [6, 7]]
665        let entry_offsets = Buffer::from([0, 3, 6, 8].to_byte_slice());
666
667        let keys_field = Arc::new(Field::new(
668            Field::MAP_KEY_FIELD_DEFAULT_NAME,
669            DataType::Int32,
670            false,
671        ));
672        let values_field = Arc::new(Field::new(
673            Field::MAP_VALUE_FIELD_DEFAULT_NAME,
674            DataType::UInt32,
675            true,
676        ));
677        let entry_struct = StructArray::from(vec![
678            (keys_field.clone(), make_array(key_data)),
679            (values_field.clone(), make_array(value_data.clone())),
680        ]);
681
682        // Construct a map array from the above two
683        let map_data_type = DataType::Map(
684            Arc::new(Field::new(
685                Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
686                entry_struct.data_type().clone(),
687                false,
688            )),
689            false,
690        );
691        let map_data = ArrayData::builder(map_data_type)
692            .len(3)
693            .add_buffer(entry_offsets)
694            .add_child_data(entry_struct.into_data())
695            .build()
696            .unwrap();
697        let map_array = MapArray::from(map_data);
698
699        assert_eq!(value_data, map_array.values().to_data());
700        assert_eq!(&DataType::UInt32, map_array.value_type());
701        assert_eq!(3, map_array.len());
702        assert_eq!(0, map_array.null_count());
703        assert_eq!(6, map_array.value_offsets()[2]);
704        assert_eq!(2, map_array.value_length(2));
705
706        let key_array = Arc::new(Int32Array::from(vec![0, 1, 2])) as ArrayRef;
707        let value_array =
708            Arc::new(UInt32Array::from(vec![None, Some(10u32), Some(20)])) as ArrayRef;
709        let struct_array = StructArray::from(vec![
710            (keys_field.clone(), key_array),
711            (values_field.clone(), value_array),
712        ]);
713        assert_eq!(
714            struct_array,
715            StructArray::from(map_array.value(0).into_data())
716        );
717        assert_eq!(
718            &struct_array,
719            unsafe { map_array.value_unchecked(0) }
720                .as_any()
721                .downcast_ref::<StructArray>()
722                .unwrap()
723        );
724        for i in 0..3 {
725            assert!(map_array.is_valid(i));
726            assert!(!map_array.is_null(i));
727        }
728
729        // Now test with a non-zero offset
730        let map_array = map_array.slice(1, 2);
731
732        assert_eq!(value_data, map_array.values().to_data());
733        assert_eq!(&DataType::UInt32, map_array.value_type());
734        assert_eq!(2, map_array.len());
735        assert_eq!(0, map_array.null_count());
736        assert_eq!(6, map_array.value_offsets()[1]);
737        assert_eq!(2, map_array.value_length(1));
738
739        let key_array = Arc::new(Int32Array::from(vec![3, 4, 5])) as ArrayRef;
740        let value_array = Arc::new(UInt32Array::from(vec![None, Some(40), None])) as ArrayRef;
741        let struct_array =
742            StructArray::from(vec![(keys_field, key_array), (values_field, value_array)]);
743        assert_eq!(
744            &struct_array,
745            map_array
746                .value(0)
747                .as_any()
748                .downcast_ref::<StructArray>()
749                .unwrap()
750        );
751        assert_eq!(
752            &struct_array,
753            unsafe { map_array.value_unchecked(0) }
754                .as_any()
755                .downcast_ref::<StructArray>()
756                .unwrap()
757        );
758    }
759
760    #[test]
761    #[ignore = "Test fails because slice of <list<struct>> is still buggy"]
762    fn test_map_array_slice() {
763        let map_array = create_from_buffers();
764
765        let sliced_array = map_array.slice(1, 2);
766        assert_eq!(2, sliced_array.len());
767        assert_eq!(1, sliced_array.offset());
768        let sliced_array_data = sliced_array.to_data();
769        for array_data in sliced_array_data.child_data() {
770            assert_eq!(array_data.offset(), 1);
771        }
772
773        // Check offset and length for each non-null value.
774        let sliced_map_array = sliced_array.as_any().downcast_ref::<MapArray>().unwrap();
775        assert_eq!(3, sliced_map_array.value_offsets()[0]);
776        assert_eq!(3, sliced_map_array.value_length(0));
777        assert_eq!(6, sliced_map_array.value_offsets()[1]);
778        assert_eq!(2, sliced_map_array.value_length(1));
779
780        // Construct key and values
781        let keys_data = ArrayData::builder(DataType::Int32)
782            .len(5)
783            .add_buffer(Buffer::from([3, 4, 5, 6, 7].to_byte_slice()))
784            .build()
785            .unwrap();
786        let values_data = ArrayData::builder(DataType::UInt32)
787            .len(5)
788            .add_buffer(Buffer::from([30u32, 40, 50, 60, 70].to_byte_slice()))
789            .build()
790            .unwrap();
791
792        // Construct a buffer for value offsets, for the nested array:
793        //  [[3, 4, 5], [6, 7]]
794        let entry_offsets = Buffer::from([0, 3, 5].to_byte_slice());
795
796        let keys = Arc::new(Field::new(
797            Field::MAP_KEY_FIELD_DEFAULT_NAME,
798            DataType::Int32,
799            false,
800        ));
801        let values = Arc::new(Field::new(
802            Field::MAP_VALUE_FIELD_DEFAULT_NAME,
803            DataType::UInt32,
804            false,
805        ));
806        let entry_struct = StructArray::from(vec![
807            (keys, make_array(keys_data)),
808            (values, make_array(values_data)),
809        ]);
810
811        // Construct a map array from the above two
812        let map_data_type = DataType::Map(
813            Arc::new(Field::new(
814                Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
815                entry_struct.data_type().clone(),
816                false,
817            )),
818            false,
819        );
820        let expected_map_data = ArrayData::builder(map_data_type)
821            .len(2)
822            .add_buffer(entry_offsets)
823            .add_child_data(entry_struct.into_data())
824            .build()
825            .unwrap();
826        let expected_map_array = MapArray::from(expected_map_data);
827
828        assert_eq!(&expected_map_array, sliced_map_array)
829    }
830
831    #[test]
832    #[should_panic(expected = "index out of bounds: the len is ")]
833    fn test_map_array_index_out_of_bound() {
834        let map_array = create_from_buffers();
835
836        map_array.value(map_array.len());
837    }
838
839    #[test]
840    #[should_panic(expected = "MapArray expected ArrayData with DataType::Map got Dictionary")]
841    fn test_from_array_data_validation() {
842        // A DictionaryArray has similar buffer layout to a MapArray
843        // but the meaning of the values differs
844        let struct_t = DataType::Struct(Fields::from(vec![
845            Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, true),
846            Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::UInt32, true),
847        ]));
848        let dict_t = DataType::Dictionary(Box::new(DataType::Int32), Box::new(struct_t));
849        let _ = MapArray::from(ArrayData::new_empty(&dict_t));
850    }
851
852    #[test]
853    fn test_new_from_strings() {
854        let keys = vec!["a", "b", "c", "d", "e", "f", "g", "h"];
855        let values_data = UInt32Array::from(vec![0u32, 10, 20, 30, 40, 50, 60, 70]);
856
857        // Construct a buffer for value offsets, for the nested array:
858        //  [[a, b, c], [d, e, f], [g, h]]
859        let entry_offsets = [0, 3, 6, 8];
860
861        let map_array =
862            MapArray::new_from_strings(keys.clone().into_iter(), &values_data, &entry_offsets)
863                .unwrap();
864
865        assert_eq!(
866            &values_data,
867            map_array.values().as_primitive::<UInt32Type>()
868        );
869        assert_eq!(&DataType::UInt32, map_array.value_type());
870        assert_eq!(3, map_array.len());
871        assert_eq!(0, map_array.null_count());
872        assert_eq!(6, map_array.value_offsets()[2]);
873        assert_eq!(2, map_array.value_length(2));
874
875        let key_array = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
876        let value_array = Arc::new(UInt32Array::from(vec![0u32, 10, 20])) as ArrayRef;
877        let keys_field = Arc::new(Field::new(
878            Field::MAP_KEY_FIELD_DEFAULT_NAME,
879            DataType::Utf8,
880            false,
881        ));
882        let values_field = Arc::new(Field::new(
883            Field::MAP_VALUE_FIELD_DEFAULT_NAME,
884            DataType::UInt32,
885            false,
886        ));
887        let struct_array =
888            StructArray::from(vec![(keys_field, key_array), (values_field, value_array)]);
889        assert_eq!(
890            struct_array,
891            StructArray::from(map_array.value(0).into_data())
892        );
893        assert_eq!(
894            &struct_array,
895            unsafe { map_array.value_unchecked(0) }
896                .as_any()
897                .downcast_ref::<StructArray>()
898                .unwrap()
899        );
900        for i in 0..3 {
901            assert!(map_array.is_valid(i));
902            assert!(!map_array.is_null(i));
903        }
904    }
905
906    #[test]
907    fn test_try_new() {
908        let offsets = OffsetBuffer::new(vec![0, 1, 4, 5].into());
909        let fields = Fields::from(vec![
910            Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
911            Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, false),
912        ]);
913        let columns = vec![
914            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _,
915            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _,
916        ];
917
918        let entries = StructArray::new(fields.clone(), columns, None);
919        let field = Arc::new(Field::new(
920            Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
921            DataType::Struct(fields),
922            false,
923        ));
924
925        MapArray::new(field.clone(), offsets.clone(), entries.clone(), None, false);
926
927        let nulls = NullBuffer::new_null(3);
928        MapArray::new(field.clone(), offsets, entries.clone(), Some(nulls), false);
929
930        let nulls = NullBuffer::new_null(3);
931        let offsets = OffsetBuffer::new(vec![0, 1, 2, 4, 5].into());
932        let err = MapArray::try_new(
933            field.clone(),
934            offsets.clone(),
935            entries.clone(),
936            Some(nulls),
937            false,
938        )
939        .unwrap_err();
940
941        assert_eq!(
942            err.to_string(),
943            "Invalid argument error: Incorrect length of null buffer for MapArray, expected 4 got 3"
944        );
945
946        let err = MapArray::try_new(field, offsets.clone(), entries.slice(0, 2), None, false)
947            .unwrap_err();
948
949        assert_eq!(
950            err.to_string(),
951            "Invalid argument error: Max offset of 5 exceeds length of entries 2"
952        );
953
954        let field = Arc::new(Field::new("element", DataType::Int64, false));
955        let err = MapArray::try_new(field, offsets.clone(), entries, None, false)
956            .unwrap_err()
957            .to_string();
958
959        assert!(
960            err.starts_with("Invalid argument error: MapArray expected data type Int64 got Struct"),
961            "{err}"
962        );
963
964        let fields = Fields::from(vec![
965            Field::new("a", DataType::Int32, false),
966            Field::new("b", DataType::Int32, false),
967            Field::new("c", DataType::Int32, false),
968        ]);
969        let columns = vec![
970            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _,
971            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _,
972            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _,
973        ];
974
975        let s = StructArray::new(fields.clone(), columns, None);
976        let field = Arc::new(Field::new(
977            Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
978            DataType::Struct(fields),
979            false,
980        ));
981        let err = MapArray::try_new(field, offsets, s, None, false).unwrap_err();
982
983        assert_eq!(
984            err.to_string(),
985            "Invalid argument error: MapArray entries must contain two children, got 3"
986        );
987    }
988
989    #[test]
990    fn test_try_new_nullable_keys_field() {
991        // https://github.com/apache/arrow-rs/issues/10268
992        let keys = Int32Array::from(vec![Some(1), None]);
993        let values = Int32Array::from(vec![None, Some(2)]);
994        let fields = Fields::from(vec![
995            Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, true),
996            Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, true),
997        ]);
998        let entries =
999            StructArray::new(fields.clone(), vec![Arc::new(keys), Arc::new(values)], None);
1000        let field = Arc::new(Field::new(
1001            Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
1002            DataType::Struct(fields),
1003            false,
1004        ));
1005
1006        let err = MapArray::try_new(field, OffsetBuffer::from_lengths([2]), entries, None, false)
1007            .unwrap_err();
1008        assert_eq!(
1009            err.to_string(),
1010            "Invalid argument error: MapArray keys field cannot be nullable"
1011        );
1012    }
1013
1014    #[test]
1015    fn test_from_vec_of_maps() {
1016        for ordered in [true, false] {
1017            let map = vec![
1018                Some(vec![]),
1019                None,
1020                Some(vec![("a", Some(1)), ("b", None), ("cd", Some(4))]),
1021                Some(vec![("e", Some(0))]),
1022            ];
1023
1024            let map_array =
1025                MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(map, ordered);
1026            assert_eq!(map_array.len(), 4);
1027
1028            let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::default());
1029
1030            // {}
1031            builder.append(true).unwrap();
1032
1033            // null
1034            builder.append_nulls(1).unwrap();
1035
1036            // {"a": 1, "b": null, "cd": 4}
1037            builder.keys().extend(["a", "b", "cd"].map(Some));
1038            builder.values().extend([Some(1), None, Some(4)]);
1039
1040            builder.append(true).unwrap();
1041
1042            // {"e": 0}
1043            builder.keys().append_value("e");
1044            builder.values().append_value(0);
1045
1046            builder.append(true).unwrap();
1047
1048            let (field, offsets, entries, null_buffer, _) = builder.finish().into_parts();
1049
1050            let expected_map = MapArray::new(field, offsets, entries, null_buffer, ordered);
1051
1052            assert_eq!(map_array, expected_map);
1053        }
1054    }
1055}
arrow_array/array/map_array.rs

arrow_array/array/
map_array.rs
