alamb commented on a change in pull request #1315:
URL: https://github.com/apache/arrow-datafusion/pull/1315#discussion_r752667648
##########
File path: datafusion/src/optimizer/single_distinct_to_groupby.rs
##########
@@ -145,19 +150,25 @@ fn optimize_children(plan: &LogicalPlan) ->
Result<LogicalPlan> {
fn is_single_distinct_agg(plan: &LogicalPlan) -> bool {
match plan {
- LogicalPlan::Aggregate { aggr_expr, .. } => {
- aggr_expr.len() == 1
- && aggr_expr
- .iter()
- .filter(|expr| {
- let mut is_distinct = false;
- if let Expr::AggregateFunction { distinct, .. } = expr
{
- is_distinct = *distinct;
- }
- is_distinct
- })
- .count()
- == 1
+ LogicalPlan::Aggregate {
+ input, aggr_expr, ..
+ } => {
+ let mut fields_set = HashSet::new();
+ aggr_expr
+ .iter()
+ .filter(|expr| {
+ let mut is_distinct = false;
+ if let Expr::AggregateFunction { distinct, args, .. } =
expr {
+ is_distinct = *distinct;
+ args.iter().for_each(|expr| {
+
fields_set.insert(expr.name(input.schema()).unwrap());
+ })
+ }
+ is_distinct
+ })
+ .count()
+ == aggr_expr.len()
+ && fields_set.len() == 1
Review comment:
👍
##########
File path: datafusion/src/optimizer/single_distinct_to_groupby.rs
##########
@@ -251,13 +266,41 @@ mod tests {
)?
.build()?;
+ // Do nothing
let expected = "Aggregate: groupBy=[[#test.a]], aggr=[[COUNT(DISTINCT
#test.b), COUNT(DISTINCT #test.c)]] [a:UInt32, COUNT(DISTINCT test.b):UInt64;N,
COUNT(DISTINCT test.c):UInt64;N]\
\n TableScan: test projection=None [a:UInt32,
b:UInt32, c:UInt32]";
assert_optimized_plan_eq(&plan, expected);
Ok(())
}
+ #[test]
+ fn one_field_two_distinct_and_groupby() -> Result<()> {
+ let table_scan = test_table_scan()?;
+
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .aggregate(
+ vec![col("a")],
+ vec![
+ count_distinct(col("b")),
+ Expr::AggregateFunction {
+ fun: aggregates::AggregateFunction::Max,
+ distinct: true,
+ args: vec![col("b")],
+ },
+ ],
+ )?
+ .build()?;
+ // Should work
+ let expected = "Projection: #test.a AS a, #COUNT(test.b) AS
COUNT(DISTINCT test.b), #MAX(test.b) AS MAX(DISTINCT test.b) [a:UInt32,
COUNT(DISTINCT test.b):UInt64;N, MAX(DISTINCT test.b):UInt32;N]\
Review comment:
nice
--
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]
