implement proto

alamb · Nov 9, 2023 · b643d72 · b643d72
1 parent 71a2742
commit b643d72
Show file tree

Hide file tree

Showing 11 changed files with 159 additions and 34 deletions.
diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs
@@ -364,7 +364,9 @@ fn create_physical_name(e: &Expr, is_first_expr: bool) -> Result<String> {
         Expr::Sort { .. } => {
             internal_err!("Create physical name does not support sort expression")
         }
-        Expr::Wildcard { .. } => internal_err!("Create physical name does not support wildcard"),
+        Expr::Wildcard { .. } => {
+            internal_err!("Create physical name does not support wildcard")
+        }
         Expr::Placeholder(_) => {
             internal_err!("Create physical name does not support placeholder")
         }

diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs
@@ -42,7 +42,11 @@ use datafusion::{assert_batches_eq, assert_batches_sorted_eq};
 use datafusion_common::{DataFusionError, ScalarValue, UnnestOptions};
 use datafusion_execution::config::SessionConfig;
 use datafusion_expr::expr::{GroupingSet, Sort};
-use datafusion_expr::{array_agg, avg, col, count, exists, expr, in_subquery, lit, max, out_ref_col, scalar_subquery, sum, AggregateFunction, Expr, ExprSchemable, WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunction, wildcard};
+use datafusion_expr::{
+    array_agg, avg, col, count, exists, expr, in_subquery, lit, max, out_ref_col,
+    scalar_subquery, sum, wildcard, AggregateFunction, Expr, ExprSchemable, WindowFrame,
+    WindowFrameBound, WindowFrameUnits, WindowFunction,
+};
 use datafusion_physical_expr::var_provider::{VarProvider, VarType};
 
 #[tokio::test]

diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs
@@ -1595,7 +1595,7 @@ mod tests {
 
         let plan = table_scan(Some("t1"), &employee_schema(), None)?
             .join_using(t2, JoinType::Inner, vec!["id"])?
-            .project(vec![Expr::Wildcard{qualifier: None}])?
+            .project(vec![Expr::Wildcard { qualifier: None }])?
             .build()?;
 
         // id column should only show up once in projection

diff --git a/datafusion/optimizer/src/analyzer/count_wildcard_rule.rs b/datafusion/optimizer/src/analyzer/count_wildcard_rule.rs
@@ -129,15 +129,17 @@ impl TreeNodeRewriter for CountWildcardRewriter {
                 order_by,
                 window_frame,
             }) if args.len() == 1 => match args[0] {
-                Expr::Wildcard { qualifier: None } => Expr::WindowFunction(expr::WindowFunction {
-                    fun: window_function::WindowFunction::AggregateFunction(
-                        aggregate_function::AggregateFunction::Count,
-                    ),
-                    args: vec![lit(COUNT_STAR_EXPANSION)],
-                    partition_by,
-                    order_by,
-                    window_frame,
-                }),
+                Expr::Wildcard { qualifier: None } => {
+                    Expr::WindowFunction(expr::WindowFunction {
+                        fun: window_function::WindowFunction::AggregateFunction(
+                            aggregate_function::AggregateFunction::Count,
+                        ),
+                        args: vec![lit(COUNT_STAR_EXPANSION)],
+                        partition_by,
+                        order_by,
+                        window_frame,
+                    })
+                }
 
                 _ => old_expr,
             },
@@ -148,13 +150,15 @@ impl TreeNodeRewriter for CountWildcardRewriter {
                 filter,
                 order_by,
             }) if args.len() == 1 => match args[0] {
-                Expr::Wildcard { qualifier: None } => Expr::AggregateFunction(AggregateFunction {
-                    fun: aggregate_function::AggregateFunction::Count,
-                    args: vec![lit(COUNT_STAR_EXPANSION)],
-                    distinct,
-                    filter,
-                    order_by,
-                }),
+                Expr::Wildcard { qualifier: None } => {
+                    Expr::AggregateFunction(AggregateFunction {
+                        fun: aggregate_function::AggregateFunction::Count,
+                        args: vec![lit(COUNT_STAR_EXPANSION)],
+                        distinct,
+                        filter,
+                        order_by,
+                    })
+                }
                 _ => old_expr,
             },
 
