Minor: Support protobuf serialization for Utf8View and BinaryView (#1…

…2165) * Minor: Support protobuf serialization for ScalarValue::Utf8View and ScalarValue::BinaryView * Use correct protobuf type, add coverage for DataType serialization --------- Co-authored-by: Andrew Lamb <[email protected]>
apache · Sep 7, 2024 · 73f3086 · 73f3086
1 parent ddbdf4b
commit 73f3086
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 8 deletions.
diff --git a/datafusion/proto-common/src/to_proto/mod.rs b/datafusion/proto-common/src/to_proto/mod.rs
@@ -153,9 +153,11 @@ impl TryFrom<&DataType> for protobuf::arrow_type::ArrowTypeEnum {
                 Self::Interval(protobuf::IntervalUnit::from(interval_unit) as i32)
             }
             DataType::Binary => Self::Binary(EmptyMessage {}),
+            DataType::BinaryView => Self::BinaryView(EmptyMessage {}),
             DataType::FixedSizeBinary(size) => Self::FixedSizeBinary(*size),
             DataType::LargeBinary => Self::LargeBinary(EmptyMessage {}),
             DataType::Utf8 => Self::Utf8(EmptyMessage {}),
+            DataType::Utf8View => Self::Utf8View(EmptyMessage {}),
             DataType::LargeUtf8 => Self::LargeUtf8(EmptyMessage {}),
             DataType::List(item_type) => Self::List(Box::new(protobuf::List {
                 field_type: Some(Box::new(item_type.as_ref().try_into()?)),
@@ -210,7 +212,7 @@ impl TryFrom<&DataType> for protobuf::arrow_type::ArrowTypeEnum {
                     "Proto serialization error: The RunEndEncoded data type is not yet supported".to_owned()
                 ))
             }
-            DataType::Utf8View | DataType::BinaryView | DataType::ListView(_) | DataType::LargeListView(_) => {
+            DataType::ListView(_) | DataType::LargeListView(_) => {
                 return Err(Error::General(format!("Proto serialization error: {val} not yet supported")))
             }
         };

diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
@@ -1191,7 +1191,7 @@ impl LogicalExtensionCodec for UDFExtensionCodec {
 }
 
 #[test]
-fn round_trip_scalar_values() {
+fn round_trip_scalar_values_and_data_types() {
     let should_pass: Vec<ScalarValue> = vec![
         ScalarValue::Boolean(None),
         ScalarValue::Float32(None),
@@ -1245,6 +1245,8 @@ fn round_trip_scalar_values() {
         ScalarValue::UInt64(Some(0)),
         ScalarValue::Utf8(Some(String::from("Test string   "))),
         ScalarValue::LargeUtf8(Some(String::from("Test Large utf8"))),
+        ScalarValue::Utf8View(Some(String::from("Test stringview"))),
+        ScalarValue::BinaryView(Some(b"binaryview".to_vec())),
         ScalarValue::Date32(Some(0)),
         ScalarValue::Date32(Some(i32::MAX)),
         ScalarValue::Date32(None),
@@ -1471,21 +1473,38 @@ fn round_trip_scalar_values() {
         ScalarValue::FixedSizeBinary(5, None),
     ];
 
-    for test_case in should_pass.into_iter() {
-        let proto: protobuf::ScalarValue = (&test_case)
-            .try_into()
-            .expect("failed conversion to protobuf");
-
+    // ScalarValue directly
+    for test_case in should_pass.iter() {
+        let proto: protobuf::ScalarValue =
+            test_case.try_into().expect("failed conversion to protobuf");
         let roundtrip: ScalarValue = (&proto)
             .try_into()
             .expect("failed conversion from protobuf");
 
         assert_eq!(
-            test_case, roundtrip,
+            test_case, &roundtrip,
             "ScalarValue was not the same after round trip!\n\n\
                         Input: {test_case:?}\n\nRoundtrip: {roundtrip:?}"
         );
     }
+
+    //  DataType conversion
+    for test_case in should_pass.iter() {
+        let dt = test_case.data_type();
+
+        let proto: protobuf::ArrowType = (&dt)
+            .try_into()
+            .expect("datatype failed conversion to protobuf");
+        let roundtrip: DataType = (&proto)
+            .try_into()
+            .expect("datatype failed conversion from protobuf");
+
+        assert_eq!(
+            dt, roundtrip,
+            "DataType was not the same after round trip!\n\n\
+                        Input: {dt:?}\n\nRoundtrip: {roundtrip:?}"
+        );
+    }
 }
 
 // create a map array [{joe:1}, {blogs:2, foo:4}, {}, null] for testing