Skip to content

Commit 75cd4d7

Browse files
fix: preserve subquery structure when unparsing SubqueryAlias over Aggregate
When the SQL unparser encountered a SubqueryAlias node whose direct child was an Aggregate (or other clause-building plan like Window, Sort, Limit, Union), it would flatten the subquery into a simple table alias, losing the aggregate entirely. For example, a plan representing: SELECT j1.col FROM j1 JOIN (SELECT max(id) AS m FROM j2) AS b ON j1.id = b.m would unparse to: SELECT j1.col FROM j1 INNER JOIN j2 AS b ON j1.id = b.m dropping the MAX aggregate and the subquery. Root cause: the SubqueryAlias handler in select_to_sql_recursively would call subquery_alias_inner_query_and_columns (which only unwraps Projection children) and unparse_table_scan_pushdown (which only handles TableScan/SubqueryAlias/Projection). When both returned nothing useful for an Aggregate child, the code recursed directly into the Aggregate, merging its GROUP BY into the outer SELECT instead of wrapping it in a derived subquery. The fix adds an early check: if the SubqueryAlias's direct child is a plan type that builds its own SELECT clauses (Aggregate, Window, Sort, Limit, Union), emit it as a derived subquery via self.derive() with the alias always attached, rather than falling through to the recursive path that would flatten it.
1 parent 3bdcdf5 commit 75cd4d7

2 files changed

Lines changed: 90 additions & 1 deletion

File tree

datafusion/sql/src/unparser/plan.rs

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -828,6 +828,27 @@ impl Unparser<'_> {
828828
Some(plan_alias.alias.clone()),
829829
select.already_projected(),
830830
)?;
831+
832+
// If the SubqueryAlias directly wraps a plan that builds its
833+
// own SELECT clauses (e.g. Aggregate adds GROUP BY, Window adds
834+
// OVER, etc.) and unparse_table_scan_pushdown couldn't reduce it,
835+
// we must emit a derived subquery: (SELECT ...) AS alias.
836+
// Without this, the recursive handler would merge those clauses
837+
// into the outer SELECT, losing the subquery structure entirely.
838+
if unparsed_table_scan.is_none()
839+
&& Self::requires_derived_subquery(plan_alias.input.as_ref())
840+
{
841+
return self.derive(
842+
&plan_alias.input,
843+
relation,
844+
Some(self.new_table_alias(
845+
plan_alias.alias.table().to_string(),
846+
columns,
847+
)),
848+
false,
849+
);
850+
}
851+
831852
// if the child plan is a TableScan with pushdown operations, we don't need to
832853
// create an additional subquery for it
833854
if !select.already_projected() && unparsed_table_scan.is_none() {
@@ -1060,6 +1081,22 @@ impl Unparser<'_> {
10601081
scan.projection.is_some() || !scan.filters.is_empty() || scan.fetch.is_some()
10611082
}
10621083

1084+
/// Returns true if a plan, when used as the direct child of a SubqueryAlias,
1085+
/// must be emitted as a derived subquery `(SELECT ...) AS alias`.
1086+
///
1087+
/// Plans like Aggregate or Window build their own SELECT clauses (GROUP BY,
1088+
/// window functions).
1089+
fn requires_derived_subquery(plan: &LogicalPlan) -> bool {
1090+
matches!(
1091+
plan,
1092+
LogicalPlan::Aggregate(_)
1093+
| LogicalPlan::Window(_)
1094+
| LogicalPlan::Sort(_)
1095+
| LogicalPlan::Limit(_)
1096+
| LogicalPlan::Union(_)
1097+
)
1098+
}
1099+
10631100
/// Try to unparse a table scan with pushdown operations into a new subquery plan.
10641101
/// If the table scan is without any pushdown operations, return None.
10651102
fn unparse_table_scan_pushdown(

datafusion/sql/tests/cases/plan_to_sql.rs

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ use datafusion_common::{
2323
};
2424
use datafusion_expr::expr::{WindowFunction, WindowFunctionParams};
2525
use datafusion_expr::test::function_stub::{
26-
count_udaf, max_udaf, min_udaf, sum, sum_udaf,
26+
count_udaf, max, max_udaf, min_udaf, sum, sum_udaf,
2727
};
2828
use datafusion_expr::{
2929
EmptyRelation, Expr, Extension, LogicalPlan, LogicalPlanBuilder, Union,
@@ -2893,3 +2893,55 @@ fn test_json_access_3() {
28932893
@r#"SELECT (j1.j1_string : 'field.inner1[''inner2'']') FROM j1"#
28942894
);
28952895
}
2896+
2897+
/// Test that unparsing a manually constructed join with a subquery aggregate
2898+
/// preserves the MAX aggregate function.
2899+
///
2900+
/// Builds the equivalent of:
2901+
/// SELECT j1.j1_string FROM j1
2902+
/// JOIN (SELECT max(j2_id) AS max_id FROM j2) AS b
2903+
/// ON j1.j1_id = b.max_id
2904+
#[test]
2905+
fn test_unparse_manual_join_with_subquery_aggregate() -> Result<()> {
2906+
let context = MockContextProvider {
2907+
state: MockSessionState::default(),
2908+
};
2909+
let j1_schema = context
2910+
.get_table_source(TableReference::bare("j1"))?
2911+
.schema();
2912+
let j2_schema = context
2913+
.get_table_source(TableReference::bare("j2"))?
2914+
.schema();
2915+
2916+
// Build the right side: SELECT max(j2_id) AS max_id FROM j2
2917+
let right_scan = table_scan(Some("j2"), &j2_schema, None)?.build()?;
2918+
let right_agg = LogicalPlanBuilder::from(right_scan)
2919+
.aggregate(vec![] as Vec<Expr>, vec![max(col("j2.j2_id")).alias("max_id")])?
2920+
.build()?;
2921+
let right_subquery = subquery_alias(right_agg, "b")?;
2922+
2923+
// Build the full plan: SELECT j1.j1_string FROM j1 JOIN (...) AS b ON j1.j1_id = b.max_id
2924+
let left_scan = table_scan(Some("j1"), &j1_schema, None)?.build()?;
2925+
let plan = LogicalPlanBuilder::from(left_scan)
2926+
.join(
2927+
right_subquery,
2928+
datafusion_expr::JoinType::Inner,
2929+
(
2930+
vec![Column::from_qualified_name("j1.j1_id")],
2931+
vec![Column::from_qualified_name("b.max_id")],
2932+
),
2933+
None,
2934+
)?
2935+
.project(vec![col("j1.j1_string")])?
2936+
.build()?;
2937+
2938+
let unparser = Unparser::default();
2939+
let sql = unparser.plan_to_sql(&plan)?.to_string();
2940+
let sql_upper = sql.to_uppercase();
2941+
assert!(
2942+
sql_upper.contains("MAX("),
2943+
"Unparsed SQL should preserve the MAX aggregate function call, got: {sql}"
2944+
);
2945+
2946+
Ok(())
2947+
}

0 commit comments

Comments
 (0)