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

arrow_array/array/
map_array.rs
