Improve AccumulatorArgs by removing the usgaes of input_types

datafusion/core/src/datasource/physical_plan/parquet/opener.rs

@@ -30,7 +30,7 @@ use crate::datasource::schema_adapter::SchemaAdapterFactory;
 use crate::physical_optimizer::pruning::PruningPredicate;
 use arrow_schema::{ArrowError, SchemaRef};
 use datafusion_common::{exec_err, Result};
-use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+use datafusion_expr::physical_expr::PhysicalExpr;
 use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
 use futures::{StreamExt, TryStreamExt};
 use log::debug;

datafusion/core/src/execution/session_state.rs

@@ -48,6 +48,7 @@ use datafusion_execution::runtime_env::RuntimeEnv;
 use datafusion_execution::TaskContext;
 use datafusion_expr::execution_props::ExecutionProps;
 use datafusion_expr::expr_rewriter::FunctionRewrite;
+use datafusion_expr::physical_expr::PhysicalExpr;
 use datafusion_expr::planner::ExprPlanner;
 use datafusion_expr::registry::{FunctionRegistry, SerializerRegistry};
 use datafusion_expr::simplify::SimplifyInfo;
@@ -61,7 +62,6 @@ use datafusion_optimizer::{
     Analyzer, AnalyzerRule, Optimizer, OptimizerConfig, OptimizerRule,
 };
 use datafusion_physical_expr::create_physical_expr;
-use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_plan::ExecutionPlan;
 use datafusion_sql::parser::{DFParser, Statement};

datafusion/core/src/physical_optimizer/aggregate_statistics.rs

@@ -324,6 +324,7 @@ pub(crate) mod tests {
     use arrow::datatypes::{DataType, Field, Schema};
     use arrow::record_batch::RecordBatch;
     use datafusion_common::cast::as_int64_array;
+    use datafusion_common::ToDFSchema;
     use datafusion_functions_aggregate::count::count_udaf;
     use datafusion_physical_expr::expressions::cast;
     use datafusion_physical_expr::PhysicalExpr;
@@ -421,7 +422,7 @@ pub(crate) mod tests {
         // Return appropriate expr depending if COUNT is for col or table (*)
         pub(crate) fn count_expr(&self, schema: &Schema) -> Arc<dyn AggregateExpr> {
             AggregateExprBuilder::new(count_udaf(), vec![self.column()])
-                .schema(Arc::new(schema.clone()))
+                .dfschema(schema.clone().to_dfschema().unwrap())
                 .name(self.column_name())
                 .build()
                 .unwrap()

datafusion/core/src/physical_optimizer/combine_partial_final_agg.rs

@@ -174,6 +174,7 @@ mod tests {
     use crate::physical_plan::{displayable, Partitioning};
 
     use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+    use datafusion_common::ToDFSchema;
     use datafusion_functions_aggregate::count::count_udaf;
     use datafusion_functions_aggregate::sum::sum_udaf;
     use datafusion_physical_expr::expressions::col;
@@ -279,7 +280,7 @@ mod tests {
         schema: &Schema,
     ) -> Arc<dyn AggregateExpr> {
         AggregateExprBuilder::new(count_udaf(), vec![expr])
-            .schema(Arc::new(schema.clone()))
+            .dfschema(schema.clone().to_dfschema().unwrap())
             .name(name)
             .build()
             .unwrap()
@@ -363,7 +364,7 @@ mod tests {
         let aggr_expr =
             vec![
                 AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?])
-                    .schema(Arc::clone(&schema))
+                    .dfschema(Arc::clone(&schema).to_dfschema()?)
                     .name("Sum(b)")
                     .build()
                     .unwrap(),

datafusion/core/src/physical_optimizer/limit_pushdown.rs

@@ -256,10 +256,10 @@ mod tests {
 
     use arrow_schema::{DataType, Field, Schema, SchemaRef};
     use datafusion_execution::{SendableRecordBatchStream, TaskContext};
+    use datafusion_expr::expressions::column::col;
     use datafusion_expr::Operator;
     use datafusion_physical_expr::expressions::BinaryExpr;
     use datafusion_physical_expr::Partitioning;
-    use datafusion_physical_expr_common::expressions::column::col;
     use datafusion_physical_expr_common::expressions::lit;
     use datafusion_physical_plan::coalesce_batches::CoalesceBatchesExec;
     use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;

datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs

@@ -32,6 +32,7 @@ use datafusion::physical_plan::memory::MemoryExec;
 use datafusion::physical_plan::{collect, displayable, ExecutionPlan};
 use datafusion::prelude::{DataFrame, SessionConfig, SessionContext};
 use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
+use datafusion_common::ToDFSchema;
 use datafusion_functions_aggregate::sum::sum_udaf;
 use datafusion_physical_expr::expressions::col;
 use datafusion_physical_expr::PhysicalSortExpr;
@@ -106,7 +107,7 @@ async fn run_aggregate_test(input1: Vec<RecordBatch>, group_by_columns: Vec<&str
     let aggregate_expr =
         vec![
             AggregateExprBuilder::new(sum_udaf(), vec![col("d", &schema).unwrap()])
-                .schema(Arc::clone(&schema))
+                .dfschema(Arc::clone(&schema).to_dfschema().unwrap())
                 .name("sum1")
                 .build()
                 .unwrap(),

datafusion/physical-expr-common/src/expressions/column.rs → datafusion/expr/src/expressions/column.rs

@@ -21,12 +21,12 @@ use std::any::Any;
 use std::hash::{Hash, Hasher};
 use std::sync::Arc;
 
+use crate::ColumnarValue;
 use arrow::{
     datatypes::{DataType, Schema},
     record_batch::RecordBatch,
 };
 use datafusion_common::{internal_err, Result};
-use datafusion_expr::ColumnarValue;
 
 use crate::physical_expr::{down_cast_any_ref, PhysicalExpr};
 
@@ -89,7 +89,7 @@ impl PhysicalExpr for Column {
     /// Evaluate the expression
     fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
         self.bounds_check(batch.schema().as_ref())?;
-        Ok(ColumnarValue::Array(batch.column(self.index).clone()))
+        Ok(ColumnarValue::Array(Arc::clone(batch.column(self.index))))
     }
 
     fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {

datafusion/expr/src/expressions/mod.rs

@@ -0,0 +1,18 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+pub mod column;

datafusion/expr/src/function.rs

@@ -17,9 +17,10 @@
 
 //! Function module contains typing and signature for built-in and user defined functions.
 
+use crate::physical_expr::PhysicalExpr;
 use crate::ColumnarValue;
 use crate::{Accumulator, Expr, PartitionEvaluator};
-use arrow::datatypes::{DataType, Field, Schema};
+use arrow::datatypes::{DataType, Field};
 use datafusion_common::{DFSchema, Result};
 use std::sync::Arc;
 
@@ -54,9 +55,6 @@ pub struct AccumulatorArgs<'a> {
     /// The return type of the aggregate function.
     pub data_type: &'a DataType,
 
-    /// The schema of the input arguments
-    pub schema: &'a Schema,
-
     /// The schema of the input arguments
     pub dfschema: &'a DFSchema,
 
@@ -94,11 +92,8 @@ pub struct AccumulatorArgs<'a> {
     /// ```
     pub is_distinct: bool,
 
-    /// The input types of the aggregate function.
-    pub input_types: &'a [DataType],
-
     /// The logical expression of arguments the aggregate function takes.
-    pub input_exprs: &'a [Expr],
+    pub input_exprs: &'a [Arc<dyn PhysicalExpr>],
 }
 
 /// [`StateFieldsArgs`] contains information about the fields that an

datafusion/expr/src/lib.rs

@@ -46,10 +46,12 @@ pub mod expr;
 pub mod expr_fn;
 pub mod expr_rewriter;
 pub mod expr_schema;
+pub mod expressions;
 pub mod function;
 pub mod groups_accumulator;
 pub mod interval_arithmetic;
 pub mod logical_plan;
+pub mod physical_expr;
 pub mod planner;
 pub mod registry;
 pub mod simplify;

datafusion/physical-expr-common/src/physical_expr.rs → datafusion/expr/src/physical_expr.rs

@@ -23,15 +23,15 @@ use std::sync::Arc;
 use crate::expressions::column::Column;
 use crate::utils::scatter;
 
+use crate::interval_arithmetic::Interval;
+use crate::sort_properties::ExprProperties;
+use crate::ColumnarValue;
 use arrow::array::BooleanArray;
 use arrow::compute::filter_record_batch;
 use arrow::datatypes::{DataType, Schema, SchemaRef};
 use arrow::record_batch::RecordBatch;
 use datafusion_common::tree_node::{Transformed, TreeNode};
 use datafusion_common::{internal_err, not_impl_err, plan_err, Result};
-use datafusion_expr::interval_arithmetic::Interval;
-use datafusion_expr::sort_properties::ExprProperties;
-use datafusion_expr::ColumnarValue;
 
 /// See [create_physical_expr](https://docs.rs/datafusion/latest/datafusion/physical_expr/fn.create_physical_expr.html)
 /// for examples of creating `PhysicalExpr` from `Expr`

datafusion/expr/src/tree_node.rs

@@ -17,15 +17,20 @@
 
 //! Tree node implementation for logical expr
 
+use std::fmt::{self, Display, Formatter};
+use std::sync::Arc;
+
 use crate::expr::{
     AggregateFunction, AggregateFunctionDefinition, Alias, Between, BinaryExpr, Case,
     Cast, GroupingSet, InList, InSubquery, Like, Placeholder, ScalarFunction, Sort,
     TryCast, Unnest, WindowFunction,
 };
+use crate::physical_expr::{with_new_children_if_necessary, PhysicalExpr};
 use crate::{Expr, ExprFunctionExt};
 
 use datafusion_common::tree_node::{
-    Transformed, TreeNode, TreeNodeIterator, TreeNodeRecursion,
+    ConcreteTreeNode, DynTreeNode, Transformed, TreeNode, TreeNodeIterator,
+    TreeNodeRecursion,
 };
 use datafusion_common::{map_until_stop_and_collect, Result};
 
@@ -401,3 +406,82 @@ fn transform_vec<F: FnMut(Expr) -> Result<Transformed<Expr>>>(
 ) -> Result<Transformed<Vec<Expr>>> {
     ve.into_iter().map_until_stop_and_collect(f)
 }
+
+impl DynTreeNode for dyn PhysicalExpr {
+    fn arc_children(&self) -> Vec<&Arc<Self>> {
+        self.children()
+    }
+
+    fn with_new_arc_children(
+        &self,
+        arc_self: Arc<Self>,
+        new_children: Vec<Arc<Self>>,
+    ) -> Result<Arc<Self>> {
+        with_new_children_if_necessary(arc_self, new_children)
+    }
+}
+
+/// A node object encapsulating a [`PhysicalExpr`] node with a payload. Since there are
+/// two ways to access child plans—directly from the plan  and through child nodes—it's
+/// recommended to perform mutable operations via [`Self::update_expr_from_children`].
+#[derive(Debug)]
+pub struct ExprContext<T: Sized> {
+    /// The physical expression associated with this context.
+    pub expr: Arc<dyn PhysicalExpr>,
+    /// Custom data payload of the node.
+    pub data: T,
+    /// Child contexts of this node.
+    pub children: Vec<Self>,
+}
+
+impl<T> ExprContext<T> {
+    pub fn new(expr: Arc<dyn PhysicalExpr>, data: T, children: Vec<Self>) -> Self {
+        Self {
+            expr,
+            data,
+            children,
+        }
+    }
+
+    pub fn update_expr_from_children(mut self) -> Result<Self> {
+        let children_expr = self.children.iter().map(|c| Arc::clone(&c.expr)).collect();
+        self.expr = with_new_children_if_necessary(self.expr, children_expr)?;
+        Ok(self)
+    }
+}
+
+impl<T: Default> ExprContext<T> {
+    pub fn new_default(plan: Arc<dyn PhysicalExpr>) -> Self {
+        let children = plan
+            .children()
+            .into_iter()
+            .cloned()
+            .map(Self::new_default)
+            .collect();
+        Self::new(plan, Default::default(), children)
+    }
+}
+
+impl<T: Display> Display for ExprContext<T> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "expr: {:?}", self.expr)?;
+        write!(f, "data:{}", self.data)?;
+        write!(f, "")
+    }
+}
+
+impl<T> ConcreteTreeNode for ExprContext<T> {
+    fn children(&self) -> &[Self] {
+        &self.children
+    }
+
+    fn take_children(mut self) -> (Self, Vec<Self>) {
+        let children = std::mem::take(&mut self.children);
+        (self, children)
+    }
+
+    fn with_new_children(mut self, children: Vec<Self>) -> Result<Self> {
+        self.children = children;
+        self.update_expr_from_children()
+    }
+}