Skip to content

Commit

Permalink
Allow aggregation without projection in Unparser (#13326)
Browse files Browse the repository at this point in the history
* Allow aggregation without projection in Unparser

* Fix formatting
  • Loading branch information
blaginin authored Nov 13, 2024
1 parent 99cdbcc commit cc96026
Show file tree
Hide file tree
Showing 2 changed files with 50 additions and 2 deletions.
22 changes: 21 additions & 1 deletion datafusion/sql/src/unparser/plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -420,7 +420,27 @@ impl Unparser<'_> {
)
}
LogicalPlan::Aggregate(agg) => {
// Aggregate nodes are handled simultaneously with Projection nodes
// Aggregation can be already handled in the projection case
if !select.already_projected() {
// The query returns aggregate and group expressions. If that weren't the case,
// the aggregate would have been placed inside a projection, making the check above^ false
let exprs: Vec<_> = agg
.aggr_expr
.iter()
.chain(agg.group_expr.iter())
.map(|expr| self.select_item_to_sql(expr))
.collect::<Result<Vec<_>>>()?;
select.projection(exprs);

select.group_by(ast::GroupByExpr::Expressions(
agg.group_expr
.iter()
.map(|expr| self.expr_to_sql(expr))
.collect::<Result<Vec<_>>>()?,
vec![],
));
}

self.select_to_sql_recursively(
agg.input.as_ref(),
query,
Expand Down
30 changes: 29 additions & 1 deletion datafusion/sql/tests/cases/plan_to_sql.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@ use std::vec;

use arrow_schema::*;
use datafusion_common::{DFSchema, Result, TableReference};
use datafusion_expr::test::function_stub::{count_udaf, max_udaf, min_udaf, sum_udaf};
use datafusion_expr::test::function_stub::{
count_udaf, max_udaf, min_udaf, sum, sum_udaf,
};
use datafusion_expr::{col, lit, table_scan, wildcard, LogicalPlanBuilder};
use datafusion_functions::unicode;
use datafusion_functions_aggregate::grouping::grouping_udaf;
Expand Down Expand Up @@ -563,6 +565,32 @@ Projection: unnest_placeholder(unnest_table.struct_col).field1, unnest_placehold
Ok(())
}

#[test]
fn test_aggregation_without_projection() -> Result<()> {
let schema = Schema::new(vec![
Field::new("name", DataType::Utf8, false),
Field::new("age", DataType::UInt8, false),
]);

let plan = LogicalPlanBuilder::from(
table_scan(Some("users"), &schema, Some(vec![0, 1]))?.build()?,
)
.aggregate(vec![col("name")], vec![sum(col("age"))])?
.build()?;

let unparser = Unparser::default();
let statement = unparser.plan_to_sql(&plan)?;

let actual = &statement.to_string();

assert_eq!(
actual,
r#"SELECT sum(users.age), users."name" FROM (SELECT users."name", users.age FROM users) GROUP BY users."name""#
);

Ok(())
}

#[test]
fn test_table_references_in_plan_to_sql() {
fn test(table_name: &str, expected_sql: &str, dialect: &impl UnparserDialect) {
Expand Down

0 comments on commit cc96026

Please sign in to comment.