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

arrow_array/array/
map_array.rs
