This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new cc960265b9 Allow aggregation without projection in `Unparser` (#13326)
cc960265b9 is described below
commit cc960265b9d88abbcc92ce6c0f78df254af67fd8
Author: Dmitrii Blaginin <[email protected]>
AuthorDate: Wed Nov 13 14:25:22 2024 +0000
Allow aggregation without projection in `Unparser` (#13326)
* Allow aggregation without projection in Unparser
* Fix formatting
---
datafusion/sql/src/unparser/plan.rs | 22 +++++++++++++++++++++-
datafusion/sql/tests/cases/plan_to_sql.rs | 30 +++++++++++++++++++++++++++++-
2 files changed, 50 insertions(+), 2 deletions(-)
diff --git a/datafusion/sql/src/unparser/plan.rs
b/datafusion/sql/src/unparser/plan.rs
index c46bc62063..433c456855 100644
--- a/datafusion/sql/src/unparser/plan.rs
+++ b/datafusion/sql/src/unparser/plan.rs
@@ -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,
diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs
b/datafusion/sql/tests/cases/plan_to_sql.rs
index 669f9f06f0..94e420066d 100644
--- a/datafusion/sql/tests/cases/plan_to_sql.rs
+++ b/datafusion/sql/tests/cases/plan_to_sql.rs
@@ -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;
@@ -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) {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
