Minor(proto): Implement TryFrom<&DFSchema> for protobuf::DfSchema (…

…#7505)
apache · Sep 8, 2023 · c7a93fc · c7a93fc
1 parent 270dc8b
commit c7a93fc
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 5 deletions.
diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs
@@ -1455,8 +1455,8 @@ mod roundtrip_tests {
     };
     use datafusion::test_util::{TestTableFactory, TestTableProvider};
     use datafusion_common::{
-        internal_err, not_impl_err, plan_err, DFSchemaRef, DataFusionError, Result,
-        ScalarValue,
+        internal_err, not_impl_err, plan_err, DFField, DFSchema, DFSchemaRef,
+        DataFusionError, Result, ScalarValue,
     };
     use datafusion_expr::expr::{
         self, Between, BinaryExpr, Case, Cast, GroupingSet, InList, Like, ScalarFunction,
@@ -2367,6 +2367,35 @@ mod roundtrip_tests {
         assert_eq!(schema, returned_schema);
     }
 
+    #[test]
+    fn roundtrip_dfschema() {
+        let dfschema = DFSchema::new_with_metadata(
+            vec![
+                DFField::new_unqualified("a", DataType::Int64, false),
+                DFField::new(Some("t"), "b", DataType::Decimal128(15, 2), true)
+                    .with_metadata(HashMap::from([(
+                        String::from("k1"),
+                        String::from("v1"),
+                    )])),
+            ],
+            HashMap::from([
+                (String::from("k2"), String::from("v2")),
+                (String::from("k3"), String::from("v3")),
+            ]),
+        )
+        .unwrap();
+        let proto_dfschema: super::protobuf::DfSchema = (&dfschema).try_into().unwrap();
+        let returned_dfschema: DFSchema = (&proto_dfschema).try_into().unwrap();
+        assert_eq!(dfschema, returned_dfschema);
+
+        let arc_dfschema = Arc::new(dfschema.clone());
+        let proto_dfschema: super::protobuf::DfSchema =
+            (&arc_dfschema).try_into().unwrap();
+        let returned_arc_dfschema: DFSchemaRef = proto_dfschema.try_into().unwrap();
+        assert_eq!(arc_dfschema, returned_arc_dfschema);
+        assert_eq!(dfschema, *returned_arc_dfschema);
+    }
+
     #[test]
     fn roundtrip_not() {
         let test_expr = Expr::Not(Box::new(lit(1.0_f32)));

diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs
@@ -34,7 +34,9 @@ use arrow::datatypes::{
     DataType, Field, IntervalMonthDayNanoType, IntervalUnit, Schema, SchemaRef, TimeUnit,
     UnionMode,
 };
-use datafusion_common::{Column, DFField, DFSchemaRef, OwnedTableReference, ScalarValue};
+use datafusion_common::{
+    Column, DFField, DFSchema, DFSchemaRef, OwnedTableReference, ScalarValue,
+};
 use datafusion_expr::expr::{
     self, Alias, Between, BinaryExpr, Cast, GetFieldAccess, GetIndexedField, GroupingSet,
     InList, Like, Placeholder, ScalarFunction, ScalarUDF, Sort,
@@ -300,10 +302,10 @@ impl TryFrom<&DFField> for protobuf::DfField {
     }
 }
 
-impl TryFrom<&DFSchemaRef> for protobuf::DfSchema {
+impl TryFrom<&DFSchema> for protobuf::DfSchema {
     type Error = Error;
 
-    fn try_from(s: &DFSchemaRef) -> Result<Self, Self::Error> {
+    fn try_from(s: &DFSchema) -> Result<Self, Self::Error> {
         let columns = s
             .fields()
             .iter()
@@ -316,6 +318,14 @@ impl TryFrom<&DFSchemaRef> for protobuf::DfSchema {
     }
 }
 
+impl TryFrom<&DFSchemaRef> for protobuf::DfSchema {
+    type Error = Error;
+
+    fn try_from(s: &DFSchemaRef) -> Result<Self, Self::Error> {
+        s.as_ref().try_into()
+    }
+}
+
 impl From<&StringifiedPlan> for protobuf::StringifiedPlan {
     fn from(stringified_plan: &StringifiedPlan) -> Self {
         Self {