github.com/matrixorigin/matrixone@v1.2.0/pkg/sql/plan/distinct_agg.go (about)

     1  // Copyright 2023 Matrix Origin
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package plan
    16  
    17  import (
    18  	"github.com/matrixorigin/matrixone/pkg/pb/plan"
    19  	"github.com/matrixorigin/matrixone/pkg/sql/plan/function"
    20  )
    21  
    22  func (builder *QueryBuilder) optimizeDistinctAgg(nodeID int32) {
    23  	node := builder.qry.Nodes[nodeID]
    24  
    25  	for _, childID := range node.Children {
    26  		builder.optimizeDistinctAgg(childID)
    27  	}
    28  
    29  	if node.NodeType == plan.Node_AGG {
    30  		if len(node.AggList) != 1 {
    31  			return
    32  		}
    33  
    34  		aggFunc := node.AggList[0].GetF()
    35  		if uint64(aggFunc.Func.Obj)&function.Distinct == 0 || (aggFunc.Func.ObjName != "count" && aggFunc.Func.ObjName != "sum") {
    36  			return
    37  		}
    38  
    39  		oldGroupLen := len(node.GroupBy)
    40  		oldGroupBy := node.GroupBy
    41  		toCount := aggFunc.Args[0]
    42  
    43  		newGroupTag := builder.genNewTag()
    44  		newAggregateTag := builder.genNewTag()
    45  		aggNodeID := builder.appendNode(&plan.Node{
    46  			NodeType:    plan.Node_AGG,
    47  			Children:    []int32{node.Children[0]},
    48  			GroupBy:     append(oldGroupBy, toCount),
    49  			BindingTags: []int32{newGroupTag, newAggregateTag},
    50  		}, builder.ctxByNode[node.Children[0]])
    51  
    52  		node.Children[0] = aggNodeID
    53  		node.GroupBy = make([]*plan.Expr, oldGroupLen)
    54  		for i := range node.GroupBy {
    55  			node.GroupBy[i] = &plan.Expr{
    56  				Typ: oldGroupBy[i].Typ,
    57  				Expr: &plan.Expr_Col{
    58  					Col: &plan.ColRef{
    59  						RelPos: newGroupTag,
    60  						ColPos: int32(i),
    61  					},
    62  				},
    63  			}
    64  		}
    65  
    66  		aggFunc.Func.Obj &= function.DistinctMask
    67  		aggFunc.Args[0] = &plan.Expr{
    68  			Typ: toCount.Typ,
    69  			Expr: &plan.Expr_Col{
    70  				Col: &plan.ColRef{
    71  					RelPos: newGroupTag,
    72  					ColPos: int32(oldGroupLen),
    73  				},
    74  			},
    75  		}
    76  	}
    77  }