Jefffrey commented on code in PR #5518:
URL: https://github.com/apache/arrow-datafusion/pull/5518#discussion_r1135089474
##########
datafusion/expr/src/logical_plan/builder.rs:
##########
@@ -986,6 +994,46 @@ impl LogicalPlanBuilder {
}
}
+//handle Count(Expr:Wildcard) with DataFrame API
+pub fn handle_wildcard(exprs: Vec<Expr>) -> Result<Vec<Expr>> {
+ let exprs: Vec<Expr> = exprs
+ .iter()
+ .map(|expr| {
+ if let Expr::AggregateFunction(AggregateFunction {
+ fun,
+ args,
+ distinct,
+ filter,
+ }) = expr
+ {
+ if let aggregate_function::AggregateFunction::Count = fun {
+ if args.len() == 1 {
+ let arg = args.get(0).unwrap().clone();
+ match arg {
+ Expr::Wildcard => {
+ Expr::AggregateFunction(AggregateFunction {
+ fun: fun.clone(),
+ args:
vec![lit(ScalarValue::UInt8(Some(1)))],
+ distinct: *distinct,
+ filter: filter.clone(),
+ })
+ }
+ _ => expr.clone(),
+ }
+ } else {
+ expr.clone()
+ }
+ } else {
+ expr.clone()
+ }
+ } else {
+ expr.clone()
+ }
+ })
+ .collect();
+ Ok(exprs)
+}
Review Comment:
the extreme nesting is not ideal, i think can simplify to this:
```suggestion
//handle Count(Expr:Wildcard) with DataFrame API
pub fn handle_wildcard(exprs: Vec<Expr>) -> Result<Vec<Expr>> {
let exprs: Vec<Expr> = exprs
.iter()
.map(|expr| match expr {
Expr::AggregateFunction(AggregateFunction {
fun: aggregate_function::AggregateFunction::Count,
args,
distinct,
filter,
}) if args.len() == 1 => match args[0] {
Expr::Wildcard => Expr::AggregateFunction(AggregateFunction {
fun: aggregate_function::AggregateFunction::Count,
// TODO: replace with the constant
args: vec![lit(ScalarValue::UInt8(Some(1)))],
distinct: *distinct,
filter: filter.clone(),
}),
_ => expr.clone(),
},
_ => expr.clone(),
})
.collect();
Ok(exprs)
}
```
unsure if can simplify more, feel free to explore that
also check my TODO comment on the constant (i think i mention it in the
original issue)
--
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]
