[GitHub] [arrow] jhorstmann commented on a change in pull request #8051: ARROW-9853: [RUST] Implement take kernel for dictionary arrays

[GitBox <gi...@apache.org>]

jhorstmann commented on a change in pull request #8051:
URL: https://github.com/apache/arrow/pull/8051#discussion_r476688318



##########
File path: rust/arrow/src/compute/kernels/take.rs
##########
@@ -657,4 +694,71 @@ mod tests {
             vec![None],
         );
     }
+
+    #[test]
+    fn test_take_dict() {
+        let keys_builder = Int16Builder::new(8);
+        let values_builder = StringBuilder::new(4);
+
+        let mut dict_builder = StringDictionaryBuilder::new(keys_builder, values_builder);
+
+        dict_builder.append("foo").unwrap();
+        dict_builder.append("bar").unwrap();
+        dict_builder.append("").unwrap();
+        dict_builder.append_null().unwrap();
+        dict_builder.append("foo").unwrap();
+        dict_builder.append("bar").unwrap();
+        dict_builder.append("bar").unwrap();
+        dict_builder.append("foo").unwrap();
+
+        let array = dict_builder.finish();
+
+        let expected_values = StringArray::try_from(vec!["foo", "bar", ""]).unwrap();
+        assert_eq!(
+            &expected_values,
+            array
+                .values()
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .unwrap()
+        );
+
+        let array_ref: Arc<dyn Array> = Arc::new(array);
+
+        let indices = UInt32Array::from(vec![
+            Some(0),
+            Some(7),
+            None,
+            Some(5),
+            Some(6),
+            Some(2),
+            Some(3),
+        ]);
+
+        let result = take(&array_ref, &indices, None).unwrap();
+        let result = result
+            .as_any()
+            .downcast_ref::<DictionaryArray<Int16Type>>()
+            .unwrap();
+
+        let result_values: StringArray = result.values().data().into();
+
+        // dictionary values should stay the same
+        assert_eq!(&expected_values, &result_values);
+
+        let result_keys: Int16Array = result.keys().collect::<Vec<_>>().into();
+        let expected_keys = Int16Array::from(vec![
+            Some(0),
+            Some(0),
+            None,
+            Some(1),

Review comment:
       Good point, I added some comments and also moved the assertion of the dictionary values closer to the assertion of the keys so hopefully their relation is now more visible.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org