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

arrow_array/array/
map_array.rs
