Merge 99f77f5 into 41fedad

arrow/src/array/array_binary.rs

@@ -201,6 +201,14 @@ impl<OffsetSize: BinaryOffsetSizeTrait> GenericBinaryArray<OffsetSize> {
         let array_data = unsafe { array_data.build_unchecked() };
         Self::from(array_data)
     }
+
+    /// Returns an iterator that returns the values of `array.value(i)` for an iterator with each element `i`
+    pub fn take_iter<'a>(
+        &'a self,
+        indexes: impl Iterator<Item = Option<usize>> + 'a,
+    ) -> impl Iterator<Item = Option<&[u8]>> + 'a {
+        indexes.map(|opt_index| opt_index.map(|index| self.value(index)))
+    }
 }
 
 impl<'a, T: BinaryOffsetSizeTrait> GenericBinaryArray<T> {

arrow/src/array/array_boolean.rs

@@ -122,6 +122,14 @@ impl BooleanArray {
         // `i < self.len()
         unsafe { self.value_unchecked(i) }
     }
+
+    /// Returns an iterator that returns the values of `array.value(i)` for an iterator with each element `i`
+    pub fn take_iter<'a>(
+        &'a self,
+        indexes: impl Iterator<Item = Option<usize>> + 'a,
+    ) -> impl Iterator<Item = Option<bool>> + 'a {
+        indexes.map(|opt_index| opt_index.map(|index| self.value(index)))
+    }
 }
 
 impl Array for BooleanArray {

arrow/src/array/array_dictionary.rs

@@ -105,10 +105,17 @@ impl<'a, K: ArrowPrimitiveType> DictionaryArray<K> {
         self.keys.is_empty()
     }
 
-    // Currently exists for compatibility purposes with Arrow IPC.
+    /// Currently exists for compatibility purposes with Arrow IPC.
     pub fn is_ordered(&self) -> bool {
         self.is_ordered
     }
+
+    /// Return an iterator over the keys (indexes into the dictionary)
+    pub fn keys_iter(&self) -> impl Iterator<Item = Option<usize>> + '_ {
+        self.keys
+            .iter()
+            .map(|key| key.map(|k| k.to_usize().expect("Dictionary index not usize")))
+    }
 }
 
 /// Constructs a `DictionaryArray` from an array data reference.
@@ -256,6 +263,7 @@ impl<T: ArrowPrimitiveType> fmt::Debug for DictionaryArray<T> {
 mod tests {
     use super::*;
 
+    use crate::array::Int8Array;
     use crate::{
         array::Int16Array,
         datatypes::{Int32Type, Int8Type, UInt32Type, UInt8Type},
@@ -422,4 +430,71 @@ mod tests {
             .validate_full()
             .expect("All null array has valid array data");
     }
+
+    #[test]
+    fn test_dictionary_iter() {
+        // Construct a value array
+        let value_data = ArrayData::builder(DataType::Int8)
+            .len(8)
+            .add_buffer(Buffer::from(
+                &[10_i8, 11, 12, 13, 14, 15, 16, 17].to_byte_slice(),
+            ))
+            .build()
+            .unwrap();
+
+        // Construct a buffer for value offsets, for the nested array:
+        let keys = Buffer::from(&[2_i16, 3, 4].to_byte_slice());
+
+        // Construct a dictionary array from the above two
+        let key_type = DataType::Int16;
+        let value_type = DataType::Int8;
+        let dict_data_type =
+            DataType::Dictionary(Box::new(key_type), Box::new(value_type));
+        let dict_data = ArrayData::builder(dict_data_type)
+            .len(3)
+            .add_buffer(keys)
+            .add_child_data(value_data)
+            .build()
+            .unwrap();
+        let dict_array = Int16DictionaryArray::from(dict_data);
+
+        let mut key_iter = dict_array.keys_iter();
+        assert_eq!(2, key_iter.next().unwrap().unwrap());
+        assert_eq!(3, key_iter.next().unwrap().unwrap());
+        assert_eq!(4, key_iter.next().unwrap().unwrap());
+        assert!(key_iter.next().is_none());
+
+        let mut iter = dict_array
+            .values()
+            .as_any()
+            .downcast_ref::<Int8Array>()
+            .unwrap()
+            .take_iter(dict_array.keys_iter());
+
+        assert_eq!(12, iter.next().unwrap().unwrap());
+        assert_eq!(13, iter.next().unwrap().unwrap());
+        assert_eq!(14, iter.next().unwrap().unwrap());
+        assert!(iter.next().is_none());
+    }
+
+    #[test]
+    fn test_dictionary_iter_with_null() {
+        let test = vec![Some("a"), None, Some("b"), None, None, Some("a")];
+        let array: DictionaryArray<Int32Type> = test.into_iter().collect();
+
+        let mut iter = array
+            .values()
+            .as_any()
+            .downcast_ref::<StringArray>()
+            .unwrap()
+            .take_iter(array.keys_iter());
+
+        assert_eq!("a", iter.next().unwrap().unwrap());
+        assert!(iter.next().unwrap().is_none());
+        assert_eq!("b", iter.next().unwrap().unwrap());
+        assert!(iter.next().unwrap().is_none());
+        assert!(iter.next().unwrap().is_none());
+        assert_eq!("a", iter.next().unwrap().unwrap());
+        assert!(iter.next().is_none());
+    }
 }

arrow/src/array/array_primitive.rs

@@ -155,6 +155,14 @@ impl<T: ArrowPrimitiveType> PrimitiveArray<T> {
         };
         PrimitiveArray::from(data)
     }
+
+    /// Returns an iterator that returns the values of `array.value(i)` for an iterator with each element `i`
+    pub fn take_iter<'a>(
+        &'a self,
+        indexes: impl Iterator<Item = Option<usize>> + 'a,
+    ) -> impl Iterator<Item = Option<T::Native>> + 'a {
+        indexes.map(|opt_index| opt_index.map(|index| self.value(index)))
+    }
 }
 
 impl<T: ArrowPrimitiveType> Array for PrimitiveArray<T> {

arrow/src/array/array_string.rs

@@ -172,6 +172,14 @@ impl<OffsetSize: StringOffsetSizeTrait> GenericStringArray<OffsetSize> {
         let array_data = unsafe { array_data.build_unchecked() };
         Self::from(array_data)
     }
+
+    /// Returns an iterator that returns the values of `array.value(i)` for an iterator with each element `i`
+    pub fn take_iter<'a>(
+        &'a self,
+        indexes: impl Iterator<Item = Option<usize>> + 'a,
+    ) -> impl Iterator<Item = Option<&str>> + 'a {
+        indexes.map(|opt_index| opt_index.map(|index| self.value(index)))
+    }
 }
 
 impl<'a, Ptr, OffsetSize: StringOffsetSizeTrait> FromIterator<&'a Option<Ptr>>