alamb commented on code in PR #10066:
URL:
https://github.com/apache/arrow-datafusion/pull/10066#discussion_r1563054898
##########
datafusion/optimizer/src/analyzer/count_wildcard_rule.rs:
##########
@@ -47,181 +38,43 @@ impl CountWildcardRule {
impl AnalyzerRule for CountWildcardRule {
fn analyze(&self, plan: LogicalPlan, _: &ConfigOptions) ->
Result<LogicalPlan> {
- plan.transform_down(&analyze_internal).data()
+ plan.transform_down_with_subqueries(&analyze_internal)
+ .data()
}
fn name(&self) -> &str {
"count_wildcard_rule"
}
}
-fn analyze_internal(plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
- let mut rewriter = CountWildcardRewriter {};
- match plan {
- LogicalPlan::Window(window) => {
- let window_expr = window
- .window_expr
- .iter()
- .map(|expr| rewrite_preserving_name(expr.clone(), &mut
rewriter))
- .collect::<Result<Vec<_>>>()?;
-
- Ok(Transformed::yes(
- LogicalPlanBuilder::from((*window.input).clone())
- .window(window_expr)?
- .build()?,
- ))
- }
- LogicalPlan::Aggregate(agg) => {
- let aggr_expr = agg
- .aggr_expr
- .iter()
- .map(|expr| rewrite_preserving_name(expr.clone(), &mut
rewriter))
- .collect::<Result<Vec<_>>>()?;
-
- Ok(Transformed::yes(LogicalPlan::Aggregate(
- Aggregate::try_new(agg.input.clone(), agg.group_expr,
aggr_expr)?,
- )))
- }
- LogicalPlan::Sort(Sort { expr, input, fetch }) => {
- let sort_expr = expr
- .iter()
- .map(|expr| rewrite_preserving_name(expr.clone(), &mut
rewriter))
- .collect::<Result<Vec<_>>>()?;
- Ok(Transformed::yes(LogicalPlan::Sort(Sort {
- expr: sort_expr,
- input,
- fetch,
- })))
- }
- LogicalPlan::Projection(projection) => {
- let projection_expr = projection
- .expr
- .iter()
- .map(|expr| rewrite_preserving_name(expr.clone(), &mut
rewriter))
- .collect::<Result<Vec<_>>>()?;
- Ok(Transformed::yes(LogicalPlan::Projection(
- Projection::try_new(projection_expr, projection.input)?,
- )))
- }
- LogicalPlan::Filter(Filter {
- predicate, input, ..
- }) => {
- let predicate = rewrite_preserving_name(predicate, &mut rewriter)?;
- Ok(Transformed::yes(LogicalPlan::Filter(Filter::try_new(
- predicate, input,
- )?)))
- }
-
- _ => Ok(Transformed::no(plan)),
- }
+fn is_wildcard(expr: &Expr) -> bool {
+ matches!(expr, Expr::Wildcard { qualifier: None })
}
-
-struct CountWildcardRewriter {}
-
-impl TreeNodeRewriter for CountWildcardRewriter {
- type Node = Expr;
-
- fn f_up(&mut self, old_expr: Expr) -> Result<Transformed<Expr>> {
- Ok(match old_expr.clone() {
- Expr::WindowFunction(expr::WindowFunction {
- fun:
- expr::WindowFunctionDefinition::AggregateFunction(
- aggregate_function::AggregateFunction::Count,
- ),
- args,
- partition_by,
- order_by,
- window_frame,
- null_treatment,
- }) if args.len() == 1 => match args[0] {
- Expr::Wildcard { qualifier: None } => {
- Transformed::yes(Expr::WindowFunction(expr::WindowFunction
{
- fun: expr::WindowFunctionDefinition::AggregateFunction(
- aggregate_function::AggregateFunction::Count,
- ),
- args: vec![lit(COUNT_STAR_EXPANSION)],
- partition_by,
- order_by,
- window_frame,
- null_treatment,
- }))
- }
-
- _ => Transformed::no(old_expr),
- },
- Expr::AggregateFunction(AggregateFunction {
- func_def:
- AggregateFunctionDefinition::BuiltIn(
- aggregate_function::AggregateFunction::Count,
- ),
- args,
- distinct,
- filter,
- order_by,
- null_treatment,
- }) if args.len() == 1 => match args[0] {
- Expr::Wildcard { qualifier: None } => {
-
Transformed::yes(Expr::AggregateFunction(AggregateFunction::new(
- aggregate_function::AggregateFunction::Count,
- vec![lit(COUNT_STAR_EXPANSION)],
- distinct,
- filter,
- order_by,
- null_treatment,
- )))
- }
- _ => Transformed::no(old_expr),
- },
-
- ScalarSubquery(Subquery {
- subquery,
- outer_ref_columns,
- }) => subquery
- .as_ref()
- .clone()
- .transform_down(&analyze_internal)?
- .update_data(|new_plan| {
- ScalarSubquery(Subquery {
- subquery: Arc::new(new_plan),
- outer_ref_columns,
- })
- }),
- Expr::InSubquery(InSubquery {
- expr,
- subquery,
- negated,
- }) => subquery
- .subquery
- .as_ref()
- .clone()
- .transform_down(&analyze_internal)?
- .update_data(|new_plan| {
- Expr::InSubquery(InSubquery::new(
- expr,
- Subquery {
- subquery: Arc::new(new_plan),
- outer_ref_columns: subquery.outer_ref_columns,
- },
- negated,
- ))
- }),
- Expr::Exists(expr::Exists { subquery, negated }) => subquery
- .subquery
- .as_ref()
- .clone()
- .transform_down(&analyze_internal)?
- .update_data(|new_plan| {
- Expr::Exists(expr::Exists {
- subquery: Subquery {
- subquery: Arc::new(new_plan),
- outer_ref_columns: subquery.outer_ref_columns,
- },
- negated,
- })
- }),
- _ => Transformed::no(old_expr),
- })
- }
+fn analyze_internal(plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
+ let name_preserver = NamePreserver::new(&plan);
+ plan.map_expressions(|expr| {
Review Comment:
Not only does this logic avoid the clone (`old_expr.clone()`) I also think
it is easier to understand what is happening
--
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]