@@ -219,7 +223,11 @@ mod tests {
     use arrow::datatypes::DataType;
     use datafusion_common::{Result, ScalarValue};
     use datafusion_expr::expr::Sort;
-    use datafusion_expr::{col, count, exists, expr, in_subquery, lit, logical_plan::LogicalPlanBuilder, max, out_ref_col, scalar_subquery, AggregateFunction, Expr, WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunction, wildcard};
+    use datafusion_expr::{
+        col, count, exists, expr, in_subquery, lit, logical_plan::LogicalPlanBuilder,
+        max, out_ref_col, scalar_subquery, wildcard, AggregateFunction, Expr,
+        WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunction,
+    };
 
     fn assert_plan_eq(plan: &LogicalPlan, expected: &str) -> Result<()> {
         assert_analyzed_plan_eq_display_indent(

diff --git a/datafusion/optimizer/src/analyzer/inline_table_scan.rs b/datafusion/optimizer/src/analyzer/inline_table_scan.rs
@@ -126,7 +126,7 @@ fn generate_projection_expr(
             ));
         }
     } else {
-        exprs.push(Expr::Wildcard { qualifier: None});
+        exprs.push(Expr::Wildcard { qualifier: None });
     }
     Ok(exprs)
 }

diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto
@@ -363,7 +363,7 @@ message LogicalExprNode {
     SortExprNode sort = 12;
     NegativeNode negative = 13;
     InListNode in_list = 14;
-    bool wildcard = 15;
+    Wildcard wildcard = 15;
     ScalarFunctionNode scalar_function = 16;
     TryCastNode try_cast = 17;
 
@@ -399,6 +399,10 @@ message LogicalExprNode {
   }
 }
 
+message Wildcard {
+  optional string qualifier = 1;
+}
+
 message PlaceholderNode {
   string id = 1;
   ArrowType data_type = 2;

diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs
diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs
diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs
@@ -1294,7 +1294,9 @@ pub fn parse_expr(
                 .collect::<Result<Vec<_>, _>>()?,
             in_list.negated,
         ))),
-        ExprType::Wildcard(_) => Ok(Expr::Wildcard),
+        ExprType::Wildcard(protobuf::Wildcard { qualifier }) => Ok(Expr::Wildcard {
+            qualifier: qualifier.clone(),
+        }),
         ExprType::ScalarFunction(expr) => {
             let scalar_function = protobuf::ScalarFunction::try_from(expr.fun)
                 .map_err(|_| Error::unknown("ScalarFunction", expr.fun))?;

diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs
@@ -960,8 +960,10 @@ impl TryFrom<&Expr> for protobuf::LogicalExprNode {
                     expr_type: Some(ExprType::InList(expr)),
                 }
             }
-            Expr::Wildcard => Self {
-                expr_type: Some(ExprType::Wildcard(true)),
+            Expr::Wildcard { qualifier } => Self {
+                expr_type: Some(ExprType::Wildcard(protobuf::Wildcard {
+                    qualifier: qualifier.clone(),
+                })),
             },
             Expr::ScalarSubquery(_)
             | Expr::InSubquery(_)
@@ -1052,11 +1054,6 @@ impl TryFrom<&Expr> for protobuf::LogicalExprNode {
                     })),
                 }
             }
-
-            Expr::QualifiedWildcard { .. } => return Err(Error::General(
-                "Proto serialization error: Expr::QualifiedWildcard { .. } not supported"
-                    .to_string(),
-            )),
         };
 
         Ok(expr_node)

diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
@@ -1147,7 +1147,17 @@ fn roundtrip_inlist() {
 
 #[test]
 fn roundtrip_wildcard() {
-    let test_expr = Expr::Wildcard;
+    let test_expr = Expr::Wildcard { qualifier: None };
+
+    let ctx = SessionContext::new();
+    roundtrip_expr_test(test_expr, ctx);
+}
+
+#[test]
+fn roundtrip_qualified_wildcard() {
+    let test_expr = Expr::Wildcard {
+        qualifier: Some("foo".into()),
+    };
 
     let ctx = SessionContext::new();
     roundtrip_expr_test(test_expr, ctx);