feat: add take for MapArray (apache#3925)

* feat: add take for MapArray * refactor: use into_builder
richox · Apr 7, 2023 · a2d4af5 · a2d4af5
1 parent bbc1469
commit a2d4af5
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 2 deletions.
diff --git a/arrow-array/src/array/map_array.rs b/arrow-array/src/array/map_array.rs
@@ -16,9 +16,9 @@
 // under the License.
 
 use crate::array::{get_offsets, print_long_array};
-use crate::{make_array, Array, ArrayRef, StringArray, StructArray};
-use arrow_buffer::buffer::OffsetBuffer;
+use crate::{make_array, Array, ArrayRef, ListArray, StringArray, StructArray};
 use arrow_buffer::{ArrowNativeType, Buffer, ToByteSlice};
+use arrow_buffer::buffer::{NullBuffer, OffsetBuffer};
 use arrow_data::ArrayData;
 use arrow_schema::{ArrowError, DataType, Field};
 use std::any::Any;
@@ -239,6 +239,20 @@ impl std::fmt::Debug for MapArray {
     }
 }
 
+impl From<MapArray> for ListArray {
+    fn from(value: MapArray) -> Self {
+        let field = match value.data_type() {
+            DataType::Map(field, _) => field,
+            _ => unreachable!("This should be a map type."),
+        };
+        let data_type = DataType::List(field.clone());
+        let builder = value.into_data().into_builder().data_type(data_type);
+        let array_data = unsafe { builder.build_unchecked() };
+
+        ListArray::from(array_data)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::cast::as_primitive_array;

diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs
@@ -158,6 +158,12 @@ where
                 *length as u32,
             )?))
         }
+        DataType::Map(_, _) => {
+            let list_arr = ListArray::from(values.as_map().clone());
+            let list_data = take_list::<_, Int32Type>(&list_arr, indices)?;
+            let builder = list_data.into_data().into_builder().data_type(values.data_type().clone());
+            Ok(Arc::new(MapArray::from(unsafe { builder.build_unchecked() })))
+        }
         DataType::Struct(fields) => {
             let struct_: &StructArray =
                 values.as_any().downcast_ref::<StructArray>().unwrap();
@@ -1935,6 +1941,30 @@ mod tests {
         take(&list_array, &index, None).unwrap();
     }
 
+    #[test]
+    fn test_take_map() {
+        let values = Int32Array::from(vec![1, 2, 3, 4]);
+        let array = MapArray::new_from_strings(
+            vec!["a", "b", "c", "a"].into_iter(),
+            &values,
+            &[0, 3, 4],
+        )
+        .unwrap();
+
+        let index = UInt32Array::from(vec![0]);
+
+        let result = take(&array, &index, None).unwrap();
+        let expected: ArrayRef = Arc::new(
+            MapArray::new_from_strings(
+                vec!["a", "b", "c"].into_iter(),
+                &values.slice(0, 3),
+                &[0, 3],
+            )
+            .unwrap(),
+        );
+        assert_eq!(&expected, &result);
+    }
+
     #[test]
     fn test_take_struct() {
         let array = create_test_struct(vec![