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

arrow_array/array/
map_array.rs
