diff --git a/pkg/sql/colexec/aggexec/empty_result.go b/pkg/sql/colexec/aggexec/empty_result.go new file mode 100644 index 0000000000000..eb8799cd78eca --- /dev/null +++ b/pkg/sql/colexec/aggexec/empty_result.go @@ -0,0 +1,51 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package aggexec + +// EmptyResultKind describes the value produced by an aggregate group that has +// no input rows. Keep this contract with the aggregate executors so planner +// rewrites do not maintain a separate function-name whitelist. +type EmptyResultKind uint8 + +const ( + EmptyResultUnsupported EmptyResultKind = iota + EmptyResultNull + EmptyResultZero + EmptyResultAllBitsSet +) + +// GetEmptyResultKind returns the aggregate executor's empty-input contract. +// Unsupported means the executor has a non-scalar or otherwise non-trivial +// empty value that a planner rewrite must not synthesize. +func GetEmptyResultKind(aggID int64) EmptyResultKind { + switch aggID { + case AggIdOfCountColumn, AggIdOfCountStar, AggIdOfApproxCount, AggIdOfApproxCountDistinct: + return EmptyResultZero + case AggIdOfBitAnd: + return EmptyResultAllBitsSet + case AggIdOfBitOr, AggIdOfBitXor: + return EmptyResultZero + case AggIdOfAny, AggIdOfAvg, AggIdOfGroupConcat, AggIdOfJsonArrayAgg, + AggIdOfJsonObjectAgg, AggIdOfMax, AggIdOfMaxBy, AggIdOfMaxByNonNull, + AggIdOfMedian, AggIdOfMin, AggIdOfApproxPercentile, AggIdOfStdDevPop, + AggIdOfStdDevSample, AggIdOfSum, AggIdOfVarPop, AggIdOfVarSample: + return EmptyResultNull + case AggIdOfAvgTwCache, AggIdOfAvgTwResult, AggIdOfBitmapConstruct, + AggIdOfBitmapOr, AggIdOfHllAdd, AggIdOfHllMerge: + return EmptyResultUnsupported + default: + return EmptyResultUnsupported + } +} diff --git a/pkg/sql/colexec/aggexec/empty_result_test.go b/pkg/sql/colexec/aggexec/empty_result_test.go new file mode 100644 index 0000000000000..9cfdec9998875 --- /dev/null +++ b/pkg/sql/colexec/aggexec/empty_result_test.go @@ -0,0 +1,65 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package aggexec + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestGetEmptyResultKind(t *testing.T) { + for _, tt := range []struct { + name string + ids []int64 + want EmptyResultKind + }{ + { + name: "null", + ids: []int64{ + AggIdOfAny, AggIdOfAvg, AggIdOfGroupConcat, AggIdOfJsonArrayAgg, + AggIdOfJsonObjectAgg, AggIdOfMax, AggIdOfMaxBy, AggIdOfMaxByNonNull, + AggIdOfMedian, AggIdOfMin, AggIdOfApproxPercentile, AggIdOfStdDevPop, + AggIdOfStdDevSample, AggIdOfSum, AggIdOfVarPop, AggIdOfVarSample, + }, + want: EmptyResultNull, + }, + { + name: "zero", + ids: []int64{ + AggIdOfCountColumn, AggIdOfCountStar, AggIdOfApproxCount, + AggIdOfApproxCountDistinct, AggIdOfBitOr, AggIdOfBitXor, + }, + want: EmptyResultZero, + }, + {name: "all bits set", ids: []int64{AggIdOfBitAnd}, want: EmptyResultAllBitsSet}, + { + name: "unsupported", + ids: []int64{ + AggIdOfAvgTwCache, AggIdOfAvgTwResult, AggIdOfBitmapConstruct, + AggIdOfBitmapOr, AggIdOfHllAdd, AggIdOfHllMerge, + }, + want: EmptyResultUnsupported, + }, + } { + t.Run(tt.name, func(t *testing.T) { + for _, id := range tt.ids { + require.Equal(t, tt.want, GetEmptyResultKind(id), "aggregate ID %d", id) + } + }) + } + + require.Equal(t, EmptyResultUnsupported, GetEmptyResultKind(-1)) +} diff --git a/pkg/sql/plan/flatten_subquery.go b/pkg/sql/plan/flatten_subquery.go index 74d3e5dc58821..6bee2fb18a07a 100644 --- a/pkg/sql/plan/flatten_subquery.go +++ b/pkg/sql/plan/flatten_subquery.go @@ -15,10 +15,14 @@ package plan import ( + "bytes" + "math" + "github.com/matrixorigin/matrixone/pkg/catalog" "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/pb/plan" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/aggexec" "github.com/matrixorigin/matrixone/pkg/sql/plan/function" ) @@ -158,11 +162,12 @@ func (builder *QueryBuilder) flattenSubquery(nodeID int32, subquery *plan.Subque switch subquery.Typ { case plan.SubqueryRef_SCALAR: - var rewrite bool + var rewriteCount bool - // Uncorrelated subquery + // Preserve the legacy COUNT fallback for plan shapes that cannot use the + // more precise empty-input projection reconstruction below. if len(joinPreds) > 0 && builder.findAggrCount(subCtx.aggregates) { - rewrite = true + rewriteCount = true } if scalarExistential { @@ -194,6 +199,12 @@ func (builder *QueryBuilder) flattenSubquery(nodeID int32, subquery *plan.Subque joinType = plan.Node_LEFT } + postJoinProjection, finalizeProjection, err := + builder.prepareCorrelatedScalarAggregatePostJoinProjection(subID, subCtx, joinPreds) + if err != nil { + return nodeID, nil, err + } + nodeID = builder.appendNode(&plan.Node{ NodeType: plan.Node_JOIN, Children: []int32{nodeID, subID}, @@ -211,7 +222,9 @@ func (builder *QueryBuilder) flattenSubquery(nodeID int32, subquery *plan.Subque } retExpr := scalarMatch - if retExpr == nil { + if finalizeProjection { + retExpr = postJoinProjection + } else if retExpr == nil { retExpr = &plan.Expr{ Typ: subCtx.results[0].Typ, Expr: &plan.Expr_Col{ @@ -232,7 +245,7 @@ func (builder *QueryBuilder) flattenSubquery(nodeID int32, subquery *plan.Subque return 0, nil, err } } - if rewrite { + if !finalizeProjection && rewriteCount { argsType := make([]types.Type, 1) argsType[0] = makeTypeByPlan2Expr(retExpr) fGet, err := function.GetFunctionByName(builder.GetContext(), "isnull", argsType) @@ -597,6 +610,189 @@ func (builder *QueryBuilder) generateRowComparison(op string, child *plan.Expr, } } +// prepareCorrelatedScalarAggregatePostJoinProjection moves the scalar final +// expression above the LEFT JOIN used to decorrelate an implicit single-group +// aggregate. pullupThroughAgg groups the inner input by the correlation key, so +// a missing key produces no right row. Evaluating COALESCE, arithmetic, or CASE +// below the join would therefore skip the expression for that outer row. +// +// The right projection is rewritten to expose only raw aggregate outputs and +// the correlation keys that pullupThroughProj already appended. Aggregate +// outputs are restored from aggexec's canonical empty-input contract after null +// extension. The saved final expression is then evaluated against those +// post-join values. +// +// This is intentionally limited to a direct AGG or the ordinary PROJECT -> AGG +// shape. Wrappers that can remove or reorder the aggregate row (for example +// HAVING, DISTINCT, SORT, or LIMIT) keep the legacy path. +func (builder *QueryBuilder) prepareCorrelatedScalarAggregatePostJoinProjection( + subID int32, + subCtx *BindContext, + joinPreds []*plan.Expr, +) (*plan.Expr, bool, error) { + if !subCtx.hasSingleRow || len(subCtx.groups) != 0 || len(subCtx.aggregates) == 0 || len(joinPreds) == 0 { + return nil, false, nil + } + + project := builder.qry.Nodes[subID] + if project.NodeType == plan.Node_AGG { + if len(project.BindingTags) < 2 || project.BindingTags[1] != subCtx.aggregateTag || + len(project.AggList) != 1 || len(subCtx.aggregates) != 1 || len(subCtx.results) != 1 { + return nil, false, nil + } + aggregate := project.AggList[0] + fn := aggregate.GetF() + if fn == nil || fn.Func == nil { + return nil, false, nil + } + projected := GetColExpr(aggregate.Typ, subCtx.aggregateTag, 0) + projected.Typ.NotNullable = false + postJoinProjection, err := builder.restoreAggregateEmptyResult(projected, aggregate, fn.Func.ObjName) + return postJoinProjection, err == nil, err + } + if project.NodeType != plan.Node_PROJECT || len(project.Children) != 1 || len(project.BindingTags) != 1 || + len(project.ProjectList) == 0 || project.Limit != nil || project.Offset != nil || project.RankOption != nil { + return nil, false, nil + } + + agg := builder.qry.Nodes[project.Children[0]] + if agg.NodeType != plan.Node_AGG || len(agg.BindingTags) < 2 || agg.BindingTags[1] != subCtx.aggregateTag || + len(agg.AggList) != len(subCtx.aggregates) { + return nil, false, nil + } + + projectTag := project.BindingTags[0] + projectedAggregates := make([]*plan.Expr, len(agg.AggList)) + rawAggregates := make([]*plan.Expr, len(agg.AggList)) + firstAppendedPos := int32(len(project.ProjectList)) + for i, aggregate := range agg.AggList { + fn := aggregate.GetF() + if fn == nil || fn.Func == nil { + return nil, false, nil + } + + projectPos := int32(0) + if i > 0 { + projectPos = firstAppendedPos + int32(i-1) + } + rawAggregates[i] = GetColExpr(aggregate.Typ, subCtx.aggregateTag, int32(i)) + projected := GetColExpr(aggregate.Typ, projectTag, projectPos) + projected.Typ.NotNullable = false + + var err error + projectedAggregates[i], err = builder.restoreAggregateEmptyResult(projected, aggregate, fn.Func.ObjName) + if err != nil { + return nil, false, err + } + } + + postJoinProjection, ok := replaceAggregateRefsForPostJoin( + DeepCopyExpr(project.ProjectList[0]), subCtx.aggregateTag, projectedAggregates) + if !ok { + return nil, false, nil + } + postJoinProjection, stillCorrelated := decreaseDepth(postJoinProjection) + if stillCorrelated { + return nil, false, nil + } + + newProjectList := make([]*plan.Expr, len(project.ProjectList), len(project.ProjectList)+len(rawAggregates)-1) + copy(newProjectList, project.ProjectList) + newProjectList[0] = rawAggregates[0] + newProjectList = append(newProjectList, rawAggregates[1:]...) + project.ProjectList = newProjectList + return postJoinProjection, true, nil +} + +func (builder *QueryBuilder) restoreAggregateEmptyResult( + aggregateExpr *plan.Expr, + aggregate *plan.Expr, + aggregateName string, +) (*plan.Expr, error) { + kind := aggexec.GetEmptyResultKind(aggregate.GetF().Func.Obj) + if kind == aggexec.EmptyResultNull { + return aggregateExpr, nil + } + if kind == aggexec.EmptyResultUnsupported { + return nil, moerr.NewNYIf(builder.GetContext(), + "aggregate %s in a correlated scalar projection will be supported in a future version", aggregateName) + } + + isNullExpr, err := BindFuncExprImplByPlanExpr(builder.GetContext(), "isnull", []*plan.Expr{aggregateExpr}) + if err != nil { + return nil, err + } + emptyExpr, err := makeAggregateEmptyResultExpr(kind, aggregate.Typ) + if err != nil { + return nil, err + } + return BindFuncExprImplByPlanExpr(builder.GetContext(), "case", []*plan.Expr{ + isNullExpr, + emptyExpr, + DeepCopyExpr(aggregateExpr), + }) +} + +func makeAggregateEmptyResultExpr(kind aggexec.EmptyResultKind, aggregateType plan.Type) (*plan.Expr, error) { + var expr *plan.Expr + switch types.T(aggregateType.Id) { + case types.T_binary, types.T_varbinary: + fill := byte(0) + if kind == aggexec.EmptyResultAllBitsSet { + fill = 0xff + } + expr = makePlan2VarBinaryConstExprWithType(string(bytes.Repeat([]byte{fill}, int(aggregateType.Width)))) + case types.T_uint64: + value := uint64(0) + if kind == aggexec.EmptyResultAllBitsSet { + value = math.MaxUint64 + } + expr = makePlan2Uint64ConstExprWithType(value) + default: + if kind == aggexec.EmptyResultAllBitsSet { + return nil, moerr.NewInternalErrorNoCtxf("all-bits-set empty aggregate result has unsupported type %s", makeTypeByPlan2Expr(&plan.Expr{Typ: aggregateType})) + } + expr = makePlan2Int64ConstExprWithType(0) + } + expr.Typ = aggregateType + expr.Typ.NotNullable = true + return expr, nil +} + +func replaceAggregateRefsForPostJoin( + expr *plan.Expr, + aggregateTag int32, + projectedAggregates []*plan.Expr, +) (*plan.Expr, bool) { + if expr == nil { + return nil, false + } + + switch item := expr.Expr.(type) { + case *plan.Expr_Col: + if item.Col.RelPos != aggregateTag { + return nil, false + } + if item.Col.ColPos < 0 || int(item.Col.ColPos) >= len(projectedAggregates) { + return nil, false + } + return DeepCopyExpr(projectedAggregates[item.Col.ColPos]), true + case *plan.Expr_F: + for i, arg := range item.F.Args { + var ok bool + item.F.Args[i], ok = replaceAggregateRefsForPostJoin(arg, aggregateTag, projectedAggregates) + if !ok { + return nil, false + } + } + return expr, true + case *plan.Expr_List, *plan.Expr_W, *plan.Expr_Sub: + return nil, false + default: + return expr, true + } +} + func (builder *QueryBuilder) findAggrCount(aggrs []*plan.Expr) bool { for _, aggr := range aggrs { switch exprImpl := aggr.Expr.(type) { diff --git a/pkg/sql/plan/flatten_subquery_test.go b/pkg/sql/plan/flatten_subquery_test.go index 9e328714a1f27..9c0644cbc3f76 100644 --- a/pkg/sql/plan/flatten_subquery_test.go +++ b/pkg/sql/plan/flatten_subquery_test.go @@ -15,10 +15,12 @@ package plan import ( + "math" "testing" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/pb/plan" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/aggexec" "github.com/stretchr/testify/require" ) @@ -220,6 +222,425 @@ func TestDirectCorrelatedScalarProjectionUsesMatchMarker(t *testing.T) { } } +func TestCorrelatedScalarAggregateProjectionRunsAfterJoin(t *testing.T) { + logicPlan, err := runOneStmt(NewMockOptimizer(true), t, + "select n.n_nationkey, (select coalesce(sum(r.r_regionkey), 0) from tpch.region r where r.r_regionkey = n.n_regionkey) as total from tpch.nation n") + require.NoError(t, err) + + query := logicPlan.GetQuery() + require.NotNil(t, query) + + var rightAggregate *plan.Node + for _, node := range query.Nodes { + if node.NodeType != plan.Node_JOIN || node.JoinType != plan.Node_LEFT || len(node.Children) != 2 { + continue + } + candidate := query.Nodes[node.Children[1]] + if candidate.NodeType != plan.Node_AGG || len(candidate.AggList) == 0 { + continue + } + rightAggregate = candidate + break + } + + require.NotNil(t, rightAggregate) + require.Equal(t, "sum", rightAggregate.AggList[0].GetF().Func.GetObjName()) + assertReachablePlanHasNoCorrelatedExpr(t, query) +} + +func assertReachablePlanHasNoCorrelatedExpr(t *testing.T, query *plan.Query) { + t.Helper() + visited := make(map[int32]bool) + var visit func(int32) + visit = func(nodeID int32) { + if visited[nodeID] { + return + } + visited[nodeID] = true + node := query.Nodes[nodeID] + exprs := make([]*plan.Expr, 0, len(node.ProjectList)+len(node.FilterList)+len(node.OnList)+len(node.GroupBy)+len(node.AggList)+len(node.WinSpecList)+2) + exprs = append(exprs, node.ProjectList...) + exprs = append(exprs, node.FilterList...) + exprs = append(exprs, node.OnList...) + exprs = append(exprs, node.GroupBy...) + exprs = append(exprs, node.AggList...) + exprs = append(exprs, node.WinSpecList...) + for _, order := range node.OrderBy { + if order != nil { + exprs = append(exprs, order.Expr) + } + } + if node.Limit != nil { + exprs = append(exprs, node.Limit) + } + if node.Offset != nil { + exprs = append(exprs, node.Offset) + } + for _, expr := range exprs { + require.False(t, hasCorrCol(expr), "reachable %s node %d contains a correlated expression", node.NodeType.String(), nodeID) + } + for _, child := range node.Children { + visit(child) + } + } + for _, root := range query.Steps { + visit(root) + } +} + +func TestCorrelatedScalarAggregatePostJoinProjectionEligibility(t *testing.T) { + for _, tt := range []struct { + name string + sql string + want bool + }{ + { + name: "ifnull aggregate", + sql: "select n.n_nationkey, (select ifnull(avg(r.r_regionkey), 7) from tpch.region r where r.r_regionkey = n.n_regionkey) from tpch.nation n", + want: true, + }, + { + name: "mixed sum and count", + sql: "select n.n_nationkey, (select sum(r.r_regionkey) + count(*) from tpch.region r where r.r_regionkey = n.n_regionkey) from tpch.nation n", + want: true, + }, + { + name: "null-safe equality", + sql: "select n.n_nationkey, (select coalesce(max(r.r_regionkey), 0) from tpch.region r where r.r_regionkey <=> n.n_regionkey) from tpch.nation n", + want: true, + }, + { + name: "cte aggregate input", + sql: "with r as (select r_regionkey from tpch.region) select n.n_nationkey, (select coalesce(min(r.r_regionkey), 0) from r where r.r_regionkey = n.n_regionkey) from tpch.nation n", + want: true, + }, + { + name: "json aggregate", + sql: "select n.n_nationkey, (select coalesce(json_arrayagg(r.r_regionkey), convert('[]', json)) from tpch.region r where r.r_regionkey = n.n_regionkey) from tpch.nation n", + want: true, + }, + { + name: "explicit group by", + sql: "select n.n_nationkey, (select coalesce(sum(r.r_regionkey), 0) from tpch.region r where r.r_regionkey = n.n_regionkey group by r.r_regionkey) from tpch.nation n", + }, + { + name: "having can remove aggregate row", + sql: "select n.n_nationkey, (select coalesce(sum(r.r_regionkey), 0) from tpch.region r where r.r_regionkey = n.n_regionkey having sum(r.r_regionkey) > 100) from tpch.nation n", + }, + { + name: "neutral aggregate", + sql: "select n.n_nationkey, (select coalesce(bit_or(r.r_regionkey), 0) from tpch.region r where r.r_regionkey = n.n_regionkey) from tpch.nation n", + want: true, + }, + { + name: "raw neutral aggregate", + sql: "select n.n_nationkey, (select bit_and(r.r_regionkey) from tpch.region r where r.r_regionkey = n.n_regionkey) from tpch.nation n", + want: true, + }, + { + name: "raw approximate count aggregate", + sql: "select n.n_nationkey, (select approx_count_distinct(r.r_regionkey) from tpch.region r where r.r_regionkey = n.n_regionkey) from tpch.nation n", + want: true, + }, + { + name: "limited aggregate", + sql: "select n.n_nationkey, (select coalesce(sum(r.r_regionkey), 0) from tpch.region r where r.r_regionkey = n.n_regionkey limit 1) from tpch.nation n", + }, + { + name: "sorted aggregate", + sql: "select n.n_nationkey, (select coalesce(sum(r.r_regionkey), 0) from tpch.region r where r.r_regionkey = n.n_regionkey order by sum(r.r_regionkey)) from tpch.nation n", + }, + } { + t.Run(tt.name, func(t *testing.T) { + logicPlan, err := runOneStmt(NewMockOptimizer(true), t, tt.sql) + require.NoError(t, err) + require.Equal(t, tt.want, hasCorrelatedAggregatePostJoinProjection(logicPlan.GetQuery())) + }) + } +} + +func TestPrepareCorrelatedScalarAggregatePostJoinProjection(t *testing.T) { + const ( + groupTag int32 = 10 + aggregateTag int32 = 11 + projectTag int32 = 12 + outerTag int32 = 13 + ) + + aggregateType := func(id types.T) plan.Type { + return plan.Type{Id: int32(id), NotNullable: true} + } + aggregates := []*plan.Expr{ + newFlattenSubqueryTestAggregate("sum", aggregateType(types.T_decimal128)), + newFlattenSubqueryTestAggregate("avg", aggregateType(types.T_float64)), + newFlattenSubqueryTestAggregate("min", aggregateType(types.T_int32)), + newFlattenSubqueryTestAggregate("max", aggregateType(types.T_int64)), + newFlattenSubqueryTestAggregate("json_arrayagg", aggregateType(types.T_json)), + newFlattenSubqueryTestAggregate("count", aggregateType(types.T_int64)), + newFlattenSubqueryTestAggregate("starcount", aggregateType(types.T_int64)), + } + projectionArgs := make([]*plan.Expr, 0, len(aggregates)+1) + for i, aggregate := range aggregates { + projectionArgs = append(projectionArgs, GetColExpr(aggregate.Typ, aggregateTag, int32(i))) + } + projectionArgs = append(projectionArgs, &plan.Expr{ + Typ: aggregateType(types.T_int64), + Expr: &plan.Expr_Corr{Corr: &plan.CorrColRef{ + RelPos: outerTag, + ColPos: 3, + Depth: 1, + }}, + }) + projection := &plan.Expr{ + Typ: aggregateType(types.T_int64), + Expr: &plan.Expr_F{F: &plan.Function{ + Func: &plan.ObjectRef{ObjName: "case"}, + Args: projectionArgs, + }}, + } + correlationKey := GetColExpr(plan.Type{Id: int32(types.T_int32)}, groupTag, 0) + + builder := NewQueryBuilder(plan.Query_SELECT, NewMockCompilerContext(true), false, true) + builder.qry.Nodes = []*plan.Node{ + { + NodeType: plan.Node_AGG, + AggList: aggregates, + BindingTags: []int32{groupTag, aggregateTag}, + }, + { + NodeType: plan.Node_PROJECT, + Children: []int32{0}, + BindingTags: []int32{projectTag}, + ProjectList: []*plan.Expr{projection, correlationKey}, + }, + } + ctx := &BindContext{ + hasSingleRow: true, + aggregateTag: aggregateTag, + aggregates: aggregates, + } + + postJoinProjection, ok, err := builder.prepareCorrelatedScalarAggregatePostJoinProjection(1, ctx, []*plan.Expr{constTrue}) + require.NoError(t, err) + require.True(t, ok) + require.Len(t, builder.qry.Nodes[1].ProjectList, len(aggregates)+1) + require.Equal(t, groupTag, builder.qry.Nodes[1].ProjectList[1].GetCol().RelPos) + require.Equal(t, int32(0), builder.qry.Nodes[1].ProjectList[1].GetCol().ColPos) + rawPositions := []int{0, 2, 3, 4, 5, 6, 7} + for i, pos := range rawPositions { + raw := builder.qry.Nodes[1].ProjectList[pos] + require.Equal(t, aggregateTag, raw.GetCol().RelPos) + require.Equal(t, int32(i), raw.GetCol().ColPos) + } + + postJoinArgs := postJoinProjection.GetF().Args + require.Len(t, postJoinArgs, len(aggregates)+1) + for i := 0; i < 5; i++ { + require.Equal(t, projectTag, postJoinArgs[i].GetCol().RelPos) + projectPos := int32(i + 1) + if i == 0 { + projectPos = 0 + } + require.Equal(t, projectPos, postJoinArgs[i].GetCol().ColPos) + require.False(t, postJoinArgs[i].Typ.NotNullable) + } + for i := 5; i < 7; i++ { + countFallback := postJoinArgs[i].GetF() + require.Equal(t, "case", countFallback.Func.GetObjName()) + require.Equal(t, projectTag, countFallback.Args[2].GetCol().RelPos) + require.Equal(t, int32(i+1), countFallback.Args[2].GetCol().ColPos) + } + require.Nil(t, postJoinArgs[7].GetCorr()) + require.Equal(t, outerTag, postJoinArgs[7].GetCol().RelPos) + require.Equal(t, int32(3), postJoinArgs[7].GetCol().ColPos) +} + +func TestMakeAggregateEmptyResultExpr(t *testing.T) { + for _, tt := range []struct { + name string + kind aggexec.EmptyResultKind + typ plan.Type + want any + }{ + {name: "uint64 zero", kind: aggexec.EmptyResultZero, typ: plan.Type{Id: int32(types.T_uint64)}, want: uint64(0)}, + {name: "uint64 all bits set", kind: aggexec.EmptyResultAllBitsSet, typ: plan.Type{Id: int32(types.T_uint64)}, want: uint64(math.MaxUint64)}, + {name: "binary zero", kind: aggexec.EmptyResultZero, typ: plan.Type{Id: int32(types.T_binary), Width: 3}, want: string([]byte{0, 0, 0})}, + {name: "varbinary all bits set", kind: aggexec.EmptyResultAllBitsSet, typ: plan.Type{Id: int32(types.T_varbinary), Width: 3}, want: string([]byte{0xff, 0xff, 0xff})}, + {name: "int64 zero", kind: aggexec.EmptyResultZero, typ: plan.Type{Id: int32(types.T_int64)}, want: int64(0)}, + } { + t.Run(tt.name, func(t *testing.T) { + expr, err := makeAggregateEmptyResultExpr(tt.kind, tt.typ) + require.NoError(t, err) + require.Equal(t, tt.typ.Id, expr.Typ.Id) + require.Equal(t, tt.typ.Width, expr.Typ.Width) + require.True(t, expr.Typ.NotNullable) + switch want := tt.want.(type) { + case uint64: + require.Equal(t, want, expr.GetLit().GetU64Val()) + case int64: + require.Equal(t, want, expr.GetLit().GetI64Val()) + case string: + require.Equal(t, want, expr.GetLit().GetSval()) + } + }) + } + + _, err := makeAggregateEmptyResultExpr(aggexec.EmptyResultAllBitsSet, plan.Type{Id: int32(types.T_int64)}) + require.Error(t, err) +} + +func TestPrepareCorrelatedScalarAggregatePostJoinProjectionRejectsUnsupportedShapes(t *testing.T) { + const ( + aggregateTag int32 = 21 + projectTag int32 = 22 + ) + + for _, tt := range []struct { + name string + aggregate string + wantErr bool + mutate func(*BindContext, []*plan.Node) + }{ + {name: "unsupported aggregate", aggregate: "hll_add_agg", wantErr: true}, + { + name: "explicit group", + aggregate: "sum", + mutate: func(ctx *BindContext, _ []*plan.Node) { + ctx.groups = []*plan.Expr{makePlan2Int64ConstExprWithType(1)} + }, + }, + { + name: "having filter", + aggregate: "sum", + mutate: func(_ *BindContext, nodes []*plan.Node) { + nodes[1].Children[0] = 2 + }, + }, + { + name: "limit", + aggregate: "sum", + mutate: func(_ *BindContext, nodes []*plan.Node) { + nodes[1].Limit = makePlan2Uint64ConstExprWithType(1) + }, + }, + { + name: "deep correlation", + aggregate: "sum", + mutate: func(_ *BindContext, nodes []*plan.Node) { + nodes[1].ProjectList[0] = &plan.Expr{ + Typ: plan.Type{Id: int32(types.T_int64)}, + Expr: &plan.Expr_Corr{Corr: &plan.CorrColRef{RelPos: 30, Depth: 2}}, + } + }, + }, + { + name: "non aggregate inner column", + aggregate: "sum", + mutate: func(_ *BindContext, nodes []*plan.Node) { + nodes[1].ProjectList[0] = GetColExpr(plan.Type{Id: int32(types.T_int64)}, 99, 0) + }, + }, + } { + t.Run(tt.name, func(t *testing.T) { + aggregate := newFlattenSubqueryTestAggregate(tt.aggregate, plan.Type{Id: int32(types.T_int64)}) + nodes := []*plan.Node{ + {NodeType: plan.Node_AGG, AggList: []*plan.Expr{aggregate}, BindingTags: []int32{20, aggregateTag}}, + { + NodeType: plan.Node_PROJECT, + Children: []int32{0}, + BindingTags: []int32{projectTag}, + ProjectList: []*plan.Expr{GetColExpr(aggregate.Typ, aggregateTag, 0)}, + }, + {NodeType: plan.Node_FILTER, Children: []int32{0}}, + } + ctx := &BindContext{hasSingleRow: true, aggregateTag: aggregateTag, aggregates: []*plan.Expr{aggregate}} + if tt.mutate != nil { + tt.mutate(ctx, nodes) + } + builder := NewQueryBuilder(plan.Query_SELECT, NewMockCompilerContext(true), false, true) + builder.qry.Nodes = nodes + + postJoinProjection, ok, err := builder.prepareCorrelatedScalarAggregatePostJoinProjection(1, ctx, []*plan.Expr{constTrue}) + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + require.False(t, ok) + require.Nil(t, postJoinProjection) + require.Len(t, builder.qry.Nodes[1].ProjectList, 1) + }) + } +} + +func TestPrepareCorrelatedScalarAggregatePostJoinProjectionRejectsUnsupportedDirectAggregate(t *testing.T) { + aggregate := newFlattenSubqueryTestAggregate("hll_add_agg", plan.Type{Id: int32(types.T_varbinary)}) + builder := NewQueryBuilder(plan.Query_SELECT, NewMockCompilerContext(true), false, true) + builder.qry.Nodes = []*plan.Node{{ + NodeType: plan.Node_AGG, + AggList: []*plan.Expr{aggregate}, + BindingTags: []int32{20, 21}, + }} + ctx := &BindContext{ + hasSingleRow: true, + aggregateTag: 21, + aggregates: []*plan.Expr{aggregate}, + results: []*plan.Expr{GetColExpr(aggregate.Typ, 21, 0)}, + } + + postJoinProjection, ok, err := builder.prepareCorrelatedScalarAggregatePostJoinProjection(0, ctx, []*plan.Expr{constTrue}) + require.Error(t, err) + require.False(t, ok) + require.Nil(t, postJoinProjection) +} + +func hasCorrelatedAggregatePostJoinProjection(query *plan.Query) bool { + if query == nil { + return false + } + for _, node := range query.Nodes { + if node.NodeType != plan.Node_JOIN || node.JoinType != plan.Node_LEFT || len(node.Children) != 2 { + continue + } + right := query.Nodes[node.Children[1]] + if right.NodeType == plan.Node_AGG { + return true + } + if right.NodeType != plan.Node_PROJECT || len(right.Children) != 1 || len(right.ProjectList) == 0 { + continue + } + agg := query.Nodes[right.Children[0]] + if agg.NodeType != plan.Node_AGG || len(agg.BindingTags) < 2 { + continue + } + first := right.ProjectList[0].GetCol() + if first != nil && first.RelPos == agg.BindingTags[1] && first.ColPos == 0 { + return true + } + } + return false +} + +func newFlattenSubqueryTestAggregate(name string, typ plan.Type) *plan.Expr { + ids := map[string]int64{ + "sum": aggexec.AggIdOfSum, + "avg": aggexec.AggIdOfAvg, + "min": aggexec.AggIdOfMin, + "max": aggexec.AggIdOfMax, + "json_arrayagg": aggexec.AggIdOfJsonArrayAgg, + "count": aggexec.AggIdOfCountColumn, + "starcount": aggexec.AggIdOfCountStar, + "bit_or": aggexec.AggIdOfBitOr, + "hll_add_agg": aggexec.AggIdOfHllAdd, + } + return &plan.Expr{ + Typ: typ, + Expr: &plan.Expr_F{F: &plan.Function{ + Func: &plan.ObjectRef{ObjName: name, Obj: ids[name]}, + }}, + } +} + func TestDirectCorrelatedScalarProjectionCasePreservesType(t *testing.T) { builder := NewQueryBuilder(plan.Query_SELECT, NewMockCompilerContext(true), false, true) diff --git a/test/distributed/cases/optimizer/associative.result b/test/distributed/cases/optimizer/associative.result index e4b84bf1c6a62..7a069708128b5 100644 --- a/test/distributed/cases/optimizer/associative.result +++ b/test/distributed/cases/optimizer/associative.result @@ -131,16 +131,15 @@ Project 𝄀 -> Aggregate 𝄀 Aggregate Functions: sum(lineitem.l_extendedprice) 𝄀 -> Filter 𝄀 - Filter Cond: (cast(lineitem.l_quantity AS DECIMAL128(38, 2)) < 0.2 * avg(lineitem.l_quantity)) + Filter Cond: (cast(lineitem.l_quantity AS DECIMAL128(38, 2)) < (0.2 * avg(lineitem.l_quantity))) -- HINT: Cast expression may prevent index usage 𝄀 -> Join 𝄀 Join Type: RIGHT 𝄀 Join Cond: (lineitem.l_partkey = part.p_partkey) 𝄀 - -> Project 𝄀 - -> Aggregate 𝄀 - Group Key: lineitem.l_partkey shuffle: range(lineitem.l_partkey) 𝄀 - Aggregate Functions: avg(lineitem.l_quantity) 𝄀 - -> Table Scan on associative.lineitem 𝄀 + -> Aggregate 𝄀 + Group Key: lineitem.l_partkey shuffle: range(lineitem.l_partkey) 𝄀 + Aggregate Functions: avg(lineitem.l_quantity) 𝄀 + -> Table Scan on associative.lineitem 𝄀 -> Join 𝄀 Join Type: INNER hashOnPK 𝄀 Join Cond: (lineitem.l_partkey = part.p_partkey) 𝄀 @@ -174,16 +173,15 @@ Project 𝄀 -> Aggregate 𝄀 Aggregate Functions: sum(lineitem.l_extendedprice) 𝄀 -> Filter 𝄀 - Filter Cond: (cast(lineitem.l_quantity AS DECIMAL128(38, 2)) < 0.2 * avg(lineitem.l_quantity)) + Filter Cond: (cast(lineitem.l_quantity AS DECIMAL128(38, 2)) < (0.2 * avg(lineitem.l_quantity))) -- HINT: Cast expression may prevent index usage 𝄀 -> Join 𝄀 Join Type: RIGHT 𝄀 Join Cond: (lineitem.l_partkey = part.p_partkey) 𝄀 - -> Project 𝄀 - -> Aggregate 𝄀 - Group Key: lineitem.l_partkey 𝄀 - Aggregate Functions: avg(lineitem.l_quantity) 𝄀 - -> Table Scan on associative.lineitem 𝄀 + -> Aggregate 𝄀 + Group Key: lineitem.l_partkey 𝄀 + Aggregate Functions: avg(lineitem.l_quantity) 𝄀 + -> Table Scan on associative.lineitem 𝄀 -> Join 𝄀 Join Type: INNER hashOnPK 𝄀 Join Cond: (lineitem.l_partkey = part.p_partkey) 𝄀 diff --git a/test/distributed/cases/subquery/scalar_correlated_projection.result b/test/distributed/cases/subquery/scalar_correlated_projection.result index 9862217b5a595..fd14e5b3f53fd 100644 --- a/test/distributed/cases/subquery/scalar_correlated_projection.result +++ b/test/distributed/cases/subquery/scalar_correlated_projection.result @@ -4,6 +4,10 @@ use test_subq_corr_project; create table t1 (a int, b int, c int); create table t2 (d int); insert into t1 values (1, 2, 3), (11, 22, 33); +create table parent_agg (id int primary key, corr_key int); +create table child_agg (corr_key int, v int); +insert into parent_agg values (1, 10), (2, 20), (3, 30), (4, 30); +insert into child_agg values (10, 5), (10, 7), (20, null); select t1.*, (select t1.a from t2 where t2.d > t1.a) as x from t1 order by t1.a; ➤ a[4,32,0] ¦ b[4,32,0] ¦ c[4,32,0] ¦ x[4,32,0] 𝄀 1 ¦ 2 ¦ 3 ¦ null 𝄀 @@ -49,4 +53,79 @@ select (select t1.a from t2 where t2.d > t1.a limit 2) as x from t1 where t1.a = 11 select (select t1.a from t2 where t2.d > t1.a limit 2) as x from t1 where t1.a = 1; Subquery returns more than 1 row +select p.id, (select coalesce(sum(c.v), 0) from child_agg c where c.corr_key = p.corr_key) as sum_value from parent_agg p order by p.id; +➤ id[4,32,0] ¦ sum_value[-5,64,0] 𝄀 +1 ¦ 12 𝄀 +2 ¦ 0 𝄀 +3 ¦ 0 𝄀 +4 ¦ 0 +select p.id, (select ifnull(avg(c.v), 7) from child_agg c where c.corr_key = p.corr_key) as avg_value, (select ifnull(min(c.v), 8) from child_agg c where c.corr_key = p.corr_key) as min_value, (select ifnull(max(c.v), 9) from child_agg c where c.corr_key = p.corr_key) as max_value from parent_agg p order by p.id; +➤ id[4,32,0] ¦ avg_value[8,54,0] ¦ min_value[-5,64,0] ¦ max_value[-5,64,0] 𝄀 +1 ¦ 6.0 ¦ 5 ¦ 7 𝄀 +2 ¦ 7.0 ¦ 8 ¦ 9 𝄀 +3 ¦ 7.0 ¦ 8 ¦ 9 𝄀 +4 ¦ 7.0 ¦ 8 ¦ 9 +select p.id, (select sum(c.v) from child_agg c where c.corr_key = p.corr_key) as raw_sum from parent_agg p order by p.id; +➤ id[4,32,0] ¦ raw_sum[-5,64,0] 𝄀 +1 ¦ 12 𝄀 +2 ¦ null 𝄀 +3 ¦ null 𝄀 +4 ¦ null +select p.id, (select count(*) from child_agg c where c.corr_key = p.corr_key) as row_count, (select count(c.v) from child_agg c where c.corr_key = p.corr_key) as value_count from parent_agg p order by p.id; +➤ id[4,32,0] ¦ row_count[-5,64,0] ¦ value_count[-5,64,0] 𝄀 +1 ¦ 2 ¦ 2 𝄀 +2 ¦ 1 ¦ 0 𝄀 +3 ¦ 0 ¦ 0 𝄀 +4 ¦ 0 ¦ 0 +select p.id, (select count(*) + 1 from child_agg c where c.corr_key = p.corr_key) as count_plus_one, (select coalesce(count(*), 5) from child_agg c where c.corr_key = p.corr_key) as count_fallback from parent_agg p order by p.id; +➤ id[4,32,0] ¦ count_plus_one[-5,64,0] ¦ count_fallback[-5,64,0] 𝄀 +1 ¦ 3 ¦ 2 𝄀 +2 ¦ 2 ¦ 1 𝄀 +3 ¦ 1 ¦ 0 𝄀 +4 ¦ 1 ¦ 0 +select p.id, (select coalesce(sum(c.v), 100) + count(*) from child_agg c where c.corr_key = p.corr_key) as mixed_value from parent_agg p order by p.id; +➤ id[4,32,0] ¦ mixed_value[-5,64,0] 𝄀 +1 ¦ 14 𝄀 +2 ¦ 101 𝄀 +3 ¦ 100 𝄀 +4 ¦ 100 +select p.id, (select case when count(*) = 0 then 42 else coalesce(sum(c.v), 0) end from child_agg c where c.corr_key = p.corr_key) as case_value from parent_agg p order by p.id; +➤ id[4,32,0] ¦ case_value[-5,64,0] 𝄀 +1 ¦ 12 𝄀 +2 ¦ 0 𝄀 +3 ¦ 42 𝄀 +4 ¦ 42 +select p.id, (select bit_and(c.v) from child_agg c where c.corr_key = p.corr_key) <=> (select bit_and(c.v) from child_agg c where false) as bit_and_matches, (select bit_or(c.v) from child_agg c where c.corr_key = p.corr_key) <=> (select bit_or(c.v) from child_agg c where false) as bit_or_matches, (select bit_xor(c.v) from child_agg c where c.corr_key = p.corr_key) <=> (select bit_xor(c.v) from child_agg c where false) as bit_xor_matches, (select approx_count_distinct(c.v) from child_agg c where c.corr_key = p.corr_key) <=> (select approx_count_distinct(c.v) from child_agg c where false) as approx_count_matches, (select sum(c.v) from child_agg c where c.corr_key = p.corr_key) <=> (select sum(c.v) from child_agg c where false) as sum_matches from parent_agg p where p.id = 3; +➤ id[4,32,0] ¦ bit_and_matches[-7,1,0] ¦ bit_or_matches[-7,1,0] ¦ bit_xor_matches[-7,1,0] ¦ approx_count_matches[-7,1,0] ¦ sum_matches[-7,1,0] 𝄀 +3 ¦ 1 ¦ 1 ¦ 1 ¦ 1 ¦ 1 +select p.id, (select coalesce(json_arrayagg(c.v), convert('[]', json)) from child_agg c where c.corr_key = p.corr_key) as json_value from parent_agg p order by p.id; +➤ id[4,32,0] ¦ json_value[-1,2147483647,0] 𝄀 +1 ¦ [5, 7] 𝄀 +2 ¦ [null] 𝄀 +3 ¦ [] 𝄀 +4 ¦ [] +with correlated_input as (select corr_key, v from child_agg) select p.id, (select coalesce(sum(c.v), 0) from correlated_input c where c.corr_key = p.corr_key) as cte_sum from parent_agg p order by p.id; +➤ id[4,32,0] ¦ cte_sum[-5,64,0] 𝄀 +1 ¦ 12 𝄀 +2 ¦ 0 𝄀 +3 ¦ 0 𝄀 +4 ¦ 0 +select p.id, (with correlated_input as (select c.v from child_agg c where c.corr_key = p.corr_key) select coalesce(sum(v), 0) from correlated_input) as cte_correlated_sum from parent_agg p order by p.id; +➤ id[4,32,0] ¦ cte_correlated_sum[-5,64,0] 𝄀 +1 ¦ 12 𝄀 +2 ¦ 0 𝄀 +3 ¦ 0 𝄀 +4 ¦ 0 +select p.id, (select sum(c.v) from child_agg c where c.corr_key = p.corr_key group by c.corr_key) as grouped_sum from parent_agg p order by p.id; +➤ id[4,32,0] ¦ grouped_sum[-5,64,0] 𝄀 +1 ¦ 12 𝄀 +2 ¦ null 𝄀 +3 ¦ null 𝄀 +4 ¦ null +select p.id, (select sum(c.v) from child_agg c where c.corr_key = p.corr_key having sum(c.v) > 100) as having_sum from parent_agg p order by p.id; +➤ id[4,32,0] ¦ having_sum[-5,64,0] 𝄀 +1 ¦ null 𝄀 +2 ¦ null 𝄀 +3 ¦ null 𝄀 +4 ¦ null drop database test_subq_corr_project; diff --git a/test/distributed/cases/subquery/scalar_correlated_projection.sql b/test/distributed/cases/subquery/scalar_correlated_projection.sql index b1bf351b87f67..3d9fdc35892b5 100644 --- a/test/distributed/cases/subquery/scalar_correlated_projection.sql +++ b/test/distributed/cases/subquery/scalar_correlated_projection.sql @@ -6,6 +6,10 @@ use test_subq_corr_project; create table t1 (a int, b int, c int); create table t2 (d int); insert into t1 values (1, 2, 3), (11, 22, 33); +create table parent_agg (id int primary key, corr_key int); +create table child_agg (corr_key int, v int); +insert into parent_agg values (1, 10), (2, 20), (3, 30), (4, 30); +insert into child_agg values (10, 5), (10, 7), (20, null); -- @case -- @desc:direct outer column projected by a correlated scalar subquery @@ -28,5 +32,22 @@ select t1.*, (select t1.a from t2 where t2.d > t1.a order by t2.d limit 1) as x select (select t1.a from t2 where t2.d > t1.a limit 2) as x from t1 where t1.a = 11; select (select t1.a from t2 where t2.d > t1.a limit 2) as x from t1 where t1.a = 1; +-- @case +-- @desc:issue #25959 - evaluate the final scalar aggregate projection after LEFT JOIN null extension +-- @label:bvt +select p.id, (select coalesce(sum(c.v), 0) from child_agg c where c.corr_key = p.corr_key) as sum_value from parent_agg p order by p.id; +select p.id, (select ifnull(avg(c.v), 7) from child_agg c where c.corr_key = p.corr_key) as avg_value, (select ifnull(min(c.v), 8) from child_agg c where c.corr_key = p.corr_key) as min_value, (select ifnull(max(c.v), 9) from child_agg c where c.corr_key = p.corr_key) as max_value from parent_agg p order by p.id; +select p.id, (select sum(c.v) from child_agg c where c.corr_key = p.corr_key) as raw_sum from parent_agg p order by p.id; +select p.id, (select count(*) from child_agg c where c.corr_key = p.corr_key) as row_count, (select count(c.v) from child_agg c where c.corr_key = p.corr_key) as value_count from parent_agg p order by p.id; +select p.id, (select count(*) + 1 from child_agg c where c.corr_key = p.corr_key) as count_plus_one, (select coalesce(count(*), 5) from child_agg c where c.corr_key = p.corr_key) as count_fallback from parent_agg p order by p.id; +select p.id, (select coalesce(sum(c.v), 100) + count(*) from child_agg c where c.corr_key = p.corr_key) as mixed_value from parent_agg p order by p.id; +select p.id, (select case when count(*) = 0 then 42 else coalesce(sum(c.v), 0) end from child_agg c where c.corr_key = p.corr_key) as case_value from parent_agg p order by p.id; +select p.id, (select bit_and(c.v) from child_agg c where c.corr_key = p.corr_key) <=> (select bit_and(c.v) from child_agg c where false) as bit_and_matches, (select bit_or(c.v) from child_agg c where c.corr_key = p.corr_key) <=> (select bit_or(c.v) from child_agg c where false) as bit_or_matches, (select bit_xor(c.v) from child_agg c where c.corr_key = p.corr_key) <=> (select bit_xor(c.v) from child_agg c where false) as bit_xor_matches, (select approx_count_distinct(c.v) from child_agg c where c.corr_key = p.corr_key) <=> (select approx_count_distinct(c.v) from child_agg c where false) as approx_count_matches, (select sum(c.v) from child_agg c where c.corr_key = p.corr_key) <=> (select sum(c.v) from child_agg c where false) as sum_matches from parent_agg p where p.id = 3; +select p.id, (select coalesce(json_arrayagg(c.v), convert('[]', json)) from child_agg c where c.corr_key = p.corr_key) as json_value from parent_agg p order by p.id; +with correlated_input as (select corr_key, v from child_agg) select p.id, (select coalesce(sum(c.v), 0) from correlated_input c where c.corr_key = p.corr_key) as cte_sum from parent_agg p order by p.id; +select p.id, (with correlated_input as (select c.v from child_agg c where c.corr_key = p.corr_key) select coalesce(sum(v), 0) from correlated_input) as cte_correlated_sum from parent_agg p order by p.id; +select p.id, (select sum(c.v) from child_agg c where c.corr_key = p.corr_key group by c.corr_key) as grouped_sum from parent_agg p order by p.id; +select p.id, (select sum(c.v) from child_agg c where c.corr_key = p.corr_key having sum(c.v) > 100) as having_sum from parent_agg p order by p.id; + -- @teardown drop database test_subq_corr_project;