revert schema change v3

xinlifoobar · Aug 5, 2024 · 4d50a94 · 4d50a94
1 parent 1ccc740
commit 4d50a94
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 4 deletions.
diff --git a/datafusion/core/src/physical_optimizer/combine_partial_final_agg.rs b/datafusion/core/src/physical_optimizer/combine_partial_final_agg.rs
@@ -279,7 +279,7 @@ mod tests {
         schema: &Schema,
     ) -> Arc<dyn AggregateExpr> {
         AggregateExprBuilder::new(count_udaf(), vec![expr])
-            .schema(schema.clone())
+            .schema(Arc::new(schema.clone()))
             .name(name)
             .build()
             .unwrap()
@@ -363,7 +363,7 @@ mod tests {
         let aggr_expr =
             vec![
                 AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?])
-                    .dfschema(Arc::clone(&schema).to_dfschema()?)
+                    .schema(Arc::clone(&schema))
                     .name("Sum(b)")
                     .build()
                     .unwrap(),

diff --git a/datafusion/physical-expr-common/src/aggregate/mod.rs b/datafusion/physical-expr-common/src/aggregate/mod.rs
@@ -108,6 +108,7 @@ pub fn create_aggregate_expr_with_dfschema(
     builder = builder.sort_exprs(sort_exprs.to_vec());
     builder = builder.order_by(ordering_req.to_vec());
     builder = builder.logical_exprs(input_exprs.to_vec());
+    builder = builder.dfschema(dfschema.clone());
     let schema: Schema = dfschema.into();
     builder = builder.schema(Arc::new(schema));
     builder = builder.name(name);

diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -1790,7 +1790,7 @@ mod tests {
     // Median(a)
     fn test_median_agg_expr(schema: SchemaRef) -> Result<Arc<dyn AggregateExpr>> {
         AggregateExprBuilder::new(median_udaf(), vec![col("a", &schema)?])
-            .schema(Arc::clone(&schema))
+            .schema(schema)
             .name("MEDIAN(a)")
             .build()
     }

diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs
@@ -489,7 +489,8 @@ impl AsExecutionPlan for protobuf::PhysicalPlanNode {
 
                                             // TODO: approx_percentile_cont and approx_percentile_cont_weight are not supported for UDAF from protobuf yet.
                                             // TODO: `order by` is not supported for UDAF yet
-                                            AggregateExprBuilder::new(agg_udf, input_phy_expr).schema(Arc::clone(&physical_schema)).name(name).with_ignore_nulls(agg_node.ignore_nulls).with_distinct(agg_node.distinct).build()                                        }
+                                            AggregateExprBuilder::new(agg_udf, input_phy_expr).schema(Arc::clone(&physical_schema)).name(name).with_ignore_nulls(agg_node.ignore_nulls).with_distinct(agg_node.distinct).build()
+                                        }
                                     }
                                 }).transpose()?.ok_or_else(|| {
                                     proto_error("Invalid AggregateExpr, missing aggregate_function")