From 51e4d1c2ea6cc6fa3189a6864acfae3e26d37dff Mon Sep 17 00:00:00 2001 From: GreatRiver Date: Wed, 8 Jul 2026 15:34:58 +0800 Subject: [PATCH 1/9] fix: rank rollup window results after union --- pkg/sql/plan/build_test.go | 39 ++ pkg/sql/plan/query_builder.go | 533 ++++++++++++++++++++ test/distributed/cases/window/rollup.result | 28 + test/distributed/cases/window/rollup.sql | 25 +- 4 files changed, 624 insertions(+), 1 deletion(-) diff --git a/pkg/sql/plan/build_test.go b/pkg/sql/plan/build_test.go index f7c3b285e39d1..6d90763941670 100644 --- a/pkg/sql/plan/build_test.go +++ b/pkg/sql/plan/build_test.go @@ -825,6 +825,45 @@ func TestSingleTableSQLBuilder(t *testing.T) { runTestShouldError(mock, t, sqls) } +func TestRollupWindowRanksAfterRollupUnion(t *testing.T) { + mock := NewMockOptimizer(false) + sql := ` + select + l_returnflag, + l_linestatus, + sum(l_quantity) as total_qty, + row_number() over (order by sum(l_quantity) desc, l_returnflag, l_linestatus) as row_num, + rank() over (order by sum(l_quantity) desc) as rank_num, + dense_rank() over (order by sum(l_quantity) desc) as dense_rank_num + from + lineitem + group by + l_returnflag, + l_linestatus with rollup + order by total_qty desc, l_returnflag, l_linestatus` + + logicPlan, err := runOneStmt(mock, t, sql) + require.NoError(t, err) + + query := logicPlan.GetQuery() + require.NotNil(t, query) + + windowCount := 0 + windowAfterUnionCount := 0 + for _, node := range query.Nodes { + if node.NodeType == plan.Node_WINDOW { + windowCount++ + require.Len(t, node.Children, 1) + if query.Nodes[node.Children[0]].NodeType == plan.Node_UNION_ALL { + windowAfterUnionCount++ + } + } + } + + require.Equal(t, 3, windowCount) + require.Equal(t, 1, windowAfterUnionCount) +} + func TestOnlyFullGroupByAllowsCorrelatedSubqueryOnGroupedColumn(t *testing.T) { sqls := []string{ ` diff --git a/pkg/sql/plan/query_builder.go b/pkg/sql/plan/query_builder.go index c1619a9eed11b..66aca09cd1904 100644 --- a/pkg/sql/plan/query_builder.go +++ b/pkg/sql/plan/query_builder.go @@ -3065,6 +3065,10 @@ func (builder *QueryBuilder) bindSelect(stmt *tree.Select, ctx *BindContext, isR } } if len(selectClause.GroupBy.GroupByExprsList) > 1 && !selectClause.GroupBy.Apart { + if rewrittenSelect, ok := rewriteRollupWindowSelect(selectClause, astOrderBy, astLimit, astRankOption); ok { + return builder.bindSelect(rewrittenSelect, ctx, isRoot) + } + groupingCount := len(selectClause.GroupBy.GroupByExprsList) selectStmts := make([]*tree.SelectClause, groupingCount) if groupingCount > 1 { @@ -3294,6 +3298,535 @@ func (builder *QueryBuilder) bindSelect(stmt *tree.Select, ctx *BindContext, isR return } +func rewriteRollupWindowSelect( + selectClause *tree.SelectClause, + astOrderBy tree.OrderBy, + astLimit *tree.Limit, + astRankOption *tree.RankOption, +) (*tree.Select, bool) { + if selectClause == nil || selectClause.GroupBy == nil || selectClause.GroupBy.Apart || selectClause.Distinct { + return nil, false + } + if selectClause.Having != nil && rollupWindowExprContainsWindow(selectClause.Having.Expr) { + return nil, false + } + + innerExprs, outerExprs, exprAliases, ok := buildRollupWindowSelectExprs(selectClause.Exprs) + if !ok { + return nil, false + } + + groupingCount := len(selectClause.GroupBy.GroupByExprsList) + if groupingCount <= 1 { + return nil, false + } + if selectClause.Having != nil { + selectClause.Having.RollupHaving = true + } + + selectStmts := make([]*tree.SelectClause, groupingCount) + for i, list := range selectClause.GroupBy.GroupByExprsList { + selectStmts[i] = &tree.SelectClause{ + Distinct: selectClause.Distinct, + Exprs: append(tree.SelectExprs(nil), innerExprs...), + From: selectClause.From, + Where: selectClause.Where, + GroupBy: &tree.GroupByClause{ + GroupByExprsList: selectClause.GroupBy.GroupByExprsList, + GroupingSet: list, + Apart: true, + Cube: false, + Rollup: false, + }, + Having: selectClause.Having, + Option: selectClause.Option, + } + } + + leftClause := &tree.UnionClause{Type: tree.UNION, Left: selectStmts[0], Right: selectStmts[1], All: true} + for i, stmt := range selectStmts { + if i == 0 || i == 1 { + continue + } + leftClause = &tree.UnionClause{Type: tree.UNION, Left: leftClause, Right: stmt, All: true} + } + + derived := &tree.Select{Select: leftClause} + return &tree.Select{ + Select: &tree.SelectClause{ + Exprs: outerExprs, + From: &tree.From{ + Tables: tree.TableExprs{ + &tree.AliasedTableExpr{ + Expr: tree.NewParenTableExpr(derived), + As: tree.AliasClause{ + Alias: tree.Identifier("__mo_rollup_window"), + }, + }, + }, + }, + }, + OrderBy: rewriteRollupWindowOrderBy(astOrderBy, exprAliases), + Limit: astLimit, + RankOption: astRankOption, + }, true +} + +func buildRollupWindowSelectExprs(selectExprs tree.SelectExprs) (tree.SelectExprs, tree.SelectExprs, map[string]string, bool) { + if len(selectExprs) == 0 { + return nil, nil, nil, false + } + + hasWindow := false + innerExprs := make(tree.SelectExprs, 0, len(selectExprs)) + outerExprs := make(tree.SelectExprs, 0, len(selectExprs)) + exprAliases := make(map[string]string) + usedAliases := make(map[string]struct{}) + aliasesByIndex := make([]string, len(selectExprs)) + + for i, selectExpr := range selectExprs { + if rollupWindowExprContainsWindow(selectExpr.Expr) { + hasWindow = true + continue + } + + alias, ok := rollupWindowSelectExprAlias(selectExpr, usedAliases) + if !ok { + return nil, nil, nil, false + } + + aliasesByIndex[i] = alias + innerExprs = append(innerExprs, selectExpr) + addRollupWindowExprAlias(exprAliases, selectExpr.Expr, alias) + addRollupWindowExprAlias(exprAliases, tree.NewUnresolvedColName(alias), alias) + if colName, ok := rollupWindowBareColumnName(selectExpr.Expr); ok { + addRollupWindowExprAlias(exprAliases, tree.NewUnresolvedColName(colName), alias) + } + } + + if !hasWindow || len(innerExprs) == 0 { + return nil, nil, nil, false + } + + for i, selectExpr := range selectExprs { + if rollupWindowExprContainsWindow(selectExpr.Expr) { + rewrittenExpr, ok := rewriteRollupWindowExpr(selectExpr.Expr, exprAliases) + if !ok { + return nil, nil, nil, false + } + outerExprs = append(outerExprs, tree.SelectExpr{ + Expr: rewrittenExpr, + As: selectExpr.As, + }) + continue + } + + outerExprs = append(outerExprs, tree.SelectExpr{ + Expr: tree.NewUnresolvedColName(aliasesByIndex[i]), + As: selectExpr.As, + }) + } + + return innerExprs, outerExprs, exprAliases, true +} + +func rollupWindowSelectExprAlias(selectExpr tree.SelectExpr, usedAliases map[string]struct{}) (string, bool) { + if selectExpr.As != nil && !selectExpr.As.Empty() { + alias := selectExpr.As.Origin() + if useRollupWindowAlias(alias, usedAliases) { + return alias, true + } + return "", false + } + + if alias, ok := rollupWindowBareColumnName(selectExpr.Expr); ok && useRollupWindowAlias(alias, usedAliases) { + return alias, true + } + + return "", false +} + +func useRollupWindowAlias(alias string, usedAliases map[string]struct{}) bool { + if alias == "" { + return false + } + key := strings.ToLower(alias) + if _, ok := usedAliases[key]; ok { + return false + } + usedAliases[key] = struct{}{} + return true +} + +func rollupWindowBareColumnName(expr tree.Expr) (string, bool) { + unresolvedName, ok := expr.(*tree.UnresolvedName) + if !ok || unresolvedName.Star || unresolvedName.NumParts == 0 { + return "", false + } + return unresolvedName.ColNameOrigin(), true +} + +func addRollupWindowExprAlias(exprAliases map[string]string, expr tree.Expr, alias string) { + exprKey := tree.String(expr, dialect.MYSQL) + exprAliases[exprKey] = alias + exprAliases[strings.ToLower(exprKey)] = alias +} + +func lookupRollupWindowExprAlias(exprAliases map[string]string, expr tree.Expr) (string, bool) { + exprKey := tree.String(expr, dialect.MYSQL) + if alias, ok := exprAliases[exprKey]; ok { + return alias, true + } + alias, ok := exprAliases[strings.ToLower(exprKey)] + return alias, ok +} + +func rewriteRollupWindowExpr(expr tree.Expr, exprAliases map[string]string) (tree.Expr, bool) { + if alias, ok := lookupRollupWindowExprAlias(exprAliases, expr); ok { + return tree.NewUnresolvedColName(alias), true + } + + switch e := expr.(type) { + case *tree.FuncExpr: + next := *e + if len(e.Exprs) > 0 { + exprs, ok := rewriteRollupWindowExprs(e.Exprs, exprAliases) + if !ok { + return nil, false + } + next.Exprs = exprs + } + if len(e.OrderBy) > 0 { + orderBy, ok := rewriteRollupWindowOrderByStrict(e.OrderBy, exprAliases) + if !ok { + return nil, false + } + next.OrderBy = orderBy + } + if e.WindowSpec != nil { + windowSpec, ok := rewriteRollupWindowSpec(e.WindowSpec, exprAliases) + if !ok { + return nil, false + } + next.WindowSpec = windowSpec + } + return &next, true + case *tree.BinaryExpr: + left, ok := rewriteRollupWindowExpr(e.Left, exprAliases) + if !ok { + return nil, false + } + right, ok := rewriteRollupWindowExpr(e.Right, exprAliases) + if !ok { + return nil, false + } + next := *e + next.Left = left + next.Right = right + return &next, true + case *tree.UnaryExpr: + child, ok := rewriteRollupWindowExpr(e.Expr, exprAliases) + if !ok { + return nil, false + } + next := *e + next.Expr = child + return &next, true + case *tree.ComparisonExpr: + next := *e + if e.Left != nil { + left, ok := rewriteRollupWindowExpr(e.Left, exprAliases) + if !ok { + return nil, false + } + next.Left = left + } + if e.Right != nil { + right, ok := rewriteRollupWindowExpr(e.Right, exprAliases) + if !ok { + return nil, false + } + next.Right = right + } + if e.Escape != nil { + escape, ok := rewriteRollupWindowExpr(e.Escape, exprAliases) + if !ok { + return nil, false + } + next.Escape = escape + } + return &next, true + case *tree.AndExpr: + left, ok := rewriteRollupWindowExpr(e.Left, exprAliases) + if !ok { + return nil, false + } + right, ok := rewriteRollupWindowExpr(e.Right, exprAliases) + if !ok { + return nil, false + } + next := *e + next.Left = left + next.Right = right + return &next, true + case *tree.XorExpr: + left, ok := rewriteRollupWindowExpr(e.Left, exprAliases) + if !ok { + return nil, false + } + right, ok := rewriteRollupWindowExpr(e.Right, exprAliases) + if !ok { + return nil, false + } + next := *e + next.Left = left + next.Right = right + return &next, true + case *tree.OrExpr: + left, ok := rewriteRollupWindowExpr(e.Left, exprAliases) + if !ok { + return nil, false + } + right, ok := rewriteRollupWindowExpr(e.Right, exprAliases) + if !ok { + return nil, false + } + next := *e + next.Left = left + next.Right = right + return &next, true + case *tree.NotExpr: + child, ok := rewriteRollupWindowExpr(e.Expr, exprAliases) + if !ok { + return nil, false + } + next := *e + next.Expr = child + return &next, true + case *tree.ParenExpr: + child, ok := rewriteRollupWindowExpr(e.Expr, exprAliases) + if !ok { + return nil, false + } + next := *e + next.Expr = child + return &next, true + case *tree.CastExpr: + child, ok := rewriteRollupWindowExpr(e.Expr, exprAliases) + if !ok { + return nil, false + } + next := *e + next.Expr = child + return &next, true + case *tree.Tuple: + exprs, ok := rewriteRollupWindowExprs(e.Exprs, exprAliases) + if !ok { + return nil, false + } + next := *e + next.Exprs = exprs + return &next, true + case *tree.CaseExpr: + next := *e + if e.Expr != nil { + caseExpr, ok := rewriteRollupWindowExpr(e.Expr, exprAliases) + if !ok { + return nil, false + } + next.Expr = caseExpr + } + if len(e.Whens) > 0 { + next.Whens = make([]*tree.When, len(e.Whens)) + for i, when := range e.Whens { + cond, ok := rewriteRollupWindowExpr(when.Cond, exprAliases) + if !ok { + return nil, false + } + val, ok := rewriteRollupWindowExpr(when.Val, exprAliases) + if !ok { + return nil, false + } + next.Whens[i] = &tree.When{Cond: cond, Val: val} + } + } + if e.Else != nil { + elseExpr, ok := rewriteRollupWindowExpr(e.Else, exprAliases) + if !ok { + return nil, false + } + next.Else = elseExpr + } + return &next, true + case *tree.NumVal, *tree.StrVal, *tree.TimeUnitExpr: + return expr, true + default: + return nil, false + } +} + +func rewriteRollupWindowExprs(exprs tree.Exprs, exprAliases map[string]string) (tree.Exprs, bool) { + if len(exprs) == 0 { + return nil, true + } + + rewrittenExprs := make(tree.Exprs, len(exprs)) + for i, expr := range exprs { + rewrittenExpr, ok := rewriteRollupWindowExpr(expr, exprAliases) + if !ok { + return nil, false + } + rewrittenExprs[i] = rewrittenExpr + } + return rewrittenExprs, true +} + +func rewriteRollupWindowSpec(windowSpec *tree.WindowSpec, exprAliases map[string]string) (*tree.WindowSpec, bool) { + next := *windowSpec + if len(windowSpec.PartitionBy) > 0 { + partitionBy, ok := rewriteRollupWindowExprs(windowSpec.PartitionBy, exprAliases) + if !ok { + return nil, false + } + next.PartitionBy = partitionBy + } + if len(windowSpec.OrderBy) > 0 { + orderBy, ok := rewriteRollupWindowOrderByStrict(windowSpec.OrderBy, exprAliases) + if !ok { + return nil, false + } + next.OrderBy = orderBy + } + if windowSpec.Frame != nil { + frame, ok := rewriteRollupWindowFrame(windowSpec.Frame, exprAliases) + if !ok { + return nil, false + } + next.Frame = frame + } + return &next, true +} + +func rewriteRollupWindowFrame(frame *tree.FrameClause, exprAliases map[string]string) (*tree.FrameClause, bool) { + next := *frame + if frame.Start != nil { + start, ok := rewriteRollupWindowFrameBound(frame.Start, exprAliases) + if !ok { + return nil, false + } + next.Start = start + } + if frame.End != nil { + end, ok := rewriteRollupWindowFrameBound(frame.End, exprAliases) + if !ok { + return nil, false + } + next.End = end + } + return &next, true +} + +func rewriteRollupWindowFrameBound(bound *tree.FrameBound, exprAliases map[string]string) (*tree.FrameBound, bool) { + next := *bound + if bound.Expr != nil { + expr, ok := rewriteRollupWindowExpr(bound.Expr, exprAliases) + if !ok { + return nil, false + } + next.Expr = expr + } + return &next, true +} + +func rewriteRollupWindowOrderBy(orderBy tree.OrderBy, exprAliases map[string]string) tree.OrderBy { + if len(orderBy) == 0 { + return nil + } + + rewrittenOrderBy := make(tree.OrderBy, len(orderBy)) + for i, order := range orderBy { + next := *order + if expr, ok := rewriteRollupWindowExpr(order.Expr, exprAliases); ok { + next.Expr = expr + } + rewrittenOrderBy[i] = &next + } + return rewrittenOrderBy +} + +func rewriteRollupWindowOrderByStrict(orderBy tree.OrderBy, exprAliases map[string]string) (tree.OrderBy, bool) { + if len(orderBy) == 0 { + return nil, true + } + + rewrittenOrderBy := make(tree.OrderBy, len(orderBy)) + for i, order := range orderBy { + expr, ok := rewriteRollupWindowExpr(order.Expr, exprAliases) + if !ok { + return nil, false + } + next := *order + next.Expr = expr + rewrittenOrderBy[i] = &next + } + return rewrittenOrderBy, true +} + +func rollupWindowExprContainsWindow(expr tree.Expr) bool { + switch e := expr.(type) { + case *tree.FuncExpr: + if e.WindowSpec != nil { + return true + } + for _, child := range e.Exprs { + if rollupWindowExprContainsWindow(child) { + return true + } + } + for _, order := range e.OrderBy { + if rollupWindowExprContainsWindow(order.Expr) { + return true + } + } + case *tree.BinaryExpr: + return rollupWindowExprContainsWindow(e.Left) || rollupWindowExprContainsWindow(e.Right) + case *tree.UnaryExpr: + return rollupWindowExprContainsWindow(e.Expr) + case *tree.ComparisonExpr: + return rollupWindowExprContainsWindow(e.Left) || + rollupWindowExprContainsWindow(e.Right) || + rollupWindowExprContainsWindow(e.Escape) + case *tree.AndExpr: + return rollupWindowExprContainsWindow(e.Left) || rollupWindowExprContainsWindow(e.Right) + case *tree.XorExpr: + return rollupWindowExprContainsWindow(e.Left) || rollupWindowExprContainsWindow(e.Right) + case *tree.OrExpr: + return rollupWindowExprContainsWindow(e.Left) || rollupWindowExprContainsWindow(e.Right) + case *tree.NotExpr: + return rollupWindowExprContainsWindow(e.Expr) + case *tree.ParenExpr: + return rollupWindowExprContainsWindow(e.Expr) + case *tree.CastExpr: + return rollupWindowExprContainsWindow(e.Expr) + case *tree.Tuple: + for _, child := range e.Exprs { + if rollupWindowExprContainsWindow(child) { + return true + } + } + case *tree.CaseExpr: + if rollupWindowExprContainsWindow(e.Expr) || rollupWindowExprContainsWindow(e.Else) { + return true + } + for _, when := range e.Whens { + if rollupWindowExprContainsWindow(when.Cond) || rollupWindowExprContainsWindow(when.Val) { + return true + } + } + } + return false +} + func (builder *QueryBuilder) bindSelectClause( ctx *BindContext, clause *tree.SelectClause, diff --git a/test/distributed/cases/window/rollup.result b/test/distributed/cases/window/rollup.result index bc23b58dd4c44..22d6c806c03f4 100644 --- a/test/distributed/cases/window/rollup.result +++ b/test/distributed/cases/window/rollup.result @@ -893,4 +893,32 @@ Year: 2021 70000.00 Year: 2022 150000.00 Year: Total 220000.00 drop table sales_mixed_types; +create table rollup_window_sales(region varchar(20), product varchar(20), qty int); +insert into rollup_window_sales values +('east', 'apple', 10), +('east', 'banana', 20), +('west', 'apple', 30), +('west', 'banana', 10); +select +region, +product, +sum(qty) as total_qty, +row_number() over (order by sum(qty) desc, region, product) as row_num, +rank() over (order by sum(qty) desc) as rank_num, +dense_rank() over (order by sum(qty) desc) as dense_rank_num +from +rollup_window_sales +group by +region, +product with rollup +order by total_qty desc, region, product; +region product total_qty row_num rank_num dense_rank_num +null null 70 1 1 1 +west null 40 2 2 2 +east null 30 3 3 3 +west apple 30 4 3 3 +east banana 20 5 5 4 +east apple 10 6 6 5 +west banana 10 7 6 5 +drop table rollup_window_sales; drop database rollup_test; diff --git a/test/distributed/cases/window/rollup.sql b/test/distributed/cases/window/rollup.sql index e58bc146cec9f..2f023abb3be18 100644 --- a/test/distributed/cases/window/rollup.sql +++ b/test/distributed/cases/window/rollup.sql @@ -592,4 +592,27 @@ group by order by year_label; drop table sales_mixed_types; -drop database rollup_test; \ No newline at end of file + +-- rollup with window functions should rank over the full rollup result +create table rollup_window_sales(region varchar(20), product varchar(20), qty int); +insert into rollup_window_sales values +('east', 'apple', 10), +('east', 'banana', 20), +('west', 'apple', 30), +('west', 'banana', 10); +select + region, + product, + sum(qty) as total_qty, + row_number() over (order by sum(qty) desc, region, product) as row_num, + rank() over (order by sum(qty) desc) as rank_num, + dense_rank() over (order by sum(qty) desc) as dense_rank_num +from + rollup_window_sales +group by + region, + product with rollup +order by total_qty desc, region, product; +drop table rollup_window_sales; + +drop database rollup_test; From f84c64cb5cfc0f76c2b93d44b106c78a26c99341 Mon Sep 17 00:00:00 2001 From: GreatRiver Date: Wed, 8 Jul 2026 19:11:27 +0800 Subject: [PATCH 2/9] test: cover rollup window rewrite helpers --- pkg/sql/plan/build_test.go | 197 +++++++++++++++++++++++++++++++++++++ 1 file changed, 197 insertions(+) diff --git a/pkg/sql/plan/build_test.go b/pkg/sql/plan/build_test.go index 6d90763941670..1ad4ff21f9173 100644 --- a/pkg/sql/plan/build_test.go +++ b/pkg/sql/plan/build_test.go @@ -31,6 +31,7 @@ import ( moruntime "github.com/matrixorigin/matrixone/pkg/common/runtime" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/pb/plan" + "github.com/matrixorigin/matrixone/pkg/sql/parsers/dialect" "github.com/matrixorigin/matrixone/pkg/sql/parsers/dialect/mysql" "github.com/matrixorigin/matrixone/pkg/sql/parsers/tree" "github.com/matrixorigin/matrixone/pkg/testutil" @@ -864,6 +865,202 @@ func TestRollupWindowRanksAfterRollupUnion(t *testing.T) { require.Equal(t, 1, windowAfterUnionCount) } +func TestRewriteRollupWindowSelectHelpers(t *testing.T) { + stmt := mustParseRollupWindowSelect(t, ` + select + l_returnflag as flag, + l_linestatus as status, + sum(l_quantity) as total_qty, + row_number() over (order by sum(l_quantity) desc) as rn + from lineitem + group by l_returnflag, l_linestatus with rollup + having sum(l_quantity) > 0 + order by sum(l_quantity) desc + limit 10`) + clause := stmt.Select.(*tree.SelectClause) + expandRollupGroupByForTest(clause.GroupBy) + + rewritten, ok := rewriteRollupWindowSelect(clause, stmt.OrderBy, stmt.Limit, stmt.RankOption) + require.True(t, ok) + require.NotNil(t, rewritten) + require.True(t, clause.Having.RollupHaving) + require.NotNil(t, rewritten.Limit) + require.Len(t, rewritten.OrderBy, 1) + require.Contains(t, tree.String(rewritten, dialect.MYSQL), "__mo_rollup_window") + require.Contains(t, tree.String(rewritten, dialect.MYSQL), "total_qty") + + innerExprs, outerExprs, exprAliases, ok := buildRollupWindowSelectExprs(clause.Exprs) + require.True(t, ok) + require.Len(t, innerExprs, 3) + require.Len(t, outerExprs, 4) + + alias, ok := lookupRollupWindowExprAlias(exprAliases, tree.NewUnresolvedColName("l_returnflag")) + require.True(t, ok) + require.Equal(t, "flag", alias) + alias, ok = lookupRollupWindowExprAlias(exprAliases, tree.NewUnresolvedColName("total_qty")) + require.True(t, ok) + require.Equal(t, "total_qty", alias) +} + +func TestRewriteRollupWindowSelectGuards(t *testing.T) { + _, ok := rewriteRollupWindowSelect(nil, nil, nil, nil) + require.False(t, ok) + + noGroup := mustParseRollupWindowSelect(t, "select a, row_number() over () from t") + _, ok = rewriteRollupWindowSelect(noGroup.Select.(*tree.SelectClause), noGroup.OrderBy, noGroup.Limit, noGroup.RankOption) + require.False(t, ok) + + distinct := mustParseRollupWindowSelect(t, "select distinct a, row_number() over () from t group by a, b with rollup") + _, ok = rewriteRollupWindowSelect(distinct.Select.(*tree.SelectClause), distinct.OrderBy, distinct.Limit, distinct.RankOption) + require.False(t, ok) + + havingWindow := mustParseRollupWindowSelect(t, "select a, b, row_number() over () from t group by a, b with rollup having row_number() over () > 0") + _, ok = rewriteRollupWindowSelect(havingWindow.Select.(*tree.SelectClause), havingWindow.OrderBy, havingWindow.Limit, havingWindow.RankOption) + require.False(t, ok) + + oneGroup := mustParseRollupWindowSelect(t, "select a, row_number() over () from t group by a") + _, ok = rewriteRollupWindowSelect(oneGroup.Select.(*tree.SelectClause), oneGroup.OrderBy, oneGroup.Limit, oneGroup.RankOption) + require.False(t, ok) + + emptyExprs := &tree.SelectClause{Exprs: nil} + _, _, _, ok = buildRollupWindowSelectExprs(emptyExprs.Exprs) + require.False(t, ok) + + noWindow := mustParseRollupWindowSelect(t, "select a from t") + _, _, _, ok = buildRollupWindowSelectExprs(noWindow.Select.(*tree.SelectClause).Exprs) + require.False(t, ok) + + unsupportedInnerExpr := mustParseRollupWindowSelect(t, "select a + b, row_number() over () from t") + _, _, _, ok = buildRollupWindowSelectExprs(unsupportedInnerExpr.Select.(*tree.SelectClause).Exprs) + require.False(t, ok) + + duplicateAlias := mustParseRollupWindowSelect(t, "select a as x, b as x, row_number() over () from t") + _, _, _, ok = buildRollupWindowSelectExprs(duplicateAlias.Select.(*tree.SelectClause).Exprs) + require.False(t, ok) +} + +func TestRewriteRollupWindowExprSupportedShapes(t *testing.T) { + exprAliases := rollupWindowAliasMap( + "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", + ) + + complexExpr := mustParseRollupWindowExpr(t, ` + case a + when 1 then not ( + ((b + -c) > cast(d as signed)) and (e like f escape g) + xor h = i + or (j, k) = (l, m) + ) + else 'fallback' + end`) + rewritten, ok := rewriteRollupWindowExpr(complexExpr, exprAliases) + require.True(t, ok) + rewrittenSQL := tree.String(rewritten, dialect.MYSQL) + require.Contains(t, rewrittenSQL, "a_alias") + require.Contains(t, rewrittenSQL, "m_alias") + + windowExpr := mustParseRollupWindowExpr(t, ` + sum(a) over ( + partition by b + order by c + rows between 1 preceding and 2 following + )`) + rewritten, ok = rewriteRollupWindowExpr(windowExpr, exprAliases) + require.True(t, ok) + rewrittenSQL = tree.String(rewritten, dialect.MYSQL) + require.Contains(t, rewrittenSQL, "a_alias") + require.Contains(t, rewrittenSQL, "partition by b_alias") + require.Contains(t, rewrittenSQL, "order by c_alias") + + funcOrderExpr := mustParseRollupWindowExpr(t, "group_concat(a order by b)") + rewritten, ok = rewriteRollupWindowExpr(funcOrderExpr, exprAliases) + require.True(t, ok) + require.Contains(t, tree.String(rewritten, dialect.MYSQL), "b_alias") + + tupleExprs, ok := rewriteRollupWindowExprs(tree.Exprs{tree.NewUnresolvedColName("a"), tree.NewUnresolvedColName("b")}, exprAliases) + require.True(t, ok) + require.Len(t, tupleExprs, 2) + + emptyExprs, ok := rewriteRollupWindowExprs(nil, exprAliases) + require.True(t, ok) + require.Nil(t, emptyExprs) + + _, ok = rewriteRollupWindowExpr(tree.NewUnresolvedColName("missing"), exprAliases) + require.False(t, ok) + + _, ok = rewriteRollupWindowOrderByStrict(tree.OrderBy{{Expr: tree.NewUnresolvedColName("missing")}}, exprAliases) + require.False(t, ok) + + orderBy := rewriteRollupWindowOrderBy(tree.OrderBy{{Expr: tree.NewUnresolvedColName("missing")}}, exprAliases) + require.Len(t, orderBy, 1) + require.Equal(t, "missing", tree.String(orderBy[0].Expr, dialect.MYSQL)) +} + +func TestRollupWindowExprContainsWindowRecursiveShapes(t *testing.T) { + for _, tc := range []struct { + name string + expr string + }{ + {name: "func args", expr: "coalesce(sum(a) over (), b)"}, + {name: "func order by", expr: "group_concat(a order by sum(b) over ())"}, + {name: "binary", expr: "a + sum(b) over ()"}, + {name: "unary", expr: "-sum(a) over ()"}, + {name: "comparison", expr: "sum(a) over () > 0"}, + {name: "and", expr: "sum(a) over () and b"}, + {name: "xor", expr: "a xor sum(b) over ()"}, + {name: "or", expr: "a or sum(b) over ()"}, + {name: "not", expr: "not sum(a) over ()"}, + {name: "paren", expr: "(sum(a) over ())"}, + {name: "cast", expr: "cast(sum(a) over () as signed)"}, + {name: "tuple", expr: "(a, sum(b) over ())"}, + {name: "case expr", expr: "case sum(a) over () when 1 then b else c end"}, + {name: "case when", expr: "case when a then sum(b) over () else c end"}, + } { + t.Run(tc.name, func(t *testing.T) { + require.True(t, rollupWindowExprContainsWindow(mustParseRollupWindowExpr(t, tc.expr))) + }) + } + + require.False(t, rollupWindowExprContainsWindow(mustParseRollupWindowExpr(t, "case when a then b else c end"))) +} + +func mustParseRollupWindowSelect(t *testing.T, sql string) *tree.Select { + t.Helper() + stmts, err := mysql.Parse(context.Background(), sql, 1) + require.NoError(t, err) + require.Len(t, stmts, 1) + stmt, ok := stmts[0].(*tree.Select) + require.True(t, ok) + return stmt +} + +func mustParseRollupWindowExpr(t *testing.T, expr string) tree.Expr { + t.Helper() + stmt := mustParseRollupWindowSelect(t, "select "+expr) + clause, ok := stmt.Select.(*tree.SelectClause) + require.True(t, ok) + require.Len(t, clause.Exprs, 1) + return clause.Exprs[0].Expr +} + +func rollupWindowAliasMap(cols ...string) map[string]string { + exprAliases := make(map[string]string) + for _, col := range cols { + addRollupWindowExprAlias(exprAliases, tree.NewUnresolvedColName(col), col+"_alias") + } + return exprAliases +} + +func expandRollupGroupByForTest(groupBy *tree.GroupByClause) { + if groupBy == nil || !groupBy.Rollup || len(groupBy.GroupByExprsList) == 0 { + return + } + for i := len(groupBy.GroupByExprsList[0]) - 1; i > 0; i-- { + groupBy.GroupByExprsList = append(groupBy.GroupByExprsList, groupBy.GroupByExprsList[0][0:i]) + } + groupBy.GroupByExprsList = append(groupBy.GroupByExprsList, nil) +} + func TestOnlyFullGroupByAllowsCorrelatedSubqueryOnGroupedColumn(t *testing.T) { sqls := []string{ ` From c4cf8a5eef4516bc694403bcf71679cd2953354c Mon Sep 17 00:00:00 2001 From: GreatRiver Date: Fri, 10 Jul 2026 17:42:46 +0800 Subject: [PATCH 3/9] fix: cover rollup window hidden dependencies --- pkg/sql/plan/build_test.go | 253 ++++++++---- pkg/sql/plan/query_builder.go | 434 ++++++++++++++------ test/distributed/cases/window/rollup.result | 41 ++ test/distributed/cases/window/rollup.sql | 29 ++ 4 files changed, 552 insertions(+), 205 deletions(-) diff --git a/pkg/sql/plan/build_test.go b/pkg/sql/plan/build_test.go index 1ad4ff21f9173..8890a1019abd9 100644 --- a/pkg/sql/plan/build_test.go +++ b/pkg/sql/plan/build_test.go @@ -828,41 +828,92 @@ func TestSingleTableSQLBuilder(t *testing.T) { func TestRollupWindowRanksAfterRollupUnion(t *testing.T) { mock := NewMockOptimizer(false) - sql := ` - select - l_returnflag, - l_linestatus, - sum(l_quantity) as total_qty, - row_number() over (order by sum(l_quantity) desc, l_returnflag, l_linestatus) as row_num, - rank() over (order by sum(l_quantity) desc) as rank_num, - dense_rank() over (order by sum(l_quantity) desc) as dense_rank_num - from - lineitem - group by - l_returnflag, - l_linestatus with rollup - order by total_qty desc, l_returnflag, l_linestatus` - - logicPlan, err := runOneStmt(mock, t, sql) - require.NoError(t, err) - - query := logicPlan.GetQuery() - require.NotNil(t, query) + for _, tc := range []struct { + name string + sql string + expectedHeadings []string + }{ + { + name: "aliased aggregate output", + sql: ` + select + l_returnflag, + l_linestatus, + sum(l_quantity) as total_qty, + row_number() over (order by sum(l_quantity) desc, l_returnflag, l_linestatus) as row_num, + rank() over (order by sum(l_quantity) desc) as rank_num, + dense_rank() over (order by sum(l_quantity) desc) as dense_rank_num + from lineitem + group by l_returnflag, l_linestatus with rollup + having total_qty > 0 + order by total_qty desc, l_returnflag, l_linestatus`, + expectedHeadings: []string{"l_returnflag", "l_linestatus", "total_qty", "row_num", "rank_num", "dense_rank_num"}, + }, + { + name: "aggregate output without alias", + sql: ` + select + l_returnflag, + l_linestatus, + sum(l_quantity), + row_number() over (order by sum(l_quantity) desc) as row_num, + rank() over (order by sum(l_quantity) desc) as rank_num, + dense_rank() over (order by sum(l_quantity) desc) as dense_rank_num + from lineitem + group by l_returnflag, l_linestatus with rollup + order by sum(l_quantity) desc, l_returnflag, l_linestatus`, + expectedHeadings: []string{"l_returnflag", "l_linestatus", "sum(l_quantity)", "row_num", "rank_num", "dense_rank_num"}, + }, + { + name: "aggregate used only by windows", + sql: ` + select + l_returnflag, + l_linestatus, + row_number() over (order by sum(l_quantity) desc) as row_num, + rank() over (order by sum(l_quantity) desc) as rank_num, + dense_rank() over (order by sum(l_quantity) desc) as dense_rank_num + from lineitem + group by l_returnflag, l_linestatus with rollup + order by row_num`, + expectedHeadings: []string{"l_returnflag", "l_linestatus", "row_num", "rank_num", "dense_rank_num"}, + }, + { + name: "window outputs only", + sql: ` + select + row_number() over (order by 1) as row_num, + rank() over (order by 1) as rank_num, + dense_rank() over (order by 1) as dense_rank_num + from lineitem + group by l_returnflag with rollup`, + expectedHeadings: []string{"row_num", "rank_num", "dense_rank_num"}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + logicPlan, err := runOneStmt(mock, t, tc.sql) + require.NoError(t, err) - windowCount := 0 - windowAfterUnionCount := 0 - for _, node := range query.Nodes { - if node.NodeType == plan.Node_WINDOW { - windowCount++ - require.Len(t, node.Children, 1) - if query.Nodes[node.Children[0]].NodeType == plan.Node_UNION_ALL { - windowAfterUnionCount++ + query := logicPlan.GetQuery() + require.NotNil(t, query) + require.Equal(t, tc.expectedHeadings, query.Headings) + + windowCount := 0 + windowAfterUnionCount := 0 + for _, node := range query.Nodes { + if node.NodeType == plan.Node_WINDOW { + windowCount++ + require.Len(t, node.Children, 1) + if query.Nodes[node.Children[0]].NodeType == plan.Node_UNION_ALL { + windowAfterUnionCount++ + } + } } - } - } - require.Equal(t, 3, windowCount) - require.Equal(t, 1, windowAfterUnionCount) + require.Equal(t, 3, windowCount) + require.Equal(t, 1, windowAfterUnionCount) + }) + } } func TestRewriteRollupWindowSelectHelpers(t *testing.T) { @@ -883,23 +934,33 @@ func TestRewriteRollupWindowSelectHelpers(t *testing.T) { rewritten, ok := rewriteRollupWindowSelect(clause, stmt.OrderBy, stmt.Limit, stmt.RankOption) require.True(t, ok) require.NotNil(t, rewritten) - require.True(t, clause.Having.RollupHaving) require.NotNil(t, rewritten.Limit) require.Len(t, rewritten.OrderBy, 1) require.Contains(t, tree.String(rewritten, dialect.MYSQL), "__mo_rollup_window") require.Contains(t, tree.String(rewritten, dialect.MYSQL), "total_qty") - innerExprs, outerExprs, exprAliases, ok := buildRollupWindowSelectExprs(clause.Exprs) + state := newRollupWindowRewriteState(clause.Exprs) + outerExprs, ok := buildRollupWindowSelectExprs(clause.Exprs, state) require.True(t, ok) - require.Len(t, innerExprs, 3) + require.Len(t, state.innerExprs, 3) require.Len(t, outerExprs, 4) + for _, innerExpr := range state.innerExprs { + require.NotNil(t, innerExpr.As) + require.True(t, strings.HasPrefix(innerExpr.As.Origin(), rollupWindowInternalAliasPrefix)) + } - alias, ok := lookupRollupWindowExprAlias(exprAliases, tree.NewUnresolvedColName("l_returnflag")) + alias, ok := state.lookupExprAlias(tree.NewUnresolvedColName("l_returnflag")) require.True(t, ok) - require.Equal(t, "flag", alias) - alias, ok = lookupRollupWindowExprAlias(exprAliases, tree.NewUnresolvedColName("total_qty")) + require.True(t, strings.HasPrefix(alias, rollupWindowInternalAliasPrefix)) + flagAlias, ok := state.lookupExprAlias(tree.NewUnresolvedColName("flag")) require.True(t, ok) - require.Equal(t, "total_qty", alias) + require.Equal(t, alias, flagAlias) + upperFlagAlias, ok := state.lookupExprAlias(tree.NewUnresolvedColName("FLAG")) + require.True(t, ok) + require.Equal(t, alias, upperFlagAlias) + alias, ok = state.lookupExprAlias(tree.NewUnresolvedColName("total_qty")) + require.True(t, ok) + require.True(t, strings.HasPrefix(alias, rollupWindowInternalAliasPrefix)) } func TestRewriteRollupWindowSelectGuards(t *testing.T) { @@ -908,39 +969,86 @@ func TestRewriteRollupWindowSelectGuards(t *testing.T) { noGroup := mustParseRollupWindowSelect(t, "select a, row_number() over () from t") _, ok = rewriteRollupWindowSelect(noGroup.Select.(*tree.SelectClause), noGroup.OrderBy, noGroup.Limit, noGroup.RankOption) - require.False(t, ok) + require.True(t, ok) distinct := mustParseRollupWindowSelect(t, "select distinct a, row_number() over () from t group by a, b with rollup") - _, ok = rewriteRollupWindowSelect(distinct.Select.(*tree.SelectClause), distinct.OrderBy, distinct.Limit, distinct.RankOption) - require.False(t, ok) + distinctClause := distinct.Select.(*tree.SelectClause) + expandRollupGroupByForTest(distinctClause.GroupBy) + rewritten, ok := rewriteRollupWindowSelect(distinctClause, distinct.OrderBy, distinct.Limit, distinct.RankOption) + require.True(t, ok) + require.Nil(t, rewritten) havingWindow := mustParseRollupWindowSelect(t, "select a, b, row_number() over () from t group by a, b with rollup having row_number() over () > 0") - _, ok = rewriteRollupWindowSelect(havingWindow.Select.(*tree.SelectClause), havingWindow.OrderBy, havingWindow.Limit, havingWindow.RankOption) - require.False(t, ok) + havingWindowClause := havingWindow.Select.(*tree.SelectClause) + expandRollupGroupByForTest(havingWindowClause.GroupBy) + rewritten, ok = rewriteRollupWindowSelect(havingWindowClause, havingWindow.OrderBy, havingWindow.Limit, havingWindow.RankOption) + require.True(t, ok) + require.Nil(t, rewritten) oneGroup := mustParseRollupWindowSelect(t, "select a, row_number() over () from t group by a") _, ok = rewriteRollupWindowSelect(oneGroup.Select.(*tree.SelectClause), oneGroup.OrderBy, oneGroup.Limit, oneGroup.RankOption) - require.False(t, ok) + require.True(t, ok) - emptyExprs := &tree.SelectClause{Exprs: nil} - _, _, _, ok = buildRollupWindowSelectExprs(emptyExprs.Exprs) + state := newRollupWindowRewriteState(nil) + _, ok = buildRollupWindowSelectExprs(nil, state) + require.False(t, ok) + star := mustParseRollupWindowSelect(t, "select *, row_number() over () from t") + starExprs := star.Select.(*tree.SelectClause).Exprs + state = newRollupWindowRewriteState(starExprs) + _, ok = buildRollupWindowSelectExprs(starExprs, state) require.False(t, ok) noWindow := mustParseRollupWindowSelect(t, "select a from t") - _, _, _, ok = buildRollupWindowSelectExprs(noWindow.Select.(*tree.SelectClause).Exprs) + _, ok = rewriteRollupWindowSelect(noWindow.Select.(*tree.SelectClause), noWindow.OrderBy, noWindow.Limit, noWindow.RankOption) require.False(t, ok) - unsupportedInnerExpr := mustParseRollupWindowSelect(t, "select a + b, row_number() over () from t") - _, _, _, ok = buildRollupWindowSelectExprs(unsupportedInnerExpr.Select.(*tree.SelectClause).Exprs) - require.False(t, ok) + complexInnerExpr := mustParseRollupWindowSelect(t, "select a + b, row_number() over () from t") + complexExprs := complexInnerExpr.Select.(*tree.SelectClause).Exprs + state = newRollupWindowRewriteState(complexExprs) + _, ok = buildRollupWindowSelectExprs(complexExprs, state) + require.True(t, ok) + require.Len(t, state.innerExprs, 1) + + windowBeforeAlias := mustParseRollupWindowSelect(t, "select row_number() over (order by total) as rn, sum(a) as total from t") + windowBeforeAliasExprs := windowBeforeAlias.Select.(*tree.SelectClause).Exprs + state = newRollupWindowRewriteState(windowBeforeAliasExprs) + outerExprs, ok := buildRollupWindowSelectExprs(windowBeforeAliasExprs, state) + require.True(t, ok) + require.Len(t, state.innerExprs, 1) + require.Contains(t, tree.String(outerExprs[0].Expr, dialect.MYSQL), state.innerExprs[0].As.Origin()) duplicateAlias := mustParseRollupWindowSelect(t, "select a as x, b as x, row_number() over () from t") - _, _, _, ok = buildRollupWindowSelectExprs(duplicateAlias.Select.(*tree.SelectClause).Exprs) - require.False(t, ok) + duplicateExprs := duplicateAlias.Select.(*tree.SelectClause).Exprs + state = newRollupWindowRewriteState(duplicateExprs) + _, ok = buildRollupWindowSelectExprs(duplicateExprs, state) + require.True(t, ok) + require.Len(t, state.innerExprs, 2) + + state = newRollupWindowRewriteState(nil) + upperLiteralAlias, ok := state.ensureInnerExpr(mustParseRollupWindowExpr(t, "'A'")) + require.True(t, ok) + lowerLiteralAlias, ok := state.ensureInnerExpr(mustParseRollupWindowExpr(t, "'a'")) + require.True(t, ok) + require.NotEqual(t, upperLiteralAlias, lowerLiteralAlias) + + parenAggregate := mustParseRollupWindowSelect(t, "select (sum(a)), row_number() over () from t") + require.Equal(t, "sum(a)", rollupWindowOutputAlias(parenAggregate.Select.(*tree.SelectClause).Exprs[0]).Origin()) +} + +func TestRewriteRollupWindowUnsupportedDoesNotFallback(t *testing.T) { + mock := NewMockOptimizer(false) + for _, sql := range []string{ + "select distinct a, row_number() over () from nation group by a, n_regionkey with rollup", + "select a, row_number() over () is null from nation group by a, n_regionkey with rollup", + } { + _, err := runOneStmt(mock, t, sql) + require.Error(t, err) + require.Contains(t, err.Error(), "window functions with ROLLUP or CUBE for this expression") + } } func TestRewriteRollupWindowExprSupportedShapes(t *testing.T) { - exprAliases := rollupWindowAliasMap( + state := rollupWindowAliasMap( "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", ) @@ -953,7 +1061,7 @@ func TestRewriteRollupWindowExprSupportedShapes(t *testing.T) { ) else 'fallback' end`) - rewritten, ok := rewriteRollupWindowExpr(complexExpr, exprAliases) + rewritten, ok := rewriteRollupWindowExpr(complexExpr, state) require.True(t, ok) rewrittenSQL := tree.String(rewritten, dialect.MYSQL) require.Contains(t, rewrittenSQL, "a_alias") @@ -965,35 +1073,31 @@ func TestRewriteRollupWindowExprSupportedShapes(t *testing.T) { order by c rows between 1 preceding and 2 following )`) - rewritten, ok = rewriteRollupWindowExpr(windowExpr, exprAliases) + rewritten, ok = rewriteRollupWindowExpr(windowExpr, state) require.True(t, ok) rewrittenSQL = tree.String(rewritten, dialect.MYSQL) require.Contains(t, rewrittenSQL, "a_alias") require.Contains(t, rewrittenSQL, "partition by b_alias") require.Contains(t, rewrittenSQL, "order by c_alias") - funcOrderExpr := mustParseRollupWindowExpr(t, "group_concat(a order by b)") - rewritten, ok = rewriteRollupWindowExpr(funcOrderExpr, exprAliases) + aggregateExpr := mustParseRollupWindowExpr(t, "sum(missing)") + rewritten, ok = rewriteRollupWindowExpr(aggregateExpr, state) require.True(t, ok) - require.Contains(t, tree.String(rewritten, dialect.MYSQL), "b_alias") + require.True(t, strings.HasPrefix(tree.String(rewritten, dialect.MYSQL), rollupWindowInternalAliasPrefix)) + require.Len(t, state.innerExprs, 1) - tupleExprs, ok := rewriteRollupWindowExprs(tree.Exprs{tree.NewUnresolvedColName("a"), tree.NewUnresolvedColName("b")}, exprAliases) + tupleExprs, ok := rewriteRollupWindowExprs(tree.Exprs{tree.NewUnresolvedColName("a"), tree.NewUnresolvedColName("b")}, state) require.True(t, ok) require.Len(t, tupleExprs, 2) - emptyExprs, ok := rewriteRollupWindowExprs(nil, exprAliases) + emptyExprs, ok := rewriteRollupWindowExprs(nil, state) require.True(t, ok) require.Nil(t, emptyExprs) - _, ok = rewriteRollupWindowExpr(tree.NewUnresolvedColName("missing"), exprAliases) - require.False(t, ok) - - _, ok = rewriteRollupWindowOrderByStrict(tree.OrderBy{{Expr: tree.NewUnresolvedColName("missing")}}, exprAliases) - require.False(t, ok) - - orderBy := rewriteRollupWindowOrderBy(tree.OrderBy{{Expr: tree.NewUnresolvedColName("missing")}}, exprAliases) + orderBy, ok := rewriteRollupWindowOrderBy(tree.OrderBy{{Expr: tree.NewUnresolvedColName("another_missing")}}, state) + require.True(t, ok) require.Len(t, orderBy, 1) - require.Equal(t, "missing", tree.String(orderBy[0].Expr, dialect.MYSQL)) + require.True(t, strings.HasPrefix(tree.String(orderBy[0].Expr, dialect.MYSQL), rollupWindowInternalAliasPrefix)) } func TestRollupWindowExprContainsWindowRecursiveShapes(t *testing.T) { @@ -1010,9 +1114,12 @@ func TestRollupWindowExprContainsWindowRecursiveShapes(t *testing.T) { {name: "xor", expr: "a xor sum(b) over ()"}, {name: "or", expr: "a or sum(b) over ()"}, {name: "not", expr: "not sum(a) over ()"}, + {name: "is null", expr: "sum(a) over () is null"}, + {name: "is not null", expr: "sum(a) over () is not null"}, {name: "paren", expr: "(sum(a) over ())"}, {name: "cast", expr: "cast(sum(a) over () as signed)"}, {name: "tuple", expr: "(a, sum(b) over ())"}, + {name: "between", expr: "sum(a) over () between 1 and 2"}, {name: "case expr", expr: "case sum(a) over () when 1 then b else c end"}, {name: "case when", expr: "case when a then sum(b) over () else c end"}, } { @@ -1043,12 +1150,12 @@ func mustParseRollupWindowExpr(t *testing.T, expr string) tree.Expr { return clause.Exprs[0].Expr } -func rollupWindowAliasMap(cols ...string) map[string]string { - exprAliases := make(map[string]string) +func rollupWindowAliasMap(cols ...string) *rollupWindowRewriteState { + state := newRollupWindowRewriteState(nil) for _, col := range cols { - addRollupWindowExprAlias(exprAliases, tree.NewUnresolvedColName(col), col+"_alias") + state.addExprAlias(tree.NewUnresolvedColName(col), col+"_alias") } - return exprAliases + return state } func expandRollupGroupByForTest(groupBy *tree.GroupByClause) { diff --git a/pkg/sql/plan/query_builder.go b/pkg/sql/plan/query_builder.go index 66aca09cd1904..315ecf03169b1 100644 --- a/pkg/sql/plan/query_builder.go +++ b/pkg/sql/plan/query_builder.go @@ -3065,7 +3065,10 @@ func (builder *QueryBuilder) bindSelect(stmt *tree.Select, ctx *BindContext, isR } } if len(selectClause.GroupBy.GroupByExprsList) > 1 && !selectClause.GroupBy.Apart { - if rewrittenSelect, ok := rewriteRollupWindowSelect(selectClause, astOrderBy, astLimit, astRankOption); ok { + if rewrittenSelect, hasWindow := rewriteRollupWindowSelect(selectClause, astOrderBy, astLimit, astRankOption); hasWindow { + if rewrittenSelect == nil { + return 0, moerr.NewNotSupported(builder.GetContext(), "window functions with ROLLUP or CUBE for this expression") + } return builder.bindSelect(rewrittenSelect, ctx, isRoot) } @@ -3304,33 +3307,65 @@ func rewriteRollupWindowSelect( astLimit *tree.Limit, astRankOption *tree.RankOption, ) (*tree.Select, bool) { - if selectClause == nil || selectClause.GroupBy == nil || selectClause.GroupBy.Apart || selectClause.Distinct { - return nil, false - } - if selectClause.Having != nil && rollupWindowExprContainsWindow(selectClause.Having.Expr) { + if selectClause == nil { return nil, false } - innerExprs, outerExprs, exprAliases, ok := buildRollupWindowSelectExprs(selectClause.Exprs) - if !ok { + hasWindow := false + for _, selectExpr := range selectClause.Exprs { + hasWindow = hasWindow || rollupWindowExprContainsWindow(selectExpr.Expr) + } + for _, order := range astOrderBy { + hasWindow = hasWindow || rollupWindowExprContainsWindow(order.Expr) + } + if selectClause.Having != nil { + hasWindow = hasWindow || rollupWindowExprContainsWindow(selectClause.Having.Expr) + } + if !hasWindow { return nil, false } + if selectClause.GroupBy == nil || selectClause.GroupBy.Apart || + len(selectClause.GroupBy.GroupByExprsList) <= 1 || + selectClause.Distinct || + (selectClause.Having != nil && rollupWindowExprContainsWindow(selectClause.Having.Expr)) { + return nil, true + } - groupingCount := len(selectClause.GroupBy.GroupByExprsList) - if groupingCount <= 1 { - return nil, false + rewriteState := newRollupWindowRewriteState(selectClause.Exprs) + outerExprs, ok := buildRollupWindowSelectExprs(selectClause.Exprs, rewriteState) + if !ok { + return nil, true } + var rewrittenHaving *tree.Where if selectClause.Having != nil { - selectClause.Having.RollupHaving = true + havingExpr, ok := rewriteRollupWindowExpr(selectClause.Having.Expr, rewriteState) + if !ok { + return nil, true + } + rewrittenHaving = &tree.Where{ + Type: selectClause.Having.Type, + RollupHaving: true, + Expr: havingExpr, + } + } + rewrittenOrderBy, ok := rewriteRollupWindowOrderBy(astOrderBy, rewriteState) + if !ok { + return nil, true + } + if len(rewriteState.innerExprs) == 0 { + _, ok = rewriteState.ensureInnerExpr(tree.NewNumVal(int64(1), "1", false, tree.P_int64)) + if !ok { + return nil, true + } } + groupingCount := len(selectClause.GroupBy.GroupByExprsList) selectStmts := make([]*tree.SelectClause, groupingCount) for i, list := range selectClause.GroupBy.GroupByExprsList { selectStmts[i] = &tree.SelectClause{ - Distinct: selectClause.Distinct, - Exprs: append(tree.SelectExprs(nil), innerExprs...), - From: selectClause.From, - Where: selectClause.Where, + Exprs: append(tree.SelectExprs(nil), rewriteState.innerExprs...), + From: selectClause.From, + Where: selectClause.Where, GroupBy: &tree.GroupByClause{ GroupByExprsList: selectClause.GroupBy.GroupByExprsList, GroupingSet: list, @@ -3338,7 +3373,7 @@ func rewriteRollupWindowSelect( Cube: false, Rollup: false, }, - Having: selectClause.Having, + Having: rewrittenHaving, Option: selectClause.Option, } } @@ -3354,7 +3389,8 @@ func rewriteRollupWindowSelect( derived := &tree.Select{Select: leftClause} return &tree.Select{ Select: &tree.SelectClause{ - Exprs: outerExprs, + Distinct: selectClause.Distinct, + Exprs: outerExprs, From: &tree.From{ Tables: tree.TableExprs{ &tree.AliasedTableExpr{ @@ -3366,96 +3402,174 @@ func rewriteRollupWindowSelect( }, }, }, - OrderBy: rewriteRollupWindowOrderBy(astOrderBy, exprAliases), + OrderBy: rewrittenOrderBy, Limit: astLimit, RankOption: astRankOption, }, true } -func buildRollupWindowSelectExprs(selectExprs tree.SelectExprs) (tree.SelectExprs, tree.SelectExprs, map[string]string, bool) { +const rollupWindowInternalAliasPrefix = "__mo_rollup_window_col_" + +type rollupWindowRewriteState struct { + innerExprs tree.SelectExprs + exprAliases map[string]string + nameAliases map[string]string + usedAliases map[string]struct{} + nextAlias int +} + +func newRollupWindowRewriteState(selectExprs tree.SelectExprs) *rollupWindowRewriteState { + state := &rollupWindowRewriteState{ + innerExprs: make(tree.SelectExprs, 0, len(selectExprs)), + exprAliases: make(map[string]string), + nameAliases: make(map[string]string), + usedAliases: make(map[string]struct{}), + } + for _, selectExpr := range selectExprs { + if selectExpr.As != nil && !selectExpr.As.Empty() { + state.usedAliases[strings.ToLower(selectExpr.As.Origin())] = struct{}{} + } + if colName, ok := rollupWindowBareColumnName(selectExpr.Expr); ok { + state.usedAliases[strings.ToLower(colName)] = struct{}{} + } + } + return state +} + +func (state *rollupWindowRewriteState) newInternalAlias() string { + for { + alias := fmt.Sprintf("%s%d", rollupWindowInternalAliasPrefix, state.nextAlias) + state.nextAlias++ + key := strings.ToLower(alias) + if _, exists := state.usedAliases[key]; exists { + continue + } + state.usedAliases[key] = struct{}{} + return alias + } +} + +func (state *rollupWindowRewriteState) ensureInnerExpr(expr tree.Expr) (string, bool) { + if rollupWindowExprIsStar(expr) { + return "", false + } + if alias, ok := state.lookupExprAlias(expr); ok { + return alias, true + } + + alias := state.newInternalAlias() + state.innerExprs = append(state.innerExprs, tree.SelectExpr{ + Expr: expr, + As: tree.NewCStr(alias, 1), + }) + state.addExprAlias(expr, alias) + if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star && unresolvedName.NumParts == 1 { + state.addNameAlias(unresolvedName.ColNameOrigin(), alias) + } + return alias, true +} + +func (state *rollupWindowRewriteState) addExprAlias(expr tree.Expr, alias string) { + key := rollupWindowExprKey(expr) + if _, exists := state.exprAliases[key]; !exists { + state.exprAliases[key] = alias + } +} + +func (state *rollupWindowRewriteState) addNameAlias(name, alias string) { + key := strings.ToLower(name) + if key == "" { + return + } + if _, exists := state.nameAliases[key]; !exists { + state.nameAliases[key] = alias + } +} + +func (state *rollupWindowRewriteState) lookupExprAlias(expr tree.Expr) (string, bool) { + if alias, ok := state.exprAliases[rollupWindowExprKey(expr)]; ok { + return alias, true + } + if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star && unresolvedName.NumParts == 1 { + alias, ok := state.nameAliases[strings.ToLower(unresolvedName.ColNameOrigin())] + return alias, ok + } + return "", false +} + +func rollupWindowExprKey(expr tree.Expr) string { + return tree.String(expr, dialect.MYSQL) +} + +func buildRollupWindowSelectExprs(selectExprs tree.SelectExprs, state *rollupWindowRewriteState) (tree.SelectExprs, bool) { if len(selectExprs) == 0 { - return nil, nil, nil, false + return nil, false } - hasWindow := false - innerExprs := make(tree.SelectExprs, 0, len(selectExprs)) outerExprs := make(tree.SelectExprs, 0, len(selectExprs)) - exprAliases := make(map[string]string) - usedAliases := make(map[string]struct{}) aliasesByIndex := make([]string, len(selectExprs)) + for _, selectExpr := range selectExprs { + if rollupWindowExprContainsWindow(selectExpr.Expr) && selectExpr.As != nil && !selectExpr.As.Empty() { + state.addNameAlias(selectExpr.As.Origin(), selectExpr.As.Origin()) + } + } for i, selectExpr := range selectExprs { if rollupWindowExprContainsWindow(selectExpr.Expr) { - hasWindow = true continue } - alias, ok := rollupWindowSelectExprAlias(selectExpr, usedAliases) + alias, ok := state.ensureInnerExpr(selectExpr.Expr) if !ok { - return nil, nil, nil, false + return nil, false } - - aliasesByIndex[i] = alias - innerExprs = append(innerExprs, selectExpr) - addRollupWindowExprAlias(exprAliases, selectExpr.Expr, alias) - addRollupWindowExprAlias(exprAliases, tree.NewUnresolvedColName(alias), alias) - if colName, ok := rollupWindowBareColumnName(selectExpr.Expr); ok { - addRollupWindowExprAlias(exprAliases, tree.NewUnresolvedColName(colName), alias) + if selectExpr.As != nil && !selectExpr.As.Empty() { + state.addNameAlias(selectExpr.As.Origin(), alias) } - } - - if !hasWindow || len(innerExprs) == 0 { - return nil, nil, nil, false + aliasesByIndex[i] = alias } for i, selectExpr := range selectExprs { if rollupWindowExprContainsWindow(selectExpr.Expr) { - rewrittenExpr, ok := rewriteRollupWindowExpr(selectExpr.Expr, exprAliases) + rewrittenExpr, ok := rewriteRollupWindowExpr(selectExpr.Expr, state) if !ok { - return nil, nil, nil, false + return nil, false } outerExprs = append(outerExprs, tree.SelectExpr{ Expr: rewrittenExpr, - As: selectExpr.As, + As: rollupWindowOutputAlias(selectExpr), }) continue } outerExprs = append(outerExprs, tree.SelectExpr{ Expr: tree.NewUnresolvedColName(aliasesByIndex[i]), - As: selectExpr.As, + As: rollupWindowOutputAlias(selectExpr), }) } - return innerExprs, outerExprs, exprAliases, true + return outerExprs, true } -func rollupWindowSelectExprAlias(selectExpr tree.SelectExpr, usedAliases map[string]struct{}) (string, bool) { +func rollupWindowOutputAlias(selectExpr tree.SelectExpr) *tree.CStr { if selectExpr.As != nil && !selectExpr.As.Empty() { - alias := selectExpr.As.Origin() - if useRollupWindowAlias(alias, usedAliases) { - return alias, true - } - return "", false + return selectExpr.As } - - if alias, ok := rollupWindowBareColumnName(selectExpr.Expr); ok && useRollupWindowAlias(alias, usedAliases) { - return alias, true + if colName, ok := rollupWindowBareColumnName(selectExpr.Expr); ok { + return tree.NewCStr(colName, 1) } - - return "", false -} - -func useRollupWindowAlias(alias string, usedAliases map[string]struct{}) bool { - if alias == "" { - return false + expr := selectExpr.Expr + for { + parenExpr, ok := expr.(*tree.ParenExpr) + if !ok { + break + } + expr = parenExpr.Expr } - key := strings.ToLower(alias) - if _, ok := usedAliases[key]; ok { - return false + if heading, ok := nameConstHeading(expr); ok { + return tree.NewCStr(heading, 1) } - usedAliases[key] = struct{}{} - return true + return tree.NewCStr(tree.String(expr, dialect.MYSQL), 1) } func rollupWindowBareColumnName(expr tree.Expr) (string, bool) { @@ -3466,45 +3580,51 @@ func rollupWindowBareColumnName(expr tree.Expr) (string, bool) { return unresolvedName.ColNameOrigin(), true } -func addRollupWindowExprAlias(exprAliases map[string]string, expr tree.Expr, alias string) { - exprKey := tree.String(expr, dialect.MYSQL) - exprAliases[exprKey] = alias - exprAliases[strings.ToLower(exprKey)] = alias -} - -func lookupRollupWindowExprAlias(exprAliases map[string]string, expr tree.Expr) (string, bool) { - exprKey := tree.String(expr, dialect.MYSQL) - if alias, ok := exprAliases[exprKey]; ok { - return alias, true +func rollupWindowExprIsStar(expr tree.Expr) bool { + switch e := expr.(type) { + case *tree.UnresolvedName: + return e.Star + case tree.UnqualifiedStar: + return true + default: + return false } - alias, ok := exprAliases[strings.ToLower(exprKey)] - return alias, ok } -func rewriteRollupWindowExpr(expr tree.Expr, exprAliases map[string]string) (tree.Expr, bool) { - if alias, ok := lookupRollupWindowExprAlias(exprAliases, expr); ok { +func rewriteRollupWindowExpr(expr tree.Expr, state *rollupWindowRewriteState) (tree.Expr, bool) { + if alias, ok := state.lookupExprAlias(expr); ok { return tree.NewUnresolvedColName(alias), true } switch e := expr.(type) { case *tree.FuncExpr: + if rollupWindowFuncMustBeMaterialized(e) { + if rollupWindowExprContainsWindow(e) { + return nil, false + } + alias, ok := state.ensureInnerExpr(e) + if !ok { + return nil, false + } + return tree.NewUnresolvedColName(alias), true + } next := *e if len(e.Exprs) > 0 { - exprs, ok := rewriteRollupWindowExprs(e.Exprs, exprAliases) + exprs, ok := rewriteRollupWindowExprs(e.Exprs, state) if !ok { return nil, false } next.Exprs = exprs } if len(e.OrderBy) > 0 { - orderBy, ok := rewriteRollupWindowOrderByStrict(e.OrderBy, exprAliases) + orderBy, ok := rewriteRollupWindowOrderBy(e.OrderBy, state) if !ok { return nil, false } next.OrderBy = orderBy } if e.WindowSpec != nil { - windowSpec, ok := rewriteRollupWindowSpec(e.WindowSpec, exprAliases) + windowSpec, ok := rewriteRollupWindowSpec(e.WindowSpec, state) if !ok { return nil, false } @@ -3512,11 +3632,11 @@ func rewriteRollupWindowExpr(expr tree.Expr, exprAliases map[string]string) (tre } return &next, true case *tree.BinaryExpr: - left, ok := rewriteRollupWindowExpr(e.Left, exprAliases) + left, ok := rewriteRollupWindowExpr(e.Left, state) if !ok { return nil, false } - right, ok := rewriteRollupWindowExpr(e.Right, exprAliases) + right, ok := rewriteRollupWindowExpr(e.Right, state) if !ok { return nil, false } @@ -3525,7 +3645,7 @@ func rewriteRollupWindowExpr(expr tree.Expr, exprAliases map[string]string) (tre next.Right = right return &next, true case *tree.UnaryExpr: - child, ok := rewriteRollupWindowExpr(e.Expr, exprAliases) + child, ok := rewriteRollupWindowExpr(e.Expr, state) if !ok { return nil, false } @@ -3535,21 +3655,21 @@ func rewriteRollupWindowExpr(expr tree.Expr, exprAliases map[string]string) (tre case *tree.ComparisonExpr: next := *e if e.Left != nil { - left, ok := rewriteRollupWindowExpr(e.Left, exprAliases) + left, ok := rewriteRollupWindowExpr(e.Left, state) if !ok { return nil, false } next.Left = left } if e.Right != nil { - right, ok := rewriteRollupWindowExpr(e.Right, exprAliases) + right, ok := rewriteRollupWindowExpr(e.Right, state) if !ok { return nil, false } next.Right = right } if e.Escape != nil { - escape, ok := rewriteRollupWindowExpr(e.Escape, exprAliases) + escape, ok := rewriteRollupWindowExpr(e.Escape, state) if !ok { return nil, false } @@ -3557,11 +3677,11 @@ func rewriteRollupWindowExpr(expr tree.Expr, exprAliases map[string]string) (tre } return &next, true case *tree.AndExpr: - left, ok := rewriteRollupWindowExpr(e.Left, exprAliases) + left, ok := rewriteRollupWindowExpr(e.Left, state) if !ok { return nil, false } - right, ok := rewriteRollupWindowExpr(e.Right, exprAliases) + right, ok := rewriteRollupWindowExpr(e.Right, state) if !ok { return nil, false } @@ -3570,11 +3690,11 @@ func rewriteRollupWindowExpr(expr tree.Expr, exprAliases map[string]string) (tre next.Right = right return &next, true case *tree.XorExpr: - left, ok := rewriteRollupWindowExpr(e.Left, exprAliases) + left, ok := rewriteRollupWindowExpr(e.Left, state) if !ok { return nil, false } - right, ok := rewriteRollupWindowExpr(e.Right, exprAliases) + right, ok := rewriteRollupWindowExpr(e.Right, state) if !ok { return nil, false } @@ -3583,11 +3703,11 @@ func rewriteRollupWindowExpr(expr tree.Expr, exprAliases map[string]string) (tre next.Right = right return &next, true case *tree.OrExpr: - left, ok := rewriteRollupWindowExpr(e.Left, exprAliases) + left, ok := rewriteRollupWindowExpr(e.Left, state) if !ok { return nil, false } - right, ok := rewriteRollupWindowExpr(e.Right, exprAliases) + right, ok := rewriteRollupWindowExpr(e.Right, state) if !ok { return nil, false } @@ -3596,7 +3716,7 @@ func rewriteRollupWindowExpr(expr tree.Expr, exprAliases map[string]string) (tre next.Right = right return &next, true case *tree.NotExpr: - child, ok := rewriteRollupWindowExpr(e.Expr, exprAliases) + child, ok := rewriteRollupWindowExpr(e.Expr, state) if !ok { return nil, false } @@ -3604,7 +3724,7 @@ func rewriteRollupWindowExpr(expr tree.Expr, exprAliases map[string]string) (tre next.Expr = child return &next, true case *tree.ParenExpr: - child, ok := rewriteRollupWindowExpr(e.Expr, exprAliases) + child, ok := rewriteRollupWindowExpr(e.Expr, state) if !ok { return nil, false } @@ -3612,7 +3732,7 @@ func rewriteRollupWindowExpr(expr tree.Expr, exprAliases map[string]string) (tre next.Expr = child return &next, true case *tree.CastExpr: - child, ok := rewriteRollupWindowExpr(e.Expr, exprAliases) + child, ok := rewriteRollupWindowExpr(e.Expr, state) if !ok { return nil, false } @@ -3620,7 +3740,7 @@ func rewriteRollupWindowExpr(expr tree.Expr, exprAliases map[string]string) (tre next.Expr = child return &next, true case *tree.Tuple: - exprs, ok := rewriteRollupWindowExprs(e.Exprs, exprAliases) + exprs, ok := rewriteRollupWindowExprs(e.Exprs, state) if !ok { return nil, false } @@ -3630,7 +3750,7 @@ func rewriteRollupWindowExpr(expr tree.Expr, exprAliases map[string]string) (tre case *tree.CaseExpr: next := *e if e.Expr != nil { - caseExpr, ok := rewriteRollupWindowExpr(e.Expr, exprAliases) + caseExpr, ok := rewriteRollupWindowExpr(e.Expr, state) if !ok { return nil, false } @@ -3639,11 +3759,11 @@ func rewriteRollupWindowExpr(expr tree.Expr, exprAliases map[string]string) (tre if len(e.Whens) > 0 { next.Whens = make([]*tree.When, len(e.Whens)) for i, when := range e.Whens { - cond, ok := rewriteRollupWindowExpr(when.Cond, exprAliases) + cond, ok := rewriteRollupWindowExpr(when.Cond, state) if !ok { return nil, false } - val, ok := rewriteRollupWindowExpr(when.Val, exprAliases) + val, ok := rewriteRollupWindowExpr(when.Val, state) if !ok { return nil, false } @@ -3651,28 +3771,58 @@ func rewriteRollupWindowExpr(expr tree.Expr, exprAliases map[string]string) (tre } } if e.Else != nil { - elseExpr, ok := rewriteRollupWindowExpr(e.Else, exprAliases) + elseExpr, ok := rewriteRollupWindowExpr(e.Else, state) if !ok { return nil, false } next.Else = elseExpr } return &next, true - case *tree.NumVal, *tree.StrVal, *tree.TimeUnitExpr: + case *tree.NumVal, *tree.StrVal, *tree.TimeUnitExpr, *tree.ParamExpr, *tree.VarExpr: + return expr, true + case *tree.UnresolvedName: + if e.Star { + return expr, true + } + alias, ok := state.ensureInnerExpr(expr) + if !ok { + return nil, false + } + return tree.NewUnresolvedColName(alias), true + case tree.UnqualifiedStar: return expr, true default: - return nil, false + if rollupWindowExprContainsWindow(expr) { + return nil, false + } + alias, ok := state.ensureInnerExpr(expr) + if !ok { + return nil, false + } + return tree.NewUnresolvedColName(alias), true + } +} + +func rollupWindowFuncMustBeMaterialized(expr *tree.FuncExpr) bool { + if expr.WindowSpec != nil { + return false + } + funcRef, ok := expr.Func.FunctionReference.(*tree.UnresolvedName) + if !ok { + return true } + funcName := funcRef.ColName() + return function.GetFunctionIsAggregateByName(funcName) || strings.EqualFold(funcName, "grouping") } -func rewriteRollupWindowExprs(exprs tree.Exprs, exprAliases map[string]string) (tree.Exprs, bool) { +func rewriteRollupWindowExprs(exprs tree.Exprs, state *rollupWindowRewriteState) (tree.Exprs, bool) { if len(exprs) == 0 { return nil, true } rewrittenExprs := make(tree.Exprs, len(exprs)) for i, expr := range exprs { - rewrittenExpr, ok := rewriteRollupWindowExpr(expr, exprAliases) + rewrittenExpr, ok := rewriteRollupWindowExpr(expr, state) if !ok { return nil, false } @@ -3681,24 +3831,24 @@ func rewriteRollupWindowExprs(exprs tree.Exprs, exprAliases map[string]string) ( return rewrittenExprs, true } -func rewriteRollupWindowSpec(windowSpec *tree.WindowSpec, exprAliases map[string]string) (*tree.WindowSpec, bool) { +func rewriteRollupWindowSpec(windowSpec *tree.WindowSpec, state *rollupWindowRewriteState) (*tree.WindowSpec, bool) { next := *windowSpec if len(windowSpec.PartitionBy) > 0 { - partitionBy, ok := rewriteRollupWindowExprs(windowSpec.PartitionBy, exprAliases) + partitionBy, ok := rewriteRollupWindowExprs(windowSpec.PartitionBy, state) if !ok { return nil, false } next.PartitionBy = partitionBy } if len(windowSpec.OrderBy) > 0 { - orderBy, ok := rewriteRollupWindowOrderByStrict(windowSpec.OrderBy, exprAliases) + orderBy, ok := rewriteRollupWindowOrderBy(windowSpec.OrderBy, state) if !ok { return nil, false } next.OrderBy = orderBy } if windowSpec.Frame != nil { - frame, ok := rewriteRollupWindowFrame(windowSpec.Frame, exprAliases) + frame, ok := rewriteRollupWindowFrame(windowSpec.Frame, state) if !ok { return nil, false } @@ -3707,17 +3857,17 @@ func rewriteRollupWindowSpec(windowSpec *tree.WindowSpec, exprAliases map[string return &next, true } -func rewriteRollupWindowFrame(frame *tree.FrameClause, exprAliases map[string]string) (*tree.FrameClause, bool) { +func rewriteRollupWindowFrame(frame *tree.FrameClause, state *rollupWindowRewriteState) (*tree.FrameClause, bool) { next := *frame if frame.Start != nil { - start, ok := rewriteRollupWindowFrameBound(frame.Start, exprAliases) + start, ok := rewriteRollupWindowFrameBound(frame.Start, state) if !ok { return nil, false } next.Start = start } if frame.End != nil { - end, ok := rewriteRollupWindowFrameBound(frame.End, exprAliases) + end, ok := rewriteRollupWindowFrameBound(frame.End, state) if !ok { return nil, false } @@ -3726,10 +3876,10 @@ func rewriteRollupWindowFrame(frame *tree.FrameClause, exprAliases map[string]st return &next, true } -func rewriteRollupWindowFrameBound(bound *tree.FrameBound, exprAliases map[string]string) (*tree.FrameBound, bool) { +func rewriteRollupWindowFrameBound(bound *tree.FrameBound, state *rollupWindowRewriteState) (*tree.FrameBound, bool) { next := *bound if bound.Expr != nil { - expr, ok := rewriteRollupWindowExpr(bound.Expr, exprAliases) + expr, ok := rewriteRollupWindowExpr(bound.Expr, state) if !ok { return nil, false } @@ -3738,30 +3888,14 @@ func rewriteRollupWindowFrameBound(bound *tree.FrameBound, exprAliases map[strin return &next, true } -func rewriteRollupWindowOrderBy(orderBy tree.OrderBy, exprAliases map[string]string) tree.OrderBy { - if len(orderBy) == 0 { - return nil - } - - rewrittenOrderBy := make(tree.OrderBy, len(orderBy)) - for i, order := range orderBy { - next := *order - if expr, ok := rewriteRollupWindowExpr(order.Expr, exprAliases); ok { - next.Expr = expr - } - rewrittenOrderBy[i] = &next - } - return rewrittenOrderBy -} - -func rewriteRollupWindowOrderByStrict(orderBy tree.OrderBy, exprAliases map[string]string) (tree.OrderBy, bool) { +func rewriteRollupWindowOrderBy(orderBy tree.OrderBy, state *rollupWindowRewriteState) (tree.OrderBy, bool) { if len(orderBy) == 0 { return nil, true } rewrittenOrderBy := make(tree.OrderBy, len(orderBy)) for i, order := range orderBy { - expr, ok := rewriteRollupWindowExpr(order.Expr, exprAliases) + expr, ok := rewriteRollupWindowExpr(order.Expr, state) if !ok { return nil, false } @@ -3804,16 +3938,46 @@ func rollupWindowExprContainsWindow(expr tree.Expr) bool { return rollupWindowExprContainsWindow(e.Left) || rollupWindowExprContainsWindow(e.Right) case *tree.NotExpr: return rollupWindowExprContainsWindow(e.Expr) + case *tree.IsNullExpr: + return rollupWindowExprContainsWindow(e.Expr) + case *tree.IsNotNullExpr: + return rollupWindowExprContainsWindow(e.Expr) + case *tree.IsUnknownExpr: + return rollupWindowExprContainsWindow(e.Expr) + case *tree.IsNotUnknownExpr: + return rollupWindowExprContainsWindow(e.Expr) + case *tree.IsTrueExpr: + return rollupWindowExprContainsWindow(e.Expr) + case *tree.IsNotTrueExpr: + return rollupWindowExprContainsWindow(e.Expr) + case *tree.IsFalseExpr: + return rollupWindowExprContainsWindow(e.Expr) + case *tree.IsNotFalseExpr: + return rollupWindowExprContainsWindow(e.Expr) + case *tree.ExprList: + for _, child := range e.Exprs { + if rollupWindowExprContainsWindow(child) { + return true + } + } case *tree.ParenExpr: return rollupWindowExprContainsWindow(e.Expr) case *tree.CastExpr: return rollupWindowExprContainsWindow(e.Expr) + case *tree.BitCastExpr: + return rollupWindowExprContainsWindow(e.Expr) + case *tree.SerialExtractExpr: + return rollupWindowExprContainsWindow(e.SerialExpr) || rollupWindowExprContainsWindow(e.IndexExpr) case *tree.Tuple: for _, child := range e.Exprs { if rollupWindowExprContainsWindow(child) { return true } } + case *tree.RangeCond: + return rollupWindowExprContainsWindow(e.Left) || + rollupWindowExprContainsWindow(e.From) || + rollupWindowExprContainsWindow(e.To) case *tree.CaseExpr: if rollupWindowExprContainsWindow(e.Expr) || rollupWindowExprContainsWindow(e.Else) { return true @@ -3823,6 +3987,12 @@ func rollupWindowExprContainsWindow(expr tree.Expr) bool { return true } } + case *tree.IntervalExpr: + return rollupWindowExprContainsWindow(e.Expr) + case *tree.DefaultVal: + return rollupWindowExprContainsWindow(e.Expr) + case *tree.VarExpr: + return rollupWindowExprContainsWindow(e.Expr) } return false } diff --git a/test/distributed/cases/window/rollup.result b/test/distributed/cases/window/rollup.result index 22d6c806c03f4..bdb2489a5435c 100644 --- a/test/distributed/cases/window/rollup.result +++ b/test/distributed/cases/window/rollup.result @@ -920,5 +920,46 @@ west apple 30 4 3 3 east banana 20 5 5 4 east apple 10 6 6 5 west banana 10 7 6 5 +select +region, +product, +sum(qty), +row_number() over (order by sum(qty) desc, region, product) as row_num, +rank() over (order by sum(qty) desc) as rank_num, +dense_rank() over (order by sum(qty) desc) as dense_rank_num +from +rollup_window_sales +group by +region, +product with rollup +order by sum(qty) desc, region, product; +region product sum(qty) row_num rank_num dense_rank_num +null null 70 1 1 1 +west null 40 2 2 2 +east null 30 3 3 3 +west apple 30 4 3 3 +east banana 20 5 5 4 +east apple 10 6 6 5 +west banana 10 7 6 5 +select +region, +product, +row_number() over (order by sum(qty) desc, region, product) as row_num, +rank() over (order by sum(qty) desc) as rank_num, +dense_rank() over (order by sum(qty) desc) as dense_rank_num +from +rollup_window_sales +group by +region, +product with rollup +order by row_num; +region product row_num rank_num dense_rank_num +null null 1 1 1 +west null 2 2 2 +east null 3 3 3 +west apple 4 3 3 +east banana 5 5 4 +east apple 6 6 5 +west banana 7 6 5 drop table rollup_window_sales; drop database rollup_test; diff --git a/test/distributed/cases/window/rollup.sql b/test/distributed/cases/window/rollup.sql index 2f023abb3be18..61f37b8e5a3c8 100644 --- a/test/distributed/cases/window/rollup.sql +++ b/test/distributed/cases/window/rollup.sql @@ -613,6 +613,35 @@ group by region, product with rollup order by total_qty desc, region, product; + +-- aggregate output without an alias must still be available to outer windows +select + region, + product, + sum(qty), + row_number() over (order by sum(qty) desc, region, product) as row_num, + rank() over (order by sum(qty) desc) as rank_num, + dense_rank() over (order by sum(qty) desc) as dense_rank_num +from + rollup_window_sales +group by + region, + product with rollup +order by sum(qty) desc, region, product; + +-- aggregates used only by window specifications are hidden inner projections +select + region, + product, + row_number() over (order by sum(qty) desc, region, product) as row_num, + rank() over (order by sum(qty) desc) as rank_num, + dense_rank() over (order by sum(qty) desc) as dense_rank_num +from + rollup_window_sales +group by + region, + product with rollup +order by row_num; drop table rollup_window_sales; drop database rollup_test; From 931c17f560fc48161bff6cb24b6051f2957e7d37 Mon Sep 17 00:00:00 2001 From: GreatRiver Date: Sun, 12 Jul 2026 16:53:14 +0800 Subject: [PATCH 4/9] fix: separate rollup window alias scopes --- pkg/sql/plan/build_test.go | 85 +++++++++++++++++++++++++++++--- pkg/sql/plan/query_builder.go | 91 +++++++++++++++++++++++------------ 2 files changed, 136 insertions(+), 40 deletions(-) diff --git a/pkg/sql/plan/build_test.go b/pkg/sql/plan/build_test.go index 8890a1019abd9..3d4ce0d3c6461 100644 --- a/pkg/sql/plan/build_test.go +++ b/pkg/sql/plan/build_test.go @@ -916,6 +916,66 @@ func TestRollupWindowRanksAfterRollupUnion(t *testing.T) { } } +func TestRollupWindowAliasCollisionsPreserveSourceScope(t *testing.T) { + mock := NewMockOptimizer(false) + tests := []struct { + name string + sql string + expectedHeadings []string + }{ + { + name: "select alias collides with source column", + sql: ` + select n_name as n_regionkey, n_regionkey, sum(n_nationkey), + row_number() over (order by n_regionkey) as rn + from nation + group by n_name, n_regionkey with rollup`, + expectedHeadings: []string{"n_regionkey", "n_regionkey", "sum(n_nationkey)", "rn"}, + }, + { + name: "window output alias collides with source column", + sql: ` + select n_name, n_regionkey, sum(n_nationkey), + row_number() over (order by n_regionkey) as n_regionkey + from nation + group by n_name, n_regionkey with rollup`, + expectedHeadings: []string{"n_name", "n_regionkey", "sum(n_nationkey)", "n_regionkey"}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + logicPlan, err := runOneStmt(mock, t, test.sql) + require.NoError(t, err) + query := logicPlan.GetQuery() + require.NotNil(t, query) + require.Equal(t, test.expectedHeadings, query.Headings) + require.NotEmpty(t, query.Steps) + + root := query.Nodes[query.Steps[len(query.Steps)-1]] + require.Len(t, root.ProjectList, 4) + require.Equal(t, int32(types.T_int32), root.ProjectList[1].Typ.Id) + + foundRowNumber := false + for _, node := range query.Nodes { + if node.NodeType != plan.Node_WINDOW { + continue + } + for _, winExpr := range node.WinSpecList { + spec := winExpr.GetW() + if spec == nil || spec.Name != "row_number" { + continue + } + foundRowNumber = true + require.Len(t, spec.OrderBy, 1) + require.Equal(t, int32(types.T_int32), spec.OrderBy[0].Expr.Typ.Id) + } + } + require.True(t, foundRowNumber) + }) + } +} + func TestRewriteRollupWindowSelectHelpers(t *testing.T) { stmt := mustParseRollupWindowSelect(t, ` select @@ -952,15 +1012,20 @@ func TestRewriteRollupWindowSelectHelpers(t *testing.T) { alias, ok := state.lookupExprAlias(tree.NewUnresolvedColName("l_returnflag")) require.True(t, ok) require.True(t, strings.HasPrefix(alias, rollupWindowInternalAliasPrefix)) - flagAlias, ok := state.lookupExprAlias(tree.NewUnresolvedColName("flag")) + _, ok = state.lookupExprAlias(tree.NewUnresolvedColName("flag")) + require.False(t, ok) + rewrittenFlag, ok := rewriteRollupWindowExprWithNameAliases( + tree.NewUnresolvedColName("flag"), state, state.havingAliases) require.True(t, ok) - require.Equal(t, alias, flagAlias) - upperFlagAlias, ok := state.lookupExprAlias(tree.NewUnresolvedColName("FLAG")) + require.Equal(t, alias, tree.String(rewrittenFlag, dialect.MYSQL)) + rewrittenFlag, ok = rewriteRollupWindowExprWithNameAliases( + tree.NewUnresolvedColName("FLAG"), state, state.havingAliases) require.True(t, ok) - require.Equal(t, alias, upperFlagAlias) - alias, ok = state.lookupExprAlias(tree.NewUnresolvedColName("total_qty")) + require.Equal(t, alias, tree.String(rewrittenFlag, dialect.MYSQL)) + rewrittenTotal, ok := rewriteRollupWindowExprWithNameAliases( + tree.NewUnresolvedColName("total_qty"), state, state.havingAliases) require.True(t, ok) - require.True(t, strings.HasPrefix(alias, rollupWindowInternalAliasPrefix)) + require.True(t, strings.HasPrefix(tree.String(rewrittenTotal, dialect.MYSQL), rollupWindowInternalAliasPrefix)) } func TestRewriteRollupWindowSelectGuards(t *testing.T) { @@ -1014,8 +1079,12 @@ func TestRewriteRollupWindowSelectGuards(t *testing.T) { state = newRollupWindowRewriteState(windowBeforeAliasExprs) outerExprs, ok := buildRollupWindowSelectExprs(windowBeforeAliasExprs, state) require.True(t, ok) - require.Len(t, state.innerExprs, 1) - require.Contains(t, tree.String(outerExprs[0].Expr, dialect.MYSQL), state.innerExprs[0].As.Origin()) + require.Len(t, state.innerExprs, 2) + sourceAlias := state.exprAliases[rollupWindowExprKey(tree.NewUnresolvedColName("total"))] + aggregateAlias := state.exprAliases[rollupWindowExprKey(mustParseRollupWindowExpr(t, "sum(a)"))] + rewrittenWindow := tree.String(outerExprs[0].Expr, dialect.MYSQL) + require.Contains(t, rewrittenWindow, sourceAlias) + require.NotContains(t, rewrittenWindow, aggregateAlias) duplicateAlias := mustParseRollupWindowSelect(t, "select a as x, b as x, row_number() over () from t") duplicateExprs := duplicateAlias.Select.(*tree.SelectClause).Exprs diff --git a/pkg/sql/plan/query_builder.go b/pkg/sql/plan/query_builder.go index 315ecf03169b1..1c45e034f6cf3 100644 --- a/pkg/sql/plan/query_builder.go +++ b/pkg/sql/plan/query_builder.go @@ -3338,7 +3338,7 @@ func rewriteRollupWindowSelect( } var rewrittenHaving *tree.Where if selectClause.Having != nil { - havingExpr, ok := rewriteRollupWindowExpr(selectClause.Having.Expr, rewriteState) + havingExpr, ok := rewriteRollupWindowExprWithNameAliases(selectClause.Having.Expr, rewriteState, rewriteState.havingAliases) if !ok { return nil, true } @@ -3348,7 +3348,7 @@ func rewriteRollupWindowSelect( Expr: havingExpr, } } - rewrittenOrderBy, ok := rewriteRollupWindowOrderBy(astOrderBy, rewriteState) + rewrittenOrderBy, ok := rewriteRollupWindowOrderByWithNameAliases(astOrderBy, rewriteState, rewriteState.outputAliases) if !ok { return nil, true } @@ -3411,23 +3411,31 @@ func rewriteRollupWindowSelect( const rollupWindowInternalAliasPrefix = "__mo_rollup_window_col_" type rollupWindowRewriteState struct { - innerExprs tree.SelectExprs - exprAliases map[string]string - nameAliases map[string]string - usedAliases map[string]struct{} - nextAlias int + innerExprs tree.SelectExprs + exprAliases map[string]string + havingAliases map[string]string + outputAliases map[string]string + activeNameAliases map[string]string + usedAliases map[string]struct{} + nextAlias int } func newRollupWindowRewriteState(selectExprs tree.SelectExprs) *rollupWindowRewriteState { state := &rollupWindowRewriteState{ - innerExprs: make(tree.SelectExprs, 0, len(selectExprs)), - exprAliases: make(map[string]string), - nameAliases: make(map[string]string), - usedAliases: make(map[string]struct{}), + innerExprs: make(tree.SelectExprs, 0, len(selectExprs)), + exprAliases: make(map[string]string), + havingAliases: make(map[string]string), + outputAliases: make(map[string]string), + usedAliases: make(map[string]struct{}), } for _, selectExpr := range selectExprs { if selectExpr.As != nil && !selectExpr.As.Empty() { - state.usedAliases[strings.ToLower(selectExpr.As.Origin())] = struct{}{} + alias := selectExpr.As.Origin() + key := strings.ToLower(alias) + state.usedAliases[key] = struct{}{} + if _, exists := state.outputAliases[key]; !exists { + state.outputAliases[key] = alias + } } if colName, ok := rollupWindowBareColumnName(selectExpr.Expr); ok { state.usedAliases[strings.ToLower(colName)] = struct{}{} @@ -3463,9 +3471,6 @@ func (state *rollupWindowRewriteState) ensureInnerExpr(expr tree.Expr) (string, As: tree.NewCStr(alias, 1), }) state.addExprAlias(expr, alias) - if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star && unresolvedName.NumParts == 1 { - state.addNameAlias(unresolvedName.ColNameOrigin(), alias) - } return alias, true } @@ -3476,25 +3481,26 @@ func (state *rollupWindowRewriteState) addExprAlias(expr tree.Expr, alias string } } -func (state *rollupWindowRewriteState) addNameAlias(name, alias string) { +func addRollupWindowNameAlias(aliases map[string]string, name, alias string) { key := strings.ToLower(name) if key == "" { return } - if _, exists := state.nameAliases[key]; !exists { - state.nameAliases[key] = alias + if _, exists := aliases[key]; !exists { + aliases[key] = alias } } func (state *rollupWindowRewriteState) lookupExprAlias(expr tree.Expr) (string, bool) { - if alias, ok := state.exprAliases[rollupWindowExprKey(expr)]; ok { - return alias, true - } - if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star && unresolvedName.NumParts == 1 { - alias, ok := state.nameAliases[strings.ToLower(unresolvedName.ColNameOrigin())] - return alias, ok + if state.activeNameAliases != nil { + if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star && unresolvedName.NumParts == 1 { + if alias, exists := state.activeNameAliases[strings.ToLower(unresolvedName.ColNameOrigin())]; exists { + return alias, true + } + } } - return "", false + alias, ok := state.exprAliases[rollupWindowExprKey(expr)] + return alias, ok } func rollupWindowExprKey(expr tree.Expr) string { @@ -3508,12 +3514,6 @@ func buildRollupWindowSelectExprs(selectExprs tree.SelectExprs, state *rollupWin outerExprs := make(tree.SelectExprs, 0, len(selectExprs)) aliasesByIndex := make([]string, len(selectExprs)) - for _, selectExpr := range selectExprs { - if rollupWindowExprContainsWindow(selectExpr.Expr) && selectExpr.As != nil && !selectExpr.As.Empty() { - state.addNameAlias(selectExpr.As.Origin(), selectExpr.As.Origin()) - } - } - for i, selectExpr := range selectExprs { if rollupWindowExprContainsWindow(selectExpr.Expr) { continue @@ -3524,7 +3524,7 @@ func buildRollupWindowSelectExprs(selectExprs tree.SelectExprs, state *rollupWin return nil, false } if selectExpr.As != nil && !selectExpr.As.Empty() { - state.addNameAlias(selectExpr.As.Origin(), alias) + addRollupWindowNameAlias(state.havingAliases, selectExpr.As.Origin(), alias) } aliasesByIndex[i] = alias } @@ -3592,6 +3592,9 @@ func rollupWindowExprIsStar(expr tree.Expr) bool { } func rewriteRollupWindowExpr(expr tree.Expr, state *rollupWindowRewriteState) (tree.Expr, bool) { + if funcExpr, ok := expr.(*tree.FuncExpr); ok && funcExpr.WindowSpec != nil && state.activeNameAliases != nil { + return rewriteRollupWindowExprWithNameAliases(expr, state, nil) + } if alias, ok := state.lookupExprAlias(expr); ok { return tree.NewUnresolvedColName(alias), true } @@ -3803,6 +3806,18 @@ func rewriteRollupWindowExpr(expr tree.Expr, state *rollupWindowRewriteState) (t } } +func rewriteRollupWindowExprWithNameAliases( + expr tree.Expr, + state *rollupWindowRewriteState, + aliases map[string]string, +) (tree.Expr, bool) { + previous := state.activeNameAliases + state.activeNameAliases = aliases + rewritten, ok := rewriteRollupWindowExpr(expr, state) + state.activeNameAliases = previous + return rewritten, ok +} + func rollupWindowFuncMustBeMaterialized(expr *tree.FuncExpr) bool { if expr.WindowSpec != nil { return false @@ -3906,6 +3921,18 @@ func rewriteRollupWindowOrderBy(orderBy tree.OrderBy, state *rollupWindowRewrite return rewrittenOrderBy, true } +func rewriteRollupWindowOrderByWithNameAliases( + orderBy tree.OrderBy, + state *rollupWindowRewriteState, + aliases map[string]string, +) (tree.OrderBy, bool) { + previous := state.activeNameAliases + state.activeNameAliases = aliases + rewritten, ok := rewriteRollupWindowOrderBy(orderBy, state) + state.activeNameAliases = previous + return rewritten, ok +} + func rollupWindowExprContainsWindow(expr tree.Expr) bool { switch e := expr.(type) { case *tree.FuncExpr: From aacefce4b13a996f295432f3ad01a17bcac97311 Mon Sep 17 00:00:00 2001 From: GreatRiver Date: Sun, 12 Jul 2026 16:57:05 +0800 Subject: [PATCH 5/9] test: cover rollup window order by alias scope --- pkg/sql/plan/build_test.go | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/pkg/sql/plan/build_test.go b/pkg/sql/plan/build_test.go index 41985994c0ce2..2050dcf85425b 100644 --- a/pkg/sql/plan/build_test.go +++ b/pkg/sql/plan/build_test.go @@ -919,9 +919,10 @@ func TestRollupWindowRanksAfterRollupUnion(t *testing.T) { func TestRollupWindowAliasCollisionsPreserveSourceScope(t *testing.T) { mock := NewMockOptimizer(false) tests := []struct { - name string - sql string - expectedHeadings []string + name string + sql string + expectedHeadings []string + expectedProjectLen int }{ { name: "select alias collides with source column", @@ -930,7 +931,8 @@ func TestRollupWindowAliasCollisionsPreserveSourceScope(t *testing.T) { row_number() over (order by n_regionkey) as rn from nation group by n_name, n_regionkey with rollup`, - expectedHeadings: []string{"n_regionkey", "n_regionkey", "sum(n_nationkey)", "rn"}, + expectedHeadings: []string{"n_regionkey", "n_regionkey", "sum(n_nationkey)", "rn"}, + expectedProjectLen: 4, }, { name: "window output alias collides with source column", @@ -939,7 +941,18 @@ func TestRollupWindowAliasCollisionsPreserveSourceScope(t *testing.T) { row_number() over (order by n_regionkey) as n_regionkey from nation group by n_name, n_regionkey with rollup`, - expectedHeadings: []string{"n_name", "n_regionkey", "sum(n_nationkey)", "n_regionkey"}, + expectedHeadings: []string{"n_name", "n_regionkey", "sum(n_nationkey)", "n_regionkey"}, + expectedProjectLen: 4, + }, + { + name: "final order by window keeps source scope", + sql: ` + select n_name as n_regionkey, n_regionkey, sum(n_nationkey) + from nation + group by n_name, n_regionkey with rollup + order by row_number() over (order by n_regionkey)`, + expectedHeadings: []string{"n_regionkey", "n_regionkey", "sum(n_nationkey)"}, + expectedProjectLen: 3, }, } @@ -953,7 +966,7 @@ func TestRollupWindowAliasCollisionsPreserveSourceScope(t *testing.T) { require.NotEmpty(t, query.Steps) root := query.Nodes[query.Steps[len(query.Steps)-1]] - require.Len(t, root.ProjectList, 4) + require.Len(t, root.ProjectList, test.expectedProjectLen) require.Equal(t, int32(types.T_int32), root.ProjectList[1].Typ.Id) foundRowNumber := false From afc33b658bf3075e9312b320f13a632aec5ddb36 Mon Sep 17 00:00:00 2001 From: GreatRiver Date: Tue, 14 Jul 2026 00:09:50 +0800 Subject: [PATCH 6/9] fix: preserve rollup having column precedence --- pkg/sql/plan/build_test.go | 60 +++++++++++++ pkg/sql/plan/query_builder.go | 94 +++++++++++++++++---- test/distributed/cases/window/rollup.result | 13 +++ test/distributed/cases/window/rollup.sql | 12 +++ 4 files changed, 161 insertions(+), 18 deletions(-) diff --git a/pkg/sql/plan/build_test.go b/pkg/sql/plan/build_test.go index 2050dcf85425b..96f8773cb3534 100644 --- a/pkg/sql/plan/build_test.go +++ b/pkg/sql/plan/build_test.go @@ -923,6 +923,7 @@ func TestRollupWindowAliasCollisionsPreserveSourceScope(t *testing.T) { sql string expectedHeadings []string expectedProjectLen int + expectedHavingType int32 }{ { name: "select alias collides with source column", @@ -954,6 +955,18 @@ func TestRollupWindowAliasCollisionsPreserveSourceScope(t *testing.T) { expectedHeadings: []string{"n_regionkey", "n_regionkey", "sum(n_nationkey)"}, expectedProjectLen: 3, }, + { + name: "having alias collision keeps source scope", + sql: ` + select sum(n_nationkey) as n_regionkey, n_regionkey, + row_number() over (order by n_regionkey) as rn + from nation + group by n_regionkey with rollup + having N_REGIONKEY > 0`, + expectedHeadings: []string{"n_regionkey", "n_regionkey", "rn"}, + expectedProjectLen: 3, + expectedHavingType: int32(types.T_int32), + }, } for _, test := range tests { @@ -985,10 +998,57 @@ func TestRollupWindowAliasCollisionsPreserveSourceScope(t *testing.T) { } } require.True(t, foundRowNumber) + + if test.expectedHavingType != 0 { + foundHaving := false + for _, node := range query.Nodes { + if node.NodeType != plan.Node_FILTER || !node.RollupFilter { + continue + } + for _, filter := range node.FilterList { + fn := filter.GetF() + require.NotNil(t, fn) + require.Equal(t, ">", fn.Func.ObjName) + require.NotEmpty(t, fn.Args) + require.Equal(t, test.expectedHavingType, fn.Args[0].Typ.Id) + foundHaving = true + } + } + require.True(t, foundHaving) + } }) } } +func TestRollupWindowHavingAliasCollisionWithHiddenGroup(t *testing.T) { + mock := NewMockOptimizer(false) + logicPlan, err := runOneStmt(mock, t, ` + select sum(n_nationkey) as n_regionkey, + row_number() over (order by sum(n_nationkey)) as rn + from nation + group by n_regionkey with rollup + having N_REGIONKEY > 0`) + require.NoError(t, err) + + query := logicPlan.GetQuery() + require.NotNil(t, query) + foundHaving := false + for _, node := range query.Nodes { + if node.NodeType != plan.Node_FILTER || !node.RollupFilter { + continue + } + for _, filter := range node.FilterList { + fn := filter.GetF() + require.NotNil(t, fn) + require.Equal(t, ">", fn.Func.ObjName) + require.NotEmpty(t, fn.Args) + require.Equal(t, int32(types.T_int32), fn.Args[0].Typ.Id) + foundHaving = true + } + } + require.True(t, foundHaving) +} + func TestRewriteRollupWindowSelectHelpers(t *testing.T) { stmt := mustParseRollupWindowSelect(t, ` select diff --git a/pkg/sql/plan/query_builder.go b/pkg/sql/plan/query_builder.go index 7cd9c8519608f..0851b793787bc 100644 --- a/pkg/sql/plan/query_builder.go +++ b/pkg/sql/plan/query_builder.go @@ -3552,13 +3552,18 @@ func rewriteRollupWindowSelect( } rewriteState := newRollupWindowRewriteState(selectClause.Exprs) + if selectClause.Having != nil && !rewriteState.addGroupedSourceExprs(selectClause.GroupBy) { + return nil, true + } outerExprs, ok := buildRollupWindowSelectExprs(selectClause.Exprs, rewriteState) if !ok { return nil, true } var rewrittenHaving *tree.Where if selectClause.Having != nil { - havingExpr, ok := rewriteRollupWindowExprWithNameAliases(selectClause.Having.Expr, rewriteState, rewriteState.havingAliases) + // HAVING uses AliasAfterColumn: a materialized source expression wins + // when its name collides with a SELECT alias. + havingExpr, ok := rewriteRollupWindowExprWithFallbackNameAliases(selectClause.Having.Expr, rewriteState, rewriteState.havingAliases) if !ok { return nil, true } @@ -3631,22 +3636,25 @@ func rewriteRollupWindowSelect( const rollupWindowInternalAliasPrefix = "__mo_rollup_window_col_" type rollupWindowRewriteState struct { - innerExprs tree.SelectExprs - exprAliases map[string]string - havingAliases map[string]string - outputAliases map[string]string - activeNameAliases map[string]string - usedAliases map[string]struct{} - nextAlias int + innerExprs tree.SelectExprs + exprAliases map[string]string + sourceNameAliases map[string]string + havingAliases map[string]string + outputAliases map[string]string + activeNameAliases map[string]string + nameAliasesAfterExpr bool + usedAliases map[string]struct{} + nextAlias int } func newRollupWindowRewriteState(selectExprs tree.SelectExprs) *rollupWindowRewriteState { state := &rollupWindowRewriteState{ - innerExprs: make(tree.SelectExprs, 0, len(selectExprs)), - exprAliases: make(map[string]string), - havingAliases: make(map[string]string), - outputAliases: make(map[string]string), - usedAliases: make(map[string]struct{}), + innerExprs: make(tree.SelectExprs, 0, len(selectExprs)), + exprAliases: make(map[string]string), + sourceNameAliases: make(map[string]string), + havingAliases: make(map[string]string), + outputAliases: make(map[string]string), + usedAliases: make(map[string]struct{}), } for _, selectExpr := range selectExprs { if selectExpr.As != nil && !selectExpr.As.Empty() { @@ -3664,6 +3672,21 @@ func newRollupWindowRewriteState(selectExprs tree.SelectExprs) *rollupWindowRewr return state } +func (state *rollupWindowRewriteState) addGroupedSourceExprs(groupBy *tree.GroupByClause) bool { + // A grouped column may only be referenced by HAVING. Materialize bare group + // names so AliasAfterColumn can still choose them ahead of SELECT aliases. + for _, groupByExprs := range groupBy.GroupByExprsList { + for _, expr := range groupByExprs { + if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star && unresolvedName.NumParts == 1 { + if _, ok := state.ensureInnerExpr(expr); !ok { + return false + } + } + } + } + return true +} + func (state *rollupWindowRewriteState) newInternalAlias() string { for { alias := fmt.Sprintf("%s%d", rollupWindowInternalAliasPrefix, state.nextAlias) @@ -3699,6 +3722,9 @@ func (state *rollupWindowRewriteState) addExprAlias(expr tree.Expr, alias string if _, exists := state.exprAliases[key]; !exists { state.exprAliases[key] = alias } + if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star && unresolvedName.NumParts == 1 { + addRollupWindowNameAlias(state.sourceNameAliases, unresolvedName.ColNameOrigin(), alias) + } } func addRollupWindowNameAlias(aliases map[string]string, name, alias string) { @@ -3712,6 +3738,17 @@ func addRollupWindowNameAlias(aliases map[string]string, name, alias string) { } func (state *rollupWindowRewriteState) lookupExprAlias(expr tree.Expr) (string, bool) { + exprKey := rollupWindowExprKey(expr) + if state.nameAliasesAfterExpr { + if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star && unresolvedName.NumParts == 1 { + if alias, exists := state.sourceNameAliases[strings.ToLower(unresolvedName.ColNameOrigin())]; exists { + return alias, true + } + } + if alias, ok := state.exprAliases[exprKey]; ok { + return alias, true + } + } if state.activeNameAliases != nil { if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star && unresolvedName.NumParts == 1 { if alias, exists := state.activeNameAliases[strings.ToLower(unresolvedName.ColNameOrigin())]; exists { @@ -3719,7 +3756,7 @@ func (state *rollupWindowRewriteState) lookupExprAlias(expr tree.Expr) (string, } } } - alias, ok := state.exprAliases[rollupWindowExprKey(expr)] + alias, ok := state.exprAliases[exprKey] return alias, ok } @@ -4031,10 +4068,28 @@ func rewriteRollupWindowExprWithNameAliases( state *rollupWindowRewriteState, aliases map[string]string, ) (tree.Expr, bool) { - previous := state.activeNameAliases + previousAliases := state.activeNameAliases + previousAfterExpr := state.nameAliasesAfterExpr + state.activeNameAliases = aliases + state.nameAliasesAfterExpr = false + rewritten, ok := rewriteRollupWindowExpr(expr, state) + state.activeNameAliases = previousAliases + state.nameAliasesAfterExpr = previousAfterExpr + return rewritten, ok +} + +func rewriteRollupWindowExprWithFallbackNameAliases( + expr tree.Expr, + state *rollupWindowRewriteState, + aliases map[string]string, +) (tree.Expr, bool) { + previousAliases := state.activeNameAliases + previousAfterExpr := state.nameAliasesAfterExpr state.activeNameAliases = aliases + state.nameAliasesAfterExpr = true rewritten, ok := rewriteRollupWindowExpr(expr, state) - state.activeNameAliases = previous + state.activeNameAliases = previousAliases + state.nameAliasesAfterExpr = previousAfterExpr return rewritten, ok } @@ -4146,10 +4201,13 @@ func rewriteRollupWindowOrderByWithNameAliases( state *rollupWindowRewriteState, aliases map[string]string, ) (tree.OrderBy, bool) { - previous := state.activeNameAliases + previousAliases := state.activeNameAliases + previousAfterExpr := state.nameAliasesAfterExpr state.activeNameAliases = aliases + state.nameAliasesAfterExpr = false rewritten, ok := rewriteRollupWindowOrderBy(orderBy, state) - state.activeNameAliases = previous + state.activeNameAliases = previousAliases + state.nameAliasesAfterExpr = previousAfterExpr return rewritten, ok } diff --git a/test/distributed/cases/window/rollup.result b/test/distributed/cases/window/rollup.result index bdb2489a5435c..11170134cca68 100644 --- a/test/distributed/cases/window/rollup.result +++ b/test/distributed/cases/window/rollup.result @@ -961,5 +961,18 @@ west apple 4 3 3 east banana 5 5 4 east apple 6 6 5 west banana 7 6 5 +select +sum(qty) as qty, +row_number() over (order by sum(qty), count(*)) as row_num +from +rollup_window_sales +group by +qty with rollup +having qty > 0 +order by row_num; +qty row_num +20 1 +20 2 +30 3 drop table rollup_window_sales; drop database rollup_test; diff --git a/test/distributed/cases/window/rollup.sql b/test/distributed/cases/window/rollup.sql index 61f37b8e5a3c8..841ef93189c6c 100644 --- a/test/distributed/cases/window/rollup.sql +++ b/test/distributed/cases/window/rollup.sql @@ -642,6 +642,18 @@ group by region, product with rollup order by row_num; + +-- HAVING keeps AliasAfterColumn precedence when an aggregate alias collides +-- with the grouped source column name +select + sum(qty) as qty, + row_number() over (order by sum(qty), count(*)) as row_num +from + rollup_window_sales +group by + qty with rollup +having qty > 0 +order by row_num; drop table rollup_window_sales; drop database rollup_test; From 5a45f9357ef3256b565b51d9ff34ee6f938d0497 Mon Sep 17 00:00:00 2001 From: GreatRiver Date: Tue, 14 Jul 2026 13:26:22 +0800 Subject: [PATCH 7/9] fix: preserve qualified rollup having scope --- pkg/sql/plan/build_test.go | 31 ++++++++++++++ pkg/sql/plan/query_builder.go | 47 +++++++++++++++------ test/distributed/cases/window/rollup.result | 8 ++-- test/distributed/cases/window/rollup.sql | 8 ++-- 4 files changed, 74 insertions(+), 20 deletions(-) diff --git a/pkg/sql/plan/build_test.go b/pkg/sql/plan/build_test.go index 96f8773cb3534..19d1e45bfc0be 100644 --- a/pkg/sql/plan/build_test.go +++ b/pkg/sql/plan/build_test.go @@ -967,6 +967,18 @@ func TestRollupWindowAliasCollisionsPreserveSourceScope(t *testing.T) { expectedProjectLen: 3, expectedHavingType: int32(types.T_int32), }, + { + name: "qualified grouped source keeps bare having source scope", + sql: ` + select sum(t.n_nationkey) as n_regionkey, t.n_regionkey, + row_number() over (order by t.n_regionkey) as rn + from nation t + group by t.n_regionkey with rollup + having n_regionkey > 0`, + expectedHeadings: []string{"n_regionkey", "n_regionkey", "rn"}, + expectedProjectLen: 3, + expectedHavingType: int32(types.T_int32), + }, } for _, test := range tests { @@ -1049,6 +1061,25 @@ func TestRollupWindowHavingAliasCollisionWithHiddenGroup(t *testing.T) { require.True(t, foundHaving) } +func TestRollupWindowSourceNameAliasAmbiguity(t *testing.T) { + state := newRollupWindowRewriteState(nil) + state.addSourceNameAlias("n_regionkey", "left_alias") + state.addSourceNameAlias("N_REGIONKEY", "right_alias") + + _, exists := state.sourceNameAliases["n_regionkey"] + require.False(t, exists) + _, ambiguous := state.ambiguousSourceNames["n_regionkey"] + require.True(t, ambiguous) + + state.havingAliases["n_regionkey"] = "aggregate_alias" + rewritten, ok := rewriteRollupWindowExprWithFallbackNameAliases( + tree.NewUnresolvedColName("n_regionkey"), state, state.havingAliases) + require.True(t, ok) + require.NotEqual(t, "aggregate_alias", tree.String(rewritten, dialect.MYSQL)) + require.Len(t, state.innerExprs, 1) + require.Equal(t, "n_regionkey", tree.String(state.innerExprs[0].Expr, dialect.MYSQL)) +} + func TestRewriteRollupWindowSelectHelpers(t *testing.T) { stmt := mustParseRollupWindowSelect(t, ` select diff --git a/pkg/sql/plan/query_builder.go b/pkg/sql/plan/query_builder.go index 0851b793787bc..452560c3b1951 100644 --- a/pkg/sql/plan/query_builder.go +++ b/pkg/sql/plan/query_builder.go @@ -3639,6 +3639,7 @@ type rollupWindowRewriteState struct { innerExprs tree.SelectExprs exprAliases map[string]string sourceNameAliases map[string]string + ambiguousSourceNames map[string]struct{} havingAliases map[string]string outputAliases map[string]string activeNameAliases map[string]string @@ -3649,12 +3650,13 @@ type rollupWindowRewriteState struct { func newRollupWindowRewriteState(selectExprs tree.SelectExprs) *rollupWindowRewriteState { state := &rollupWindowRewriteState{ - innerExprs: make(tree.SelectExprs, 0, len(selectExprs)), - exprAliases: make(map[string]string), - sourceNameAliases: make(map[string]string), - havingAliases: make(map[string]string), - outputAliases: make(map[string]string), - usedAliases: make(map[string]struct{}), + innerExprs: make(tree.SelectExprs, 0, len(selectExprs)), + exprAliases: make(map[string]string), + sourceNameAliases: make(map[string]string), + ambiguousSourceNames: make(map[string]struct{}), + havingAliases: make(map[string]string), + outputAliases: make(map[string]string), + usedAliases: make(map[string]struct{}), } for _, selectExpr := range selectExprs { if selectExpr.As != nil && !selectExpr.As.Empty() { @@ -3673,11 +3675,12 @@ func newRollupWindowRewriteState(selectExprs tree.SelectExprs) *rollupWindowRewr } func (state *rollupWindowRewriteState) addGroupedSourceExprs(groupBy *tree.GroupByClause) bool { - // A grouped column may only be referenced by HAVING. Materialize bare group - // names so AliasAfterColumn can still choose them ahead of SELECT aliases. + // A grouped column may only be referenced by HAVING. Materialize grouped + // names so AliasAfterColumn can still choose their terminal column name + // ahead of SELECT aliases. for _, groupByExprs := range groupBy.GroupByExprsList { for _, expr := range groupByExprs { - if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star && unresolvedName.NumParts == 1 { + if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star { if _, ok := state.ensureInnerExpr(expr); !ok { return false } @@ -3722,11 +3725,27 @@ func (state *rollupWindowRewriteState) addExprAlias(expr tree.Expr, alias string if _, exists := state.exprAliases[key]; !exists { state.exprAliases[key] = alias } - if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star && unresolvedName.NumParts == 1 { - addRollupWindowNameAlias(state.sourceNameAliases, unresolvedName.ColNameOrigin(), alias) + if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star { + state.addSourceNameAlias(unresolvedName.ColNameOrigin(), alias) } } +func (state *rollupWindowRewriteState) addSourceNameAlias(name, alias string) { + key := strings.ToLower(name) + if key == "" { + return + } + if _, ambiguous := state.ambiguousSourceNames[key]; ambiguous { + return + } + if existing, exists := state.sourceNameAliases[key]; exists && existing != alias { + delete(state.sourceNameAliases, key) + state.ambiguousSourceNames[key] = struct{}{} + return + } + state.sourceNameAliases[key] = alias +} + func addRollupWindowNameAlias(aliases map[string]string, name, alias string) { key := strings.ToLower(name) if key == "" { @@ -3741,7 +3760,11 @@ func (state *rollupWindowRewriteState) lookupExprAlias(expr tree.Expr) (string, exprKey := rollupWindowExprKey(expr) if state.nameAliasesAfterExpr { if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star && unresolvedName.NumParts == 1 { - if alias, exists := state.sourceNameAliases[strings.ToLower(unresolvedName.ColNameOrigin())]; exists { + name := strings.ToLower(unresolvedName.ColNameOrigin()) + if _, ambiguous := state.ambiguousSourceNames[name]; ambiguous { + return "", false + } + if alias, exists := state.sourceNameAliases[name]; exists { return alias, true } } diff --git a/test/distributed/cases/window/rollup.result b/test/distributed/cases/window/rollup.result index 11170134cca68..2f4c38a21d0a0 100644 --- a/test/distributed/cases/window/rollup.result +++ b/test/distributed/cases/window/rollup.result @@ -962,12 +962,12 @@ east banana 5 5 4 east apple 6 6 5 west banana 7 6 5 select -sum(qty) as qty, -row_number() over (order by sum(qty), count(*)) as row_num +sum(r.qty) as qty, +row_number() over (order by sum(r.qty), count(*)) as row_num from -rollup_window_sales +rollup_window_sales r group by -qty with rollup +r.qty with rollup having qty > 0 order by row_num; qty row_num diff --git a/test/distributed/cases/window/rollup.sql b/test/distributed/cases/window/rollup.sql index 841ef93189c6c..c78d6b9844fe1 100644 --- a/test/distributed/cases/window/rollup.sql +++ b/test/distributed/cases/window/rollup.sql @@ -646,12 +646,12 @@ order by row_num; -- HAVING keeps AliasAfterColumn precedence when an aggregate alias collides -- with the grouped source column name select - sum(qty) as qty, - row_number() over (order by sum(qty), count(*)) as row_num + sum(r.qty) as qty, + row_number() over (order by sum(r.qty), count(*)) as row_num from - rollup_window_sales + rollup_window_sales r group by - qty with rollup + r.qty with rollup having qty > 0 order by row_num; drop table rollup_window_sales; From fd363915d92ad0928139abe1ad098ef73a30e3c8 Mon Sep 17 00:00:00 2001 From: GreatRiver Date: Wed, 15 Jul 2026 10:20:30 +0800 Subject: [PATCH 8/9] fix: preserve rollup having binding scope --- pkg/sql/plan/build_test.go | 89 ++++++++++++++------- pkg/sql/plan/query_builder.go | 140 +++++++++++----------------------- 2 files changed, 104 insertions(+), 125 deletions(-) diff --git a/pkg/sql/plan/build_test.go b/pkg/sql/plan/build_test.go index d11355e4382d9..de49ac309e9fd 100644 --- a/pkg/sql/plan/build_test.go +++ b/pkg/sql/plan/build_test.go @@ -1061,23 +1061,62 @@ func TestRollupWindowHavingAliasCollisionWithHiddenGroup(t *testing.T) { require.True(t, foundHaving) } -func TestRollupWindowSourceNameAliasAmbiguity(t *testing.T) { - state := newRollupWindowRewriteState(nil) - state.addSourceNameAlias("n_regionkey", "left_alias") - state.addSourceNameAlias("N_REGIONKEY", "right_alias") - - _, exists := state.sourceNameAliases["n_regionkey"] - require.False(t, exists) - _, ambiguous := state.ambiguousSourceNames["n_regionkey"] - require.True(t, ambiguous) +func TestRollupWindowHavingPreservesFromScopeErrors(t *testing.T) { + mock := NewMockOptimizer(false) + tests := []struct { + name string + sql string + wantError string + }{ + { + name: "ambiguous source without window", + sql: ` + select n1.n_regionkey + from nation n1 + join nation n2 on n1.n_nationkey = n2.n_nationkey + group by n1.n_regionkey with rollup + having n_regionkey > 0`, + wantError: "ambiguous column reference", + }, + { + name: "ambiguous source with window", + sql: ` + select n1.n_regionkey, + row_number() over (order by n1.n_regionkey) as rn + from nation n1 + join nation n2 on n1.n_nationkey = n2.n_nationkey + group by n1.n_regionkey with rollup + having n_regionkey > 0`, + wantError: "ambiguous column reference", + }, + { + name: "non-grouped source without window", + sql: ` + select sum(n_nationkey) as n_name + from nation + group by n_regionkey with rollup + having n_name <> ''`, + wantError: "must appear in the GROUP BY clause", + }, + { + name: "non-grouped source with window", + sql: ` + select sum(n_nationkey) as n_name, + row_number() over (order by sum(n_nationkey)) as rn + from nation + group by n_regionkey with rollup + having n_name <> ''`, + wantError: "must appear in the GROUP BY clause", + }, + } - state.havingAliases["n_regionkey"] = "aggregate_alias" - rewritten, ok := rewriteRollupWindowExprWithFallbackNameAliases( - tree.NewUnresolvedColName("n_regionkey"), state, state.havingAliases) - require.True(t, ok) - require.NotEqual(t, "aggregate_alias", tree.String(rewritten, dialect.MYSQL)) - require.Len(t, state.innerExprs, 1) - require.Equal(t, "n_regionkey", tree.String(state.innerExprs[0].Expr, dialect.MYSQL)) + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, err := runOneStmt(mock, t, test.sql) + require.Error(t, err) + require.Contains(t, err.Error(), test.wantError) + }) + } } func TestRewriteRollupWindowSelectHelpers(t *testing.T) { @@ -1118,18 +1157,12 @@ func TestRewriteRollupWindowSelectHelpers(t *testing.T) { require.True(t, strings.HasPrefix(alias, rollupWindowInternalAliasPrefix)) _, ok = state.lookupExprAlias(tree.NewUnresolvedColName("flag")) require.False(t, ok) - rewrittenFlag, ok := rewriteRollupWindowExprWithNameAliases( - tree.NewUnresolvedColName("flag"), state, state.havingAliases) - require.True(t, ok) - require.Equal(t, alias, tree.String(rewrittenFlag, dialect.MYSQL)) - rewrittenFlag, ok = rewriteRollupWindowExprWithNameAliases( - tree.NewUnresolvedColName("FLAG"), state, state.havingAliases) - require.True(t, ok) - require.Equal(t, alias, tree.String(rewrittenFlag, dialect.MYSQL)) - rewrittenTotal, ok := rewriteRollupWindowExprWithNameAliases( - tree.NewUnresolvedColName("total_qty"), state, state.havingAliases) - require.True(t, ok) - require.True(t, strings.HasPrefix(tree.String(rewrittenTotal, dialect.MYSQL), rollupWindowInternalAliasPrefix)) + + state.addHavingAliasExprs(clause.Exprs) + require.Len(t, state.innerExprs, 6) + require.Equal(t, "flag", state.innerExprs[3].As.Origin()) + require.Equal(t, "status", state.innerExprs[4].As.Origin()) + require.Equal(t, "total_qty", state.innerExprs[5].As.Origin()) } func TestRewriteRollupWindowSelectGuards(t *testing.T) { diff --git a/pkg/sql/plan/query_builder.go b/pkg/sql/plan/query_builder.go index 55377a336bd4c..793701fc4c4ea 100644 --- a/pkg/sql/plan/query_builder.go +++ b/pkg/sql/plan/query_builder.go @@ -3557,26 +3557,24 @@ func rewriteRollupWindowSelect( } rewriteState := newRollupWindowRewriteState(selectClause.Exprs) - if selectClause.Having != nil && !rewriteState.addGroupedSourceExprs(selectClause.GroupBy) { - return nil, true - } outerExprs, ok := buildRollupWindowSelectExprs(selectClause.Exprs, rewriteState) if !ok { return nil, true } var rewrittenHaving *tree.Where if selectClause.Having != nil { - // HAVING uses AliasAfterColumn: a materialized source expression wins - // when its name collides with a SELECT alias. - havingExpr, ok := rewriteRollupWindowExprWithFallbackNameAliases(selectClause.Having.Expr, rewriteState, rewriteState.havingAliases) + // Keep HAVING in the original FROM scope so AliasAfterColumn can resolve + // source columns (including ambiguity) before SELECT aliases. Hidden + // projections retain the original non-window aliases for the fallback. + rewriteState.addHavingAliasExprs(selectClause.Exprs) + nextHaving := *selectClause.Having + normalizedHaving, ok := selectClause.Having.Expr.Accept(rollupWindowHavingNameNormalizer{}) if !ok { return nil, true } - rewrittenHaving = &tree.Where{ - Type: selectClause.Having.Type, - RollupHaving: true, - Expr: havingExpr, - } + nextHaving.Expr = normalizedHaving + nextHaving.RollupHaving = true + rewrittenHaving = &nextHaving } rewrittenOrderBy, ok := rewriteRollupWindowOrderByWithNameAliases(astOrderBy, rewriteState, rewriteState.outputAliases) if !ok { @@ -3640,28 +3638,37 @@ func rewriteRollupWindowSelect( const rollupWindowInternalAliasPrefix = "__mo_rollup_window_col_" +type rollupWindowHavingNameNormalizer struct{} + +func (rollupWindowHavingNameNormalizer) Enter(expr tree.Expr) (tree.Expr, bool) { + name, ok := expr.(*tree.UnresolvedName) + if !ok || name.Star || name.NumParts == 0 { + return expr, false + } + next := *name + next.CStrParts[0] = tree.NewCStr(name.ColName(), 1) + return &next, true +} + +func (rollupWindowHavingNameNormalizer) Exit(expr tree.Expr) (tree.Expr, bool) { + return expr, true +} + type rollupWindowRewriteState struct { - innerExprs tree.SelectExprs - exprAliases map[string]string - sourceNameAliases map[string]string - ambiguousSourceNames map[string]struct{} - havingAliases map[string]string - outputAliases map[string]string - activeNameAliases map[string]string - nameAliasesAfterExpr bool - usedAliases map[string]struct{} - nextAlias int + innerExprs tree.SelectExprs + exprAliases map[string]string + outputAliases map[string]string + activeNameAliases map[string]string + usedAliases map[string]struct{} + nextAlias int } func newRollupWindowRewriteState(selectExprs tree.SelectExprs) *rollupWindowRewriteState { state := &rollupWindowRewriteState{ - innerExprs: make(tree.SelectExprs, 0, len(selectExprs)), - exprAliases: make(map[string]string), - sourceNameAliases: make(map[string]string), - ambiguousSourceNames: make(map[string]struct{}), - havingAliases: make(map[string]string), - outputAliases: make(map[string]string), - usedAliases: make(map[string]struct{}), + innerExprs: make(tree.SelectExprs, 0, len(selectExprs)), + exprAliases: make(map[string]string), + outputAliases: make(map[string]string), + usedAliases: make(map[string]struct{}), } for _, selectExpr := range selectExprs { if selectExpr.As != nil && !selectExpr.As.Empty() { @@ -3679,20 +3686,16 @@ func newRollupWindowRewriteState(selectExprs tree.SelectExprs) *rollupWindowRewr return state } -func (state *rollupWindowRewriteState) addGroupedSourceExprs(groupBy *tree.GroupByClause) bool { - // A grouped column may only be referenced by HAVING. Materialize grouped - // names so AliasAfterColumn can still choose their terminal column name - // ahead of SELECT aliases. - for _, groupByExprs := range groupBy.GroupByExprsList { - for _, expr := range groupByExprs { - if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star { - if _, ok := state.ensureInnerExpr(expr); !ok { - return false - } - } +func (state *rollupWindowRewriteState) addHavingAliasExprs(selectExprs tree.SelectExprs) { + for _, selectExpr := range selectExprs { + if selectExpr.As == nil || selectExpr.As.Empty() || rollupWindowExprContainsWindow(selectExpr.Expr) { + continue } + state.innerExprs = append(state.innerExprs, tree.SelectExpr{ + Expr: selectExpr.Expr, + As: selectExpr.As, + }) } - return true } func (state *rollupWindowRewriteState) newInternalAlias() string { @@ -3730,25 +3733,6 @@ func (state *rollupWindowRewriteState) addExprAlias(expr tree.Expr, alias string if _, exists := state.exprAliases[key]; !exists { state.exprAliases[key] = alias } - if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star { - state.addSourceNameAlias(unresolvedName.ColNameOrigin(), alias) - } -} - -func (state *rollupWindowRewriteState) addSourceNameAlias(name, alias string) { - key := strings.ToLower(name) - if key == "" { - return - } - if _, ambiguous := state.ambiguousSourceNames[key]; ambiguous { - return - } - if existing, exists := state.sourceNameAliases[key]; exists && existing != alias { - delete(state.sourceNameAliases, key) - state.ambiguousSourceNames[key] = struct{}{} - return - } - state.sourceNameAliases[key] = alias } func addRollupWindowNameAlias(aliases map[string]string, name, alias string) { @@ -3763,20 +3747,6 @@ func addRollupWindowNameAlias(aliases map[string]string, name, alias string) { func (state *rollupWindowRewriteState) lookupExprAlias(expr tree.Expr) (string, bool) { exprKey := rollupWindowExprKey(expr) - if state.nameAliasesAfterExpr { - if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star && unresolvedName.NumParts == 1 { - name := strings.ToLower(unresolvedName.ColNameOrigin()) - if _, ambiguous := state.ambiguousSourceNames[name]; ambiguous { - return "", false - } - if alias, exists := state.sourceNameAliases[name]; exists { - return alias, true - } - } - if alias, ok := state.exprAliases[exprKey]; ok { - return alias, true - } - } if state.activeNameAliases != nil { if unresolvedName, ok := expr.(*tree.UnresolvedName); ok && !unresolvedName.Star && unresolvedName.NumParts == 1 { if alias, exists := state.activeNameAliases[strings.ToLower(unresolvedName.ColNameOrigin())]; exists { @@ -3808,9 +3778,6 @@ func buildRollupWindowSelectExprs(selectExprs tree.SelectExprs, state *rollupWin if !ok { return nil, false } - if selectExpr.As != nil && !selectExpr.As.Empty() { - addRollupWindowNameAlias(state.havingAliases, selectExpr.As.Origin(), alias) - } aliasesByIndex[i] = alias } @@ -4097,27 +4064,9 @@ func rewriteRollupWindowExprWithNameAliases( aliases map[string]string, ) (tree.Expr, bool) { previousAliases := state.activeNameAliases - previousAfterExpr := state.nameAliasesAfterExpr - state.activeNameAliases = aliases - state.nameAliasesAfterExpr = false - rewritten, ok := rewriteRollupWindowExpr(expr, state) - state.activeNameAliases = previousAliases - state.nameAliasesAfterExpr = previousAfterExpr - return rewritten, ok -} - -func rewriteRollupWindowExprWithFallbackNameAliases( - expr tree.Expr, - state *rollupWindowRewriteState, - aliases map[string]string, -) (tree.Expr, bool) { - previousAliases := state.activeNameAliases - previousAfterExpr := state.nameAliasesAfterExpr state.activeNameAliases = aliases - state.nameAliasesAfterExpr = true rewritten, ok := rewriteRollupWindowExpr(expr, state) state.activeNameAliases = previousAliases - state.nameAliasesAfterExpr = previousAfterExpr return rewritten, ok } @@ -4230,12 +4179,9 @@ func rewriteRollupWindowOrderByWithNameAliases( aliases map[string]string, ) (tree.OrderBy, bool) { previousAliases := state.activeNameAliases - previousAfterExpr := state.nameAliasesAfterExpr state.activeNameAliases = aliases - state.nameAliasesAfterExpr = false rewritten, ok := rewriteRollupWindowOrderBy(orderBy, state) state.activeNameAliases = previousAliases - state.nameAliasesAfterExpr = previousAfterExpr return rewritten, ok } From 766f6dd1d390826355829570de6257059f692b85 Mon Sep 17 00:00:00 2001 From: GreatRiver Date: Wed, 15 Jul 2026 11:47:15 +0800 Subject: [PATCH 9/9] fix(plan): remove unused rollup window helper --- pkg/sql/plan/query_builder.go | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/pkg/sql/plan/query_builder.go b/pkg/sql/plan/query_builder.go index 793701fc4c4ea..7e79712b56af9 100644 --- a/pkg/sql/plan/query_builder.go +++ b/pkg/sql/plan/query_builder.go @@ -3735,16 +3735,6 @@ func (state *rollupWindowRewriteState) addExprAlias(expr tree.Expr, alias string } } -func addRollupWindowNameAlias(aliases map[string]string, name, alias string) { - key := strings.ToLower(name) - if key == "" { - return - } - if _, exists := aliases[key]; !exists { - aliases[key] = alias - } -} - func (state *rollupWindowRewriteState) lookupExprAlias(expr tree.Expr) (string, bool) { exprKey := rollupWindowExprKey(expr) if state.activeNameAliases != nil {