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

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

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



##########
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:
       For clarity, I would try to describe in comments how we arrive to one of these values, e.g.
   
   ```
   // index number 5 corresponds to "bar", that first appears in index 1, thus 1
   ```




----------------------------------------------------------------
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