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

arrow_array/array/
map_array.rs
