diff --git a/pkg/sql/plan/bind_context.go b/pkg/sql/plan/bind_context.go index ec29e7953f787..65cd5155e58cb 100644 --- a/pkg/sql/plan/bind_context.go +++ b/pkg/sql/plan/bind_context.go @@ -544,6 +544,10 @@ func (bc *BindContext) qualifyColumnNames(astExpr tree.Expr, expandAlias ExpandA exprImpl.To, err = bc.qualifyColumnNames(exprImpl.To, expandAlias) case *tree.UnresolvedName: + if havingBinder, ok := bc.binder.(*HavingBinder); ok && havingBinder.rollupHaving && + !exprImpl.Star && exprImpl.NumParts > 0 { + exprImpl.CStrParts[0] = tree.NewCStr(exprImpl.ColName(), 1) + } if !exprImpl.Star && exprImpl.NumParts == 1 { col := exprImpl.ColName() if expandAlias == AliasBeforeColumn { diff --git a/pkg/sql/plan/build_test.go b/pkg/sql/plan/build_test.go index 131ad35e34bea..1efe1759398cb 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" @@ -872,6 +873,632 @@ func TestSingleTableSQLBuilder(t *testing.T) { runTestShouldError(mock, t, sqls) } +func TestRollupWindowRanksAfterRollupUnion(t *testing.T) { + mock := NewMockOptimizer(false) + 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) + + 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) + }) + } +} + +func TestRollupWindowAliasCollisionsPreserveSourceScope(t *testing.T) { + mock := NewMockOptimizer(false) + tests := []struct { + name string + sql string + expectedHeadings []string + expectedProjectLen int + expectedHavingType int32 + }{ + { + 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"}, + expectedProjectLen: 4, + }, + { + 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"}, + 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, + }, + { + 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), + }, + { + 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 { + 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, test.expectedProjectLen) + 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) + + 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 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", + }, + } + + 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 TestRollupWindowHavingAvoidsUnsafeASTVisitor(t *testing.T) { + for _, test := range []struct { + name string + sql string + }{ + { + name: "searched case", + sql: ` + select n_regionkey, + row_number() over (order by n_regionkey) as rn + from nation + group by n_regionkey with rollup + having case when n_regionkey > 0 then 1 else 0 end = 1`, + }, + { + name: "prepared parameter", + sql: ` + prepare stmt1 from select n_regionkey, + row_number() over (order by n_regionkey) as rn + from nation + group by n_regionkey with rollup + having sum(n_nationkey) > ?`, + }, + { + name: "case insensitive name in expression", + sql: ` + select n_regionkey, + row_number() over (order by n_regionkey) as rn + from nation + group by n_regionkey with rollup + having abs(N_REGIONKEY) >= 0`, + }, + } { + t.Run(test.name, func(t *testing.T) { + mock := NewMockOptimizer(false) + _, err := runOneStmt(mock, t, test.sql) + require.NoError(t, err) + }) + } +} + +func TestRollupWindowVolatileExprOccurrencesStayIndependent(t *testing.T) { + stmt := mustParseRollupWindowSelect(t, ` + select rand() as r1, rand() as r2, + row_number() over (order by rand()) as rn + from nation + group by n_regionkey with rollup`) + 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) + + outerClause := rewritten.Select.(*tree.SelectClause) + require.Len(t, outerClause.Exprs, 3) + firstRand, ok := outerClause.Exprs[0].Expr.(*tree.UnresolvedName) + require.True(t, ok) + secondRand, ok := outerClause.Exprs[1].Expr.(*tree.UnresolvedName) + require.True(t, ok) + require.NotEqual(t, firstRand.ColNameOrigin(), secondRand.ColNameOrigin()) + + rowNumber, ok := outerClause.Exprs[2].Expr.(*tree.FuncExpr) + require.True(t, ok) + require.NotNil(t, rowNumber.WindowSpec) + require.Len(t, rowNumber.WindowSpec.OrderBy, 1) + _, ok = rowNumber.WindowSpec.OrderBy[0].Expr.(*tree.FuncExpr) + require.True(t, ok, "the window's rand occurrence must not reuse either SELECT projection") +} + +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.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") + + state := newRollupWindowRewriteState(clause.Exprs) + outerExprs, ok := buildRollupWindowSelectExprs(clause.Exprs, state) + require.True(t, ok) + require.Len(t, state.innerExprs, 4) + 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 := state.lookupExprAlias(clause.Exprs[0].Expr) + require.True(t, ok) + require.True(t, strings.HasPrefix(alias, rollupWindowInternalAliasPrefix)) + _, ok = state.lookupExprAlias(tree.NewUnresolvedColName("l_returnflag")) + require.False(t, ok) + _, ok = state.lookupExprAlias(tree.NewUnresolvedColName("flag")) + require.False(t, ok) + + state.addHavingAliasExprs(clause.Exprs) + require.Len(t, state.innerExprs, 7) + require.Equal(t, "flag", state.innerExprs[4].As.Origin()) + require.Equal(t, "status", state.innerExprs[5].As.Origin()) + require.Equal(t, "total_qty", state.innerExprs[6].As.Origin()) +} + +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.True(t, ok) + + distinct := mustParseRollupWindowSelect(t, "select distinct a, row_number() over () from t group by a, b with rollup") + 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") + 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.True(t, ok) + + 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 = rewriteRollupWindowSelect(noWindow.Select.(*tree.SelectClause), noWindow.OrderBy, noWindow.Limit, noWindow.RankOption) + 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, 2) + aggregateAlias := state.innerExprs[0].As.Origin() + sourceAlias := state.innerExprs[1].As.Origin() + 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 + 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) { + state := 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, state) + 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 + )`) + windowState := newRollupWindowRewriteState(nil) + rewritten, ok = rewriteRollupWindowExpr(windowExpr, windowState) + require.True(t, ok) + rewrittenSQL = tree.String(rewritten, dialect.MYSQL) + require.Len(t, windowState.innerExprs, 3) + require.Contains(t, rewrittenSQL, windowState.innerExprs[0].As.Origin()) + require.Contains(t, rewrittenSQL, "partition by "+windowState.innerExprs[1].As.Origin()) + require.Contains(t, rewrittenSQL, "order by "+windowState.innerExprs[2].As.Origin()) + + state = newRollupWindowRewriteState(nil) + aggregateExpr := mustParseRollupWindowExpr(t, "sum(missing)") + rewritten, ok = rewriteRollupWindowExpr(aggregateExpr, state) + require.True(t, ok) + require.True(t, strings.HasPrefix(tree.String(rewritten, dialect.MYSQL), rollupWindowInternalAliasPrefix)) + require.Len(t, state.innerExprs, 1) + + state = rollupWindowAliasMap("a", "b") + 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, state) + require.True(t, ok) + require.Nil(t, emptyExprs) + + state = newRollupWindowRewriteState(nil) + orderBy, ok := rewriteRollupWindowOrderBy(tree.OrderBy{{Expr: tree.NewUnresolvedColName("another_missing")}}, state) + require.True(t, ok) + require.Len(t, orderBy, 1) + require.True(t, strings.HasPrefix(tree.String(orderBy[0].Expr, dialect.MYSQL), rollupWindowInternalAliasPrefix)) +} + +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: "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"}, + } { + 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) *rollupWindowRewriteState { + state := newRollupWindowRewriteState(nil) + state.activeNameAliases = make(map[string]string, len(cols)) + for _, col := range cols { + state.activeNameAliases[col] = col + "_alias" + } + return state +} + +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{ ` diff --git a/pkg/sql/plan/query_builder.go b/pkg/sql/plan/query_builder.go index f9b8906fa52e5..ab8d46e6ed93f 100644 --- a/pkg/sql/plan/query_builder.go +++ b/pkg/sql/plan/query_builder.go @@ -19,6 +19,7 @@ import ( "encoding/json" "fmt" "maps" + "reflect" "slices" "strconv" "strings" @@ -3505,6 +3506,13 @@ func (builder *QueryBuilder) bindSelect(stmt *tree.Select, ctx *BindContext, isR } } if len(selectClause.GroupBy.GroupByExprsList) > 1 && !selectClause.GroupBy.Apart { + 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) + } + groupingCount := len(selectClause.GroupBy.GroupByExprsList) selectStmts := make([]*tree.SelectClause, groupingCount) if groupingCount > 1 { @@ -3740,6 +3748,724 @@ 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 { + return nil, false + } + + 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 + } + + rewriteState := newRollupWindowRewriteState(selectClause.Exprs) + outerExprs, ok := buildRollupWindowSelectExprs(selectClause.Exprs, rewriteState) + if !ok { + return nil, true + } + var rewrittenHaving *tree.Where + if selectClause.Having != nil { + // 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 + nextHaving.RollupHaving = true + rewrittenHaving = &nextHaving + } + rewrittenOrderBy, ok := rewriteRollupWindowOrderByWithNameAliases(astOrderBy, rewriteState, rewriteState.outputAliases) + 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{ + Exprs: append(tree.SelectExprs(nil), rewriteState.innerExprs...), + From: selectClause.From, + Where: selectClause.Where, + GroupBy: &tree.GroupByClause{ + GroupByExprsList: selectClause.GroupBy.GroupByExprsList, + GroupingSet: list, + Apart: true, + Cube: false, + Rollup: false, + }, + Having: rewrittenHaving, + 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{ + Distinct: selectClause.Distinct, + Exprs: outerExprs, + From: &tree.From{ + Tables: tree.TableExprs{ + &tree.AliasedTableExpr{ + Expr: tree.NewParenTableExpr(derived), + As: tree.AliasClause{ + Alias: tree.Identifier("__mo_rollup_window"), + }, + }, + }, + }, + }, + OrderBy: rewrittenOrderBy, + Limit: astLimit, + RankOption: astRankOption, + }, true +} + +const rollupWindowInternalAliasPrefix = "__mo_rollup_window_col_" + +type rollupWindowRewriteState struct { + innerExprs tree.SelectExprs + exprAliases map[tree.Expr]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[tree.Expr]string), + outputAliases: make(map[string]string), + usedAliases: make(map[string]struct{}), + } + for _, selectExpr := range selectExprs { + if selectExpr.As != nil && !selectExpr.As.Empty() { + 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{}{} + } + } + return state +} + +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, + }) + } +} + +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) + return alias, true +} + +func (state *rollupWindowRewriteState) addExprAlias(expr tree.Expr, alias string) { + // Reuse only the exact AST occurrence. Structural keys such as the SQL + // text would incorrectly merge independent volatile calls like rand(). + if exprType := reflect.TypeOf(expr); exprType != nil && exprType.Comparable() { + state.exprAliases[expr] = alias + } +} + +func (state *rollupWindowRewriteState) lookupExprAlias(expr tree.Expr) (string, bool) { + 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 + } + } + } + if exprType := reflect.TypeOf(expr); exprType == nil || !exprType.Comparable() { + return "", false + } + alias, ok := state.exprAliases[expr] + return alias, ok +} + +func buildRollupWindowSelectExprs(selectExprs tree.SelectExprs, state *rollupWindowRewriteState) (tree.SelectExprs, bool) { + if len(selectExprs) == 0 { + return nil, false + } + + outerExprs := make(tree.SelectExprs, 0, len(selectExprs)) + aliasesByIndex := make([]string, len(selectExprs)) + for i, selectExpr := range selectExprs { + if rollupWindowExprContainsWindow(selectExpr.Expr) { + continue + } + + alias, ok := state.ensureInnerExpr(selectExpr.Expr) + if !ok { + return nil, false + } + aliasesByIndex[i] = alias + } + + for i, selectExpr := range selectExprs { + if rollupWindowExprContainsWindow(selectExpr.Expr) { + rewrittenExpr, ok := rewriteRollupWindowExpr(selectExpr.Expr, state) + if !ok { + return nil, false + } + outerExprs = append(outerExprs, tree.SelectExpr{ + Expr: rewrittenExpr, + As: rollupWindowOutputAlias(selectExpr), + }) + continue + } + + outerExprs = append(outerExprs, tree.SelectExpr{ + Expr: tree.NewUnresolvedColName(aliasesByIndex[i]), + As: rollupWindowOutputAlias(selectExpr), + }) + } + + return outerExprs, true +} + +func rollupWindowOutputAlias(selectExpr tree.SelectExpr) *tree.CStr { + if selectExpr.As != nil && !selectExpr.As.Empty() { + return selectExpr.As + } + if colName, ok := rollupWindowBareColumnName(selectExpr.Expr); ok { + return tree.NewCStr(colName, 1) + } + expr := selectExpr.Expr + for { + parenExpr, ok := expr.(*tree.ParenExpr) + if !ok { + break + } + expr = parenExpr.Expr + } + if heading, ok := nameConstHeading(expr); ok { + return tree.NewCStr(heading, 1) + } + return tree.NewCStr(tree.String(expr, dialect.MYSQL), 1) +} + +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 rollupWindowExprIsStar(expr tree.Expr) bool { + switch e := expr.(type) { + case *tree.UnresolvedName: + return e.Star + case tree.UnqualifiedStar: + return true + default: + return false + } +} + +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 + } + + 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, state) + if !ok { + return nil, false + } + next.Exprs = exprs + } + if len(e.OrderBy) > 0 { + orderBy, ok := rewriteRollupWindowOrderBy(e.OrderBy, state) + if !ok { + return nil, false + } + next.OrderBy = orderBy + } + if e.WindowSpec != nil { + windowSpec, ok := rewriteRollupWindowSpec(e.WindowSpec, state) + if !ok { + return nil, false + } + next.WindowSpec = windowSpec + } + return &next, true + case *tree.BinaryExpr: + left, ok := rewriteRollupWindowExpr(e.Left, state) + if !ok { + return nil, false + } + right, ok := rewriteRollupWindowExpr(e.Right, state) + if !ok { + return nil, false + } + next := *e + next.Left = left + next.Right = right + return &next, true + case *tree.UnaryExpr: + child, ok := rewriteRollupWindowExpr(e.Expr, state) + 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, state) + if !ok { + return nil, false + } + next.Left = left + } + if e.Right != nil { + right, ok := rewriteRollupWindowExpr(e.Right, state) + if !ok { + return nil, false + } + next.Right = right + } + if e.Escape != nil { + escape, ok := rewriteRollupWindowExpr(e.Escape, state) + if !ok { + return nil, false + } + next.Escape = escape + } + return &next, true + case *tree.AndExpr: + left, ok := rewriteRollupWindowExpr(e.Left, state) + if !ok { + return nil, false + } + right, ok := rewriteRollupWindowExpr(e.Right, state) + if !ok { + return nil, false + } + next := *e + next.Left = left + next.Right = right + return &next, true + case *tree.XorExpr: + left, ok := rewriteRollupWindowExpr(e.Left, state) + if !ok { + return nil, false + } + right, ok := rewriteRollupWindowExpr(e.Right, state) + if !ok { + return nil, false + } + next := *e + next.Left = left + next.Right = right + return &next, true + case *tree.OrExpr: + left, ok := rewriteRollupWindowExpr(e.Left, state) + if !ok { + return nil, false + } + right, ok := rewriteRollupWindowExpr(e.Right, state) + if !ok { + return nil, false + } + next := *e + next.Left = left + next.Right = right + return &next, true + case *tree.NotExpr: + child, ok := rewriteRollupWindowExpr(e.Expr, state) + if !ok { + return nil, false + } + next := *e + next.Expr = child + return &next, true + case *tree.ParenExpr: + child, ok := rewriteRollupWindowExpr(e.Expr, state) + if !ok { + return nil, false + } + next := *e + next.Expr = child + return &next, true + case *tree.CastExpr: + child, ok := rewriteRollupWindowExpr(e.Expr, state) + if !ok { + return nil, false + } + next := *e + next.Expr = child + return &next, true + case *tree.Tuple: + exprs, ok := rewriteRollupWindowExprs(e.Exprs, state) + 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, state) + 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, state) + if !ok { + return nil, false + } + val, ok := rewriteRollupWindowExpr(when.Val, state) + if !ok { + return nil, false + } + next.Whens[i] = &tree.When{Cond: cond, Val: val} + } + } + if e.Else != nil { + elseExpr, ok := rewriteRollupWindowExpr(e.Else, state) + if !ok { + return nil, false + } + next.Else = elseExpr + } + return &next, true + 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: + if rollupWindowExprContainsWindow(expr) { + return nil, false + } + alias, ok := state.ensureInnerExpr(expr) + if !ok { + return nil, false + } + return tree.NewUnresolvedColName(alias), true + } +} + +func rewriteRollupWindowExprWithNameAliases( + expr tree.Expr, + state *rollupWindowRewriteState, + aliases map[string]string, +) (tree.Expr, bool) { + previousAliases := state.activeNameAliases + state.activeNameAliases = aliases + rewritten, ok := rewriteRollupWindowExpr(expr, state) + state.activeNameAliases = previousAliases + return rewritten, ok +} + +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, 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, state) + if !ok { + return nil, false + } + rewrittenExprs[i] = rewrittenExpr + } + return rewrittenExprs, true +} + +func rewriteRollupWindowSpec(windowSpec *tree.WindowSpec, state *rollupWindowRewriteState) (*tree.WindowSpec, bool) { + next := *windowSpec + if len(windowSpec.PartitionBy) > 0 { + partitionBy, ok := rewriteRollupWindowExprs(windowSpec.PartitionBy, state) + if !ok { + return nil, false + } + next.PartitionBy = partitionBy + } + if len(windowSpec.OrderBy) > 0 { + orderBy, ok := rewriteRollupWindowOrderBy(windowSpec.OrderBy, state) + if !ok { + return nil, false + } + next.OrderBy = orderBy + } + if windowSpec.Frame != nil { + frame, ok := rewriteRollupWindowFrame(windowSpec.Frame, state) + if !ok { + return nil, false + } + next.Frame = frame + } + return &next, true +} + +func rewriteRollupWindowFrame(frame *tree.FrameClause, state *rollupWindowRewriteState) (*tree.FrameClause, bool) { + next := *frame + if frame.Start != nil { + start, ok := rewriteRollupWindowFrameBound(frame.Start, state) + if !ok { + return nil, false + } + next.Start = start + } + if frame.End != nil { + end, ok := rewriteRollupWindowFrameBound(frame.End, state) + if !ok { + return nil, false + } + next.End = end + } + return &next, true +} + +func rewriteRollupWindowFrameBound(bound *tree.FrameBound, state *rollupWindowRewriteState) (*tree.FrameBound, bool) { + next := *bound + if bound.Expr != nil { + expr, ok := rewriteRollupWindowExpr(bound.Expr, state) + if !ok { + return nil, false + } + next.Expr = expr + } + return &next, true +} + +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, state) + if !ok { + return nil, false + } + next := *order + next.Expr = expr + rewrittenOrderBy[i] = &next + } + return rewrittenOrderBy, true +} + +func rewriteRollupWindowOrderByWithNameAliases( + orderBy tree.OrderBy, + state *rollupWindowRewriteState, + aliases map[string]string, +) (tree.OrderBy, bool) { + previousAliases := state.activeNameAliases + state.activeNameAliases = aliases + rewritten, ok := rewriteRollupWindowOrderBy(orderBy, state) + state.activeNameAliases = previousAliases + return rewritten, ok +} + +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.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 + } + for _, when := range e.Whens { + if rollupWindowExprContainsWindow(when.Cond) || rollupWindowExprContainsWindow(when.Val) { + 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 +} + func (builder *QueryBuilder) bindSelectClause( ctx *BindContext, clause *tree.SelectClause, @@ -4080,6 +4806,7 @@ func (builder *QueryBuilder) bindHaving( err = moerr.NewParseErrorf(builder.GetContext(), "not support having in recursive cte: '%v'", tree.String(clause, dialect.MYSQL)) return } + havingBinder.rollupHaving = clause.RollupHaving ctx.binder = havingBinder return splitAndBindCondition(clause.Expr, AliasAfterColumn, ctx) } diff --git a/pkg/sql/plan/types.go b/pkg/sql/plan/types.go index 4a12c51911685..e1b6af6ed4b20 100644 --- a/pkg/sql/plan/types.go +++ b/pkg/sql/plan/types.go @@ -502,7 +502,8 @@ type GroupBinder struct { type HavingBinder struct { baseBinder - insideAgg bool + insideAgg bool + rollupHaving bool } type ProjectionBinder struct { diff --git a/test/distributed/cases/window/rollup.result b/test/distributed/cases/window/rollup.result index bc23b58dd4c44..2f4c38a21d0a0 100644 --- a/test/distributed/cases/window/rollup.result +++ b/test/distributed/cases/window/rollup.result @@ -893,4 +893,86 @@ 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 +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 +select +sum(r.qty) as qty, +row_number() over (order by sum(r.qty), count(*)) as row_num +from +rollup_window_sales r +group by +r.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 e58bc146cec9f..c78d6b9844fe1 100644 --- a/test/distributed/cases/window/rollup.sql +++ b/test/distributed/cases/window/rollup.sql @@ -592,4 +592,68 @@ 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; + +-- 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; + +-- HAVING keeps AliasAfterColumn precedence when an aggregate alias collides +-- with the grouped source column name +select + sum(r.qty) as qty, + row_number() over (order by sum(r.qty), count(*)) as row_num +from + rollup_window_sales r +group by + r.qty with rollup +having qty > 0 +order by row_num; +drop table rollup_window_sales; + +drop database rollup_test;