1use crate::array::{get_offsets, print_long_array};
19use crate::iterator::MapArrayIter;
20use crate::{make_array, Array, ArrayAccessor, ArrayRef, ListArray, StringArray, StructArray};
21use arrow_buffer::{ArrowNativeType, Buffer, NullBuffer, OffsetBuffer, ToByteSlice};
22use arrow_data::{ArrayData, ArrayDataBuilder};
23use arrow_schema::{ArrowError, DataType, Field, FieldRef};
24use std::any::Any;
25use std::sync::Arc;
26
27#[derive(Clone)]
36pub struct MapArray {
37 data_type: DataType,
38 nulls: Option<NullBuffer>,
39 entries: StructArray,
41 value_offsets: OffsetBuffer<i32>,
43}
44
45impl MapArray {
46 pub fn try_new(
62 field: FieldRef,
63 offsets: OffsetBuffer<i32>,
64 entries: StructArray,
65 nulls: Option<NullBuffer>,
66 ordered: bool,
67 ) -> Result<Self, ArrowError> {
68 let len = offsets.len() - 1; let end_offset = offsets.last().unwrap().as_usize();
70 if end_offset > entries.len() {
73 return Err(ArrowError::InvalidArgumentError(format!(
74 "Max offset of {end_offset} exceeds length of entries {}",
75 entries.len()
76 )));
77 }
78
79 if let Some(n) = nulls.as_ref() {
80 if n.len() != len {
81 return Err(ArrowError::InvalidArgumentError(format!(
82 "Incorrect length of null buffer for MapArray, expected {len} got {}",
83 n.len(),
84 )));
85 }
86 }
87 if field.is_nullable() || entries.null_count() != 0 {
88 return Err(ArrowError::InvalidArgumentError(
89 "MapArray entries cannot contain nulls".to_string(),
90 ));
91 }
92
93 if field.data_type() != entries.data_type() {
94 return Err(ArrowError::InvalidArgumentError(format!(
95 "MapArray expected data type {} got {} for {:?}",
96 field.data_type(),
97 entries.data_type(),
98 field.name()
99 )));
100 }
101
102 if entries.columns().len() != 2 {
103 return Err(ArrowError::InvalidArgumentError(format!(
104 "MapArray entries must contain two children, got {}",
105 entries.columns().len()
106 )));
107 }
108
109 Ok(Self {
110 data_type: DataType::Map(field, ordered),
111 nulls,
112 entries,
113 value_offsets: offsets,
114 })
115 }
116
117 pub fn new(
126 field: FieldRef,
127 offsets: OffsetBuffer<i32>,
128 entries: StructArray,
129 nulls: Option<NullBuffer>,
130 ordered: bool,
131 ) -> Self {
132 Self::try_new(field, offsets, entries, nulls, ordered).unwrap()
133 }
134
135 pub fn into_parts(
137 self,
138 ) -> (
139 FieldRef,
140 OffsetBuffer<i32>,
141 StructArray,
142 Option<NullBuffer>,
143 bool,
144 ) {
145 let (f, ordered) = match self.data_type {
146 DataType::Map(f, ordered) => (f, ordered),
147 _ => unreachable!(),
148 };
149 (f, self.value_offsets, self.entries, self.nulls, ordered)
150 }
151
152 #[inline]
157 pub fn offsets(&self) -> &OffsetBuffer<i32> {
158 &self.value_offsets
159 }
160
161 pub fn keys(&self) -> &ArrayRef {
163 self.entries.column(0)
164 }
165
166 pub fn values(&self) -> &ArrayRef {
168 self.entries.column(1)
169 }
170
171 pub fn entries(&self) -> &StructArray {
173 &self.entries
174 }
175
176 pub fn key_type(&self) -> &DataType {
178 self.keys().data_type()
179 }
180
181 pub fn value_type(&self) -> &DataType {
183 self.values().data_type()
184 }
185
186 pub unsafe fn value_unchecked(&self, i: usize) -> StructArray {
191 let end = *self.value_offsets().get_unchecked(i + 1);
192 let start = *self.value_offsets().get_unchecked(i);
193 self.entries
194 .slice(start.to_usize().unwrap(), (end - start).to_usize().unwrap())
195 }
196
197 pub fn value(&self, i: usize) -> StructArray {
201 let end = self.value_offsets()[i + 1] as usize;
202 let start = self.value_offsets()[i] as usize;
203 self.entries.slice(start, end - start)
204 }
205
206 #[inline]
208 pub fn value_offsets(&self) -> &[i32] {
209 &self.value_offsets
210 }
211
212 #[inline]
214 pub fn value_length(&self, i: usize) -> i32 {
215 let offsets = self.value_offsets();
216 offsets[i + 1] - offsets[i]
217 }
218
219 pub fn slice(&self, offset: usize, length: usize) -> Self {
221 Self {
222 data_type: self.data_type.clone(),
223 nulls: self.nulls.as_ref().map(|n| n.slice(offset, length)),
224 entries: self.entries.clone(),
225 value_offsets: self.value_offsets.slice(offset, length),
226 }
227 }
228
229 pub fn iter(&self) -> MapArrayIter<'_> {
231 MapArrayIter::new(self)
232 }
233}
234
235impl From<ArrayData> for MapArray {
236 fn from(data: ArrayData) -> Self {
237 Self::try_new_from_array_data(data)
238 .expect("Expected infallible creation of MapArray from ArrayData failed")
239 }
240}
241
242impl From<MapArray> for ArrayData {
243 fn from(array: MapArray) -> Self {
244 let len = array.len();
245 let builder = ArrayDataBuilder::new(array.data_type)
246 .len(len)
247 .nulls(array.nulls)
248 .buffers(vec![array.value_offsets.into_inner().into_inner()])
249 .child_data(vec![array.entries.to_data()]);
250
251 unsafe { builder.build_unchecked() }
252 }
253}
254
255impl MapArray {
256 fn try_new_from_array_data(data: ArrayData) -> Result<Self, ArrowError> {
257 if !matches!(data.data_type(), DataType::Map(_, _)) {
258 return Err(ArrowError::InvalidArgumentError(format!(
259 "MapArray expected ArrayData with DataType::Map got {}",
260 data.data_type()
261 )));
262 }
263
264 if data.buffers().len() != 1 {
265 return Err(ArrowError::InvalidArgumentError(format!(
266 "MapArray data should contain a single buffer only (value offsets), had {}",
267 data.len()
268 )));
269 }
270
271 if data.child_data().len() != 1 {
272 return Err(ArrowError::InvalidArgumentError(format!(
273 "MapArray should contain a single child array (values array), had {}",
274 data.child_data().len()
275 )));
276 }
277
278 let entries = data.child_data()[0].clone();
279
280 if let DataType::Struct(fields) = entries.data_type() {
281 if fields.len() != 2 {
282 return Err(ArrowError::InvalidArgumentError(format!(
283 "MapArray should contain a struct array with 2 fields, have {} fields",
284 fields.len()
285 )));
286 }
287 } else {
288 return Err(ArrowError::InvalidArgumentError(format!(
289 "MapArray should contain a struct array child, found {:?}",
290 entries.data_type()
291 )));
292 }
293 let entries = entries.into();
294
295 let value_offsets = unsafe { get_offsets(&data) };
298
299 Ok(Self {
300 data_type: data.data_type().clone(),
301 nulls: data.nulls().cloned(),
302 entries,
303 value_offsets,
304 })
305 }
306
307 pub fn new_from_strings<'a>(
309 keys: impl Iterator<Item = &'a str>,
310 values: &dyn Array,
311 entry_offsets: &[u32],
312 ) -> Result<Self, ArrowError> {
313 let entry_offsets_buffer = Buffer::from(entry_offsets.to_byte_slice());
314 let keys_data = StringArray::from_iter_values(keys);
315
316 let keys_field = Arc::new(Field::new("keys", DataType::Utf8, false));
317 let values_field = Arc::new(Field::new(
318 "values",
319 values.data_type().clone(),
320 values.null_count() > 0,
321 ));
322
323 let entry_struct = StructArray::from(vec![
324 (keys_field, Arc::new(keys_data) as ArrayRef),
325 (values_field, make_array(values.to_data())),
326 ]);
327
328 let map_data_type = DataType::Map(
329 Arc::new(Field::new(
330 "entries",
331 entry_struct.data_type().clone(),
332 false,
333 )),
334 false,
335 );
336 let map_data = ArrayData::builder(map_data_type)
337 .len(entry_offsets.len() - 1)
338 .add_buffer(entry_offsets_buffer)
339 .add_child_data(entry_struct.into_data())
340 .build()?;
341
342 Ok(MapArray::from(map_data))
343 }
344}
345
346impl Array for MapArray {
347 fn as_any(&self) -> &dyn Any {
348 self
349 }
350
351 fn to_data(&self) -> ArrayData {
352 self.clone().into_data()
353 }
354
355 fn into_data(self) -> ArrayData {
356 self.into()
357 }
358
359 fn data_type(&self) -> &DataType {
360 &self.data_type
361 }
362
363 fn slice(&self, offset: usize, length: usize) -> ArrayRef {
364 Arc::new(self.slice(offset, length))
365 }
366
367 fn len(&self) -> usize {
368 self.value_offsets.len() - 1
369 }
370
371 fn is_empty(&self) -> bool {
372 self.value_offsets.len() <= 1
373 }
374
375 fn offset(&self) -> usize {
376 0
377 }
378
379 fn nulls(&self) -> Option<&NullBuffer> {
380 self.nulls.as_ref()
381 }
382
383 fn logical_null_count(&self) -> usize {
384 self.null_count()
386 }
387
388 fn get_buffer_memory_size(&self) -> usize {
389 let mut size = self.entries.get_buffer_memory_size();
390 size += self.value_offsets.inner().inner().capacity();
391 if let Some(n) = self.nulls.as_ref() {
392 size += n.buffer().capacity();
393 }
394 size
395 }
396
397 fn get_array_memory_size(&self) -> usize {
398 let mut size = std::mem::size_of::<Self>() + self.entries.get_array_memory_size();
399 size += self.value_offsets.inner().inner().capacity();
400 if let Some(n) = self.nulls.as_ref() {
401 size += n.buffer().capacity();
402 }
403 size
404 }
405}
406
407impl ArrayAccessor for &MapArray {
408 type Item = StructArray;
409
410 fn value(&self, index: usize) -> Self::Item {
411 MapArray::value(self, index)
412 }
413
414 unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
415 MapArray::value(self, index)
416 }
417}
418
419impl std::fmt::Debug for MapArray {
420 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
421 write!(f, "MapArray\n[\n")?;
422 print_long_array(self, f, |array, index, f| {
423 std::fmt::Debug::fmt(&array.value(index), f)
424 })?;
425 write!(f, "]")
426 }
427}
428
429impl From<MapArray> for ListArray {
430 fn from(value: MapArray) -> Self {
431 let field = match value.data_type() {
432 DataType::Map(field, _) => field,
433 _ => unreachable!("This should be a map type."),
434 };
435 let data_type = DataType::List(field.clone());
436 let builder = value.into_data().into_builder().data_type(data_type);
437 let array_data = unsafe { builder.build_unchecked() };
438
439 ListArray::from(array_data)
440 }
441}
442
443#[cfg(test)]
444mod tests {
445 use crate::cast::AsArray;
446 use crate::types::UInt32Type;
447 use crate::{Int32Array, UInt32Array};
448 use arrow_schema::Fields;
449
450 use super::*;
451
452 fn create_from_buffers() -> MapArray {
453 let keys_data = ArrayData::builder(DataType::Int32)
455 .len(8)
456 .add_buffer(Buffer::from([0, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
457 .build()
458 .unwrap();
459 let values_data = ArrayData::builder(DataType::UInt32)
460 .len(8)
461 .add_buffer(Buffer::from(
462 [0u32, 10, 20, 30, 40, 50, 60, 70].to_byte_slice(),
463 ))
464 .build()
465 .unwrap();
466
467 let entry_offsets = Buffer::from([0, 3, 6, 8].to_byte_slice());
470
471 let keys = Arc::new(Field::new("keys", DataType::Int32, false));
472 let values = Arc::new(Field::new("values", DataType::UInt32, false));
473 let entry_struct = StructArray::from(vec![
474 (keys, make_array(keys_data)),
475 (values, make_array(values_data)),
476 ]);
477
478 let map_data_type = DataType::Map(
480 Arc::new(Field::new(
481 "entries",
482 entry_struct.data_type().clone(),
483 false,
484 )),
485 false,
486 );
487 let map_data = ArrayData::builder(map_data_type)
488 .len(3)
489 .add_buffer(entry_offsets)
490 .add_child_data(entry_struct.into_data())
491 .build()
492 .unwrap();
493 MapArray::from(map_data)
494 }
495
496 #[test]
497 fn test_map_array() {
498 let key_data = ArrayData::builder(DataType::Int32)
500 .len(8)
501 .add_buffer(Buffer::from([0, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
502 .build()
503 .unwrap();
504 let value_data = ArrayData::builder(DataType::UInt32)
505 .len(8)
506 .add_buffer(Buffer::from(
507 [0u32, 10, 20, 0, 40, 0, 60, 70].to_byte_slice(),
508 ))
509 .null_bit_buffer(Some(Buffer::from(&[0b11010110])))
510 .build()
511 .unwrap();
512
513 let entry_offsets = Buffer::from([0, 3, 6, 8].to_byte_slice());
516
517 let keys_field = Arc::new(Field::new("keys", DataType::Int32, false));
518 let values_field = Arc::new(Field::new("values", DataType::UInt32, true));
519 let entry_struct = StructArray::from(vec![
520 (keys_field.clone(), make_array(key_data)),
521 (values_field.clone(), make_array(value_data.clone())),
522 ]);
523
524 let map_data_type = DataType::Map(
526 Arc::new(Field::new(
527 "entries",
528 entry_struct.data_type().clone(),
529 false,
530 )),
531 false,
532 );
533 let map_data = ArrayData::builder(map_data_type)
534 .len(3)
535 .add_buffer(entry_offsets)
536 .add_child_data(entry_struct.into_data())
537 .build()
538 .unwrap();
539 let map_array = MapArray::from(map_data);
540
541 assert_eq!(value_data, map_array.values().to_data());
542 assert_eq!(&DataType::UInt32, map_array.value_type());
543 assert_eq!(3, map_array.len());
544 assert_eq!(0, map_array.null_count());
545 assert_eq!(6, map_array.value_offsets()[2]);
546 assert_eq!(2, map_array.value_length(2));
547
548 let key_array = Arc::new(Int32Array::from(vec![0, 1, 2])) as ArrayRef;
549 let value_array =
550 Arc::new(UInt32Array::from(vec![None, Some(10u32), Some(20)])) as ArrayRef;
551 let struct_array = StructArray::from(vec![
552 (keys_field.clone(), key_array),
553 (values_field.clone(), value_array),
554 ]);
555 assert_eq!(
556 struct_array,
557 StructArray::from(map_array.value(0).into_data())
558 );
559 assert_eq!(
560 &struct_array,
561 unsafe { map_array.value_unchecked(0) }
562 .as_any()
563 .downcast_ref::<StructArray>()
564 .unwrap()
565 );
566 for i in 0..3 {
567 assert!(map_array.is_valid(i));
568 assert!(!map_array.is_null(i));
569 }
570
571 let map_array = map_array.slice(1, 2);
573
574 assert_eq!(value_data, map_array.values().to_data());
575 assert_eq!(&DataType::UInt32, map_array.value_type());
576 assert_eq!(2, map_array.len());
577 assert_eq!(0, map_array.null_count());
578 assert_eq!(6, map_array.value_offsets()[1]);
579 assert_eq!(2, map_array.value_length(1));
580
581 let key_array = Arc::new(Int32Array::from(vec![3, 4, 5])) as ArrayRef;
582 let value_array = Arc::new(UInt32Array::from(vec![None, Some(40), None])) as ArrayRef;
583 let struct_array =
584 StructArray::from(vec![(keys_field, key_array), (values_field, value_array)]);
585 assert_eq!(
586 &struct_array,
587 map_array
588 .value(0)
589 .as_any()
590 .downcast_ref::<StructArray>()
591 .unwrap()
592 );
593 assert_eq!(
594 &struct_array,
595 unsafe { map_array.value_unchecked(0) }
596 .as_any()
597 .downcast_ref::<StructArray>()
598 .unwrap()
599 );
600 }
601
602 #[test]
603 #[ignore = "Test fails because slice of <list<struct>> is still buggy"]
604 fn test_map_array_slice() {
605 let map_array = create_from_buffers();
606
607 let sliced_array = map_array.slice(1, 2);
608 assert_eq!(2, sliced_array.len());
609 assert_eq!(1, sliced_array.offset());
610 let sliced_array_data = sliced_array.to_data();
611 for array_data in sliced_array_data.child_data() {
612 assert_eq!(array_data.offset(), 1);
613 }
614
615 let sliced_map_array = sliced_array.as_any().downcast_ref::<MapArray>().unwrap();
617 assert_eq!(3, sliced_map_array.value_offsets()[0]);
618 assert_eq!(3, sliced_map_array.value_length(0));
619 assert_eq!(6, sliced_map_array.value_offsets()[1]);
620 assert_eq!(2, sliced_map_array.value_length(1));
621
622 let keys_data = ArrayData::builder(DataType::Int32)
624 .len(5)
625 .add_buffer(Buffer::from([3, 4, 5, 6, 7].to_byte_slice()))
626 .build()
627 .unwrap();
628 let values_data = ArrayData::builder(DataType::UInt32)
629 .len(5)
630 .add_buffer(Buffer::from([30u32, 40, 50, 60, 70].to_byte_slice()))
631 .build()
632 .unwrap();
633
634 let entry_offsets = Buffer::from([0, 3, 5].to_byte_slice());
637
638 let keys = Arc::new(Field::new("keys", DataType::Int32, false));
639 let values = Arc::new(Field::new("values", DataType::UInt32, false));
640 let entry_struct = StructArray::from(vec![
641 (keys, make_array(keys_data)),
642 (values, make_array(values_data)),
643 ]);
644
645 let map_data_type = DataType::Map(
647 Arc::new(Field::new(
648 "entries",
649 entry_struct.data_type().clone(),
650 false,
651 )),
652 false,
653 );
654 let expected_map_data = ArrayData::builder(map_data_type)
655 .len(2)
656 .add_buffer(entry_offsets)
657 .add_child_data(entry_struct.into_data())
658 .build()
659 .unwrap();
660 let expected_map_array = MapArray::from(expected_map_data);
661
662 assert_eq!(&expected_map_array, sliced_map_array)
663 }
664
665 #[test]
666 #[should_panic(expected = "index out of bounds: the len is ")]
667 fn test_map_array_index_out_of_bound() {
668 let map_array = create_from_buffers();
669
670 map_array.value(map_array.len());
671 }
672
673 #[test]
674 #[should_panic(expected = "MapArray expected ArrayData with DataType::Map got Dictionary")]
675 fn test_from_array_data_validation() {
676 let struct_t = DataType::Struct(Fields::from(vec![
679 Field::new("keys", DataType::Int32, true),
680 Field::new("values", DataType::UInt32, true),
681 ]));
682 let dict_t = DataType::Dictionary(Box::new(DataType::Int32), Box::new(struct_t));
683 let _ = MapArray::from(ArrayData::new_empty(&dict_t));
684 }
685
686 #[test]
687 fn test_new_from_strings() {
688 let keys = vec!["a", "b", "c", "d", "e", "f", "g", "h"];
689 let values_data = UInt32Array::from(vec![0u32, 10, 20, 30, 40, 50, 60, 70]);
690
691 let entry_offsets = [0, 3, 6, 8];
694
695 let map_array =
696 MapArray::new_from_strings(keys.clone().into_iter(), &values_data, &entry_offsets)
697 .unwrap();
698
699 assert_eq!(
700 &values_data,
701 map_array.values().as_primitive::<UInt32Type>()
702 );
703 assert_eq!(&DataType::UInt32, map_array.value_type());
704 assert_eq!(3, map_array.len());
705 assert_eq!(0, map_array.null_count());
706 assert_eq!(6, map_array.value_offsets()[2]);
707 assert_eq!(2, map_array.value_length(2));
708
709 let key_array = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
710 let value_array = Arc::new(UInt32Array::from(vec![0u32, 10, 20])) as ArrayRef;
711 let keys_field = Arc::new(Field::new("keys", DataType::Utf8, false));
712 let values_field = Arc::new(Field::new("values", DataType::UInt32, false));
713 let struct_array =
714 StructArray::from(vec![(keys_field, key_array), (values_field, value_array)]);
715 assert_eq!(
716 struct_array,
717 StructArray::from(map_array.value(0).into_data())
718 );
719 assert_eq!(
720 &struct_array,
721 unsafe { map_array.value_unchecked(0) }
722 .as_any()
723 .downcast_ref::<StructArray>()
724 .unwrap()
725 );
726 for i in 0..3 {
727 assert!(map_array.is_valid(i));
728 assert!(!map_array.is_null(i));
729 }
730 }
731
732 #[test]
733 fn test_try_new() {
734 let offsets = OffsetBuffer::new(vec![0, 1, 4, 5].into());
735 let fields = Fields::from(vec![
736 Field::new("key", DataType::Int32, false),
737 Field::new("values", DataType::Int32, false),
738 ]);
739 let columns = vec![
740 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _,
741 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _,
742 ];
743
744 let entries = StructArray::new(fields.clone(), columns, None);
745 let field = Arc::new(Field::new("entries", DataType::Struct(fields), false));
746
747 MapArray::new(field.clone(), offsets.clone(), entries.clone(), None, false);
748
749 let nulls = NullBuffer::new_null(3);
750 MapArray::new(field.clone(), offsets, entries.clone(), Some(nulls), false);
751
752 let nulls = NullBuffer::new_null(3);
753 let offsets = OffsetBuffer::new(vec![0, 1, 2, 4, 5].into());
754 let err = MapArray::try_new(
755 field.clone(),
756 offsets.clone(),
757 entries.clone(),
758 Some(nulls),
759 false,
760 )
761 .unwrap_err();
762
763 assert_eq!(
764 err.to_string(),
765 "Invalid argument error: Incorrect length of null buffer for MapArray, expected 4 got 3"
766 );
767
768 let err = MapArray::try_new(field, offsets.clone(), entries.slice(0, 2), None, false)
769 .unwrap_err();
770
771 assert_eq!(
772 err.to_string(),
773 "Invalid argument error: Max offset of 5 exceeds length of entries 2"
774 );
775
776 let field = Arc::new(Field::new("element", DataType::Int64, false));
777 let err = MapArray::try_new(field, offsets.clone(), entries, None, false)
778 .unwrap_err()
779 .to_string();
780
781 assert!(
782 err.starts_with("Invalid argument error: MapArray expected data type Int64 got Struct"),
783 "{err}"
784 );
785
786 let fields = Fields::from(vec![
787 Field::new("a", DataType::Int32, false),
788 Field::new("b", DataType::Int32, false),
789 Field::new("c", DataType::Int32, false),
790 ]);
791 let columns = vec![
792 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _,
793 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _,
794 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _,
795 ];
796
797 let s = StructArray::new(fields.clone(), columns, None);
798 let field = Arc::new(Field::new("entries", DataType::Struct(fields), false));
799 let err = MapArray::try_new(field, offsets, s, None, false).unwrap_err();
800
801 assert_eq!(
802 err.to_string(),
803 "Invalid argument error: MapArray entries must contain two children, got 3"
804 );
805 }
806}