alamb commented on code in PR #10032:
URL:
https://github.com/apache/arrow-datafusion/pull/10032#discussion_r1560001477
##########
datafusion/optimizer/src/analyzer/function_rewrite.rs:
##########
@@ -37,86 +39,53 @@ impl ApplyFunctionRewrites {
pub fn new(function_rewrites: Vec<Arc<dyn FunctionRewrite + Send + Sync>>)
-> Self {
Self { function_rewrites }
}
-}
-
-impl AnalyzerRule for ApplyFunctionRewrites {
- fn name(&self) -> &str {
- "apply_function_rewrites"
- }
-
- fn analyze(&self, plan: LogicalPlan, options: &ConfigOptions) ->
Result<LogicalPlan> {
- self.analyze_internal(&plan, options)
- }
-}
-impl ApplyFunctionRewrites {
- fn analyze_internal(
+ /// Rewrite a single plan, and all its expressions using the provided
rewriters
+ fn rewrite_plan(
&self,
- plan: &LogicalPlan,
+ plan: LogicalPlan,
options: &ConfigOptions,
- ) -> Result<LogicalPlan> {
- // optimize child plans first
- let new_inputs = plan
- .inputs()
- .iter()
- .map(|p| self.analyze_internal(p, options))
- .collect::<Result<Vec<_>>>()?;
-
+ ) -> Result<Transformed<LogicalPlan>> {
// get schema representing all available input fields. This is used
for data type
// resolution only, so order does not matter here
- let mut schema = merge_schema(new_inputs.iter().collect());
+ let mut schema = merge_schema(plan.inputs());
- if let LogicalPlan::TableScan(ts) = plan {
+ if let LogicalPlan::TableScan(ts) = &plan {
let source_schema = DFSchema::try_from_qualified_schema(
ts.table_name.clone(),
&ts.source.schema(),
)?;
schema.merge(&source_schema);
}
- let mut expr_rewrite = OperatorToFunctionRewriter {
- function_rewrites: &self.function_rewrites,
- options,
- schema: &schema,
- };
+ let name_preserver = NamePreserver::new(&plan);
+
+ plan.map_expressions(|expr| {
+ let original_name = name_preserver.save(&expr)?;
- let new_expr = plan
- .expressions()
- .into_iter()
- .map(|expr| {
- // ensure names don't change:
- // https://github.com/apache/arrow-datafusion/issues/3555
- rewrite_preserving_name(expr, &mut expr_rewrite)
- })
- .collect::<Result<Vec<_>>>()?;
+ // recursively transform the expression, applying the rewrites at
each step
+ let result = expr.transform_up(&|expr| {
+ let mut result = Transformed::no(expr);
+ for rewriter in self.function_rewrites.iter() {
+ result = result.transform_data(|expr| {
+ rewriter.rewrite(expr, &schema, options)
+ })?;
+ }
+ Ok(result)
+ })?;
- plan.with_new_exprs(new_expr, new_inputs)
Review Comment:
this copies the plan + expressions
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
