alamb commented on code in PR #4896:
URL: https://github.com/apache/arrow-rs/pull/4896#discussion_r1353171861


##########
arrow-flight/src/encode.rs:
##########
@@ -520,14 +542,92 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_dictionary_hydration() {
+        let arr: DictionaryArray<UInt16Type> = vec!["a", "a", 
"b"].into_iter().collect();
+        let schema = Arc::new(Schema::new(vec![Field::new_dictionary(
+            "dict",
+            DataType::UInt16,
+            DataType::Utf8,
+            false,
+        )]));
+        let batch = RecordBatch::try_new(schema, vec![Arc::new(arr)]).unwrap();
+        let encoder = FlightDataEncoderBuilder::default()
+            .build(futures::stream::once(async { Ok(batch) }));
+        let mut decoder = FlightDataDecoder::new(encoder);
+        let expected_schema =
+            Schema::new(vec![Field::new("dict", DataType::Utf8, false)]);
+        let expected_schema = Arc::new(expected_schema);
+        while let Some(decoded) = decoder.next().await {
+            let decoded = decoded.unwrap();
+            match decoded.payload {
+                DecodedPayload::None => {}
+                DecodedPayload::Schema(s) => assert_eq!(s, expected_schema),
+                DecodedPayload::RecordBatch(b) => {
+                    assert_eq!(b.schema(), expected_schema);
+                    let expected_array = StringArray::from(vec!["a", "a", 
"b"]);
+                    let actual_array = b.column_by_name("dict").unwrap();
+                    let actual_array = 
downcast_array::<StringArray>(actual_array);
+
+                    assert_eq!(actual_array, expected_array);
+                }
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_send_dictionaries() {
+        let schema = Arc::new(Schema::new(vec![Field::new_dictionary(
+            "dict",
+            DataType::UInt16,
+            DataType::Utf8,
+            false,
+        )]));
+
+        let arr_one: Arc<DictionaryArray<UInt16Type>> =
+            Arc::new(vec!["a", "a", "b"].into_iter().collect());
+        let arr_two: Arc<DictionaryArray<UInt16Type>> =

Review Comment:
   👌 



##########
arrow-flight/src/encode.rs:
##########
@@ -520,14 +542,92 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_dictionary_hydration() {

Review Comment:
   👍 



##########
arrow-flight/src/encode.rs:
##########
@@ -520,14 +542,92 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_dictionary_hydration() {
+        let arr: DictionaryArray<UInt16Type> = vec!["a", "a", 
"b"].into_iter().collect();
+        let schema = Arc::new(Schema::new(vec![Field::new_dictionary(
+            "dict",
+            DataType::UInt16,
+            DataType::Utf8,
+            false,
+        )]));
+        let batch = RecordBatch::try_new(schema, vec![Arc::new(arr)]).unwrap();
+        let encoder = FlightDataEncoderBuilder::default()
+            .build(futures::stream::once(async { Ok(batch) }));
+        let mut decoder = FlightDataDecoder::new(encoder);
+        let expected_schema =
+            Schema::new(vec![Field::new("dict", DataType::Utf8, false)]);
+        let expected_schema = Arc::new(expected_schema);
+        while let Some(decoded) = decoder.next().await {
+            let decoded = decoded.unwrap();
+            match decoded.payload {
+                DecodedPayload::None => {}
+                DecodedPayload::Schema(s) => assert_eq!(s, expected_schema),
+                DecodedPayload::RecordBatch(b) => {
+                    assert_eq!(b.schema(), expected_schema);
+                    let expected_array = StringArray::from(vec!["a", "a", 
"b"]);
+                    let actual_array = b.column_by_name("dict").unwrap();
+                    let actual_array = 
downcast_array::<StringArray>(actual_array);
+
+                    assert_eq!(actual_array, expected_array);
+                }
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_send_dictionaries() {
+        let schema = Arc::new(Schema::new(vec![Field::new_dictionary(
+            "dict",
+            DataType::UInt16,
+            DataType::Utf8,
+            false,
+        )]));
+
+        let arr_one: Arc<DictionaryArray<UInt16Type>> =
+            Arc::new(vec!["a", "a", "b"].into_iter().collect());
+        let arr_two: Arc<DictionaryArray<UInt16Type>> =

Review Comment:
   👌 



##########
arrow-flight/src/encode.rs:
##########
@@ -520,14 +542,92 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_dictionary_hydration() {

Review Comment:
   👍 



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

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to