Cannot Count(Expr:Wildcard) with DataFrame API

[jiangzhx]

Which issue does this PR close?

Closes #5473

[mingmwang]

I think such logic should be moved to the Analyzer.

[jiangzhx]

thanks @mingmwang , as @Jefffrey mention
SQL substitute function is here:

https://github.com/apache/arrow-datafusion/blob/e9852074bacd8c891d84eba38b3417aa16a2d18c/datafusion/sql/src/expr/function.rs#L193-L218

but use dafaframe will got error

    ctx.table("alltypes_plain")
        .await?
        .aggregate(vec![], vec![count(Expr::Wildcard)])?
        .explain(false, false)?
        .show()
        .await?;

I'm trying another way following @Jefffrey's method

[jiangzhx]

As everyone's suggestions, I made some modifications.
Can @mingmwang and @Jefffrey take a look and help with code review when you have time.
thanks a lot.

[Jefffrey]

the extreme nesting is not ideal, i think can simplify to this:

Suggested change

	//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)
	}
	//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)

[jiangzhx]

I think such logic should be moved to the Analyzer.

I saw @mingmwang's PR #5570, maybe I can migrate the current logic to @mingmwang's solution.
I don't know if my idea is correct.

[jiangzhx]

try adding AnalyzerRule #5627 to do the samething

[Jefffrey]

does look good, though worth checking other opinions about where the logic should reside (whether like this or in analyzer as mentioned)

[alamb]

LGTM -- I also think #5627 looks good to but defer to @mingmwang

[alamb]

👍

[alamb]

Closing in favor of #5518 (so it is clear which we are trying to merge)

Thank you @jiangzhx and @Jefffrey