Use DictionaryArray's iterator (#1330)

apache · Feb 28, 2022 · f58a8d9 · f58a8d9
1 parent 5737449
commit f58a8d9
Showing 1 changed file with 16 additions and 17 deletions.
diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs
@@ -2213,27 +2213,26 @@ macro_rules! compare_dict_op {
                     .to_string(),
             ));
         }
-        let left_values = $left.values().as_any().downcast_ref::<$value_ty>().unwrap();
-        let right_values = $right
+
+        let left_iter = $left
             .values()
             .as_any()
             .downcast_ref::<$value_ty>()
-            .unwrap();
+            .unwrap()
+            .take_iter($left.keys_iter());
 
-        let result = $left
-            .keys()
-            .iter()
-            .zip($right.keys().iter())
-            .map(|(left_key, right_key)| {
-                if let (Some(left_k), Some(right_k)) = (left_key, right_key) {
-                    let left_key = left_k.to_usize().expect("Dictionary index not usize");
-                    let right_key =
-                        right_k.to_usize().expect("Dictionary index not usize");
-                    unsafe {
-                        let left_value = left_values.value_unchecked(left_key);
-                        let right_value = right_values.value_unchecked(right_key);
-                        Some($op(left_value, right_value))
-                    }
+        let right_iter = $right
+            .values()
+            .as_any()
+            .downcast_ref::<$value_ty>()
+            .unwrap()
+            .take_iter($right.keys_iter());
+
+        let result = left_iter
+            .zip(right_iter)
+            .map(|(left_value, right_value)| {
+                if let (Some(left), Some(right)) = (left_value, right_value) {
+                    Some($op(left, right))
                 } else {
                     None
                 }