aggfuncs: implement bit-or with new aggregation framework (#6975)

executor/aggfuncs/aggfuncs.go

@@ -32,11 +32,13 @@ var (
 	_ AggFunc = (*avgOriginal4Float64)(nil)
 	_ AggFunc = (*avgPartial4Float64)(nil)
 
-// All the AggFunc implementations for "FIRSTROW" are listed here.
-// All the AggFunc implementations for "MAX" are listed here.
-// All the AggFunc implementations for "MIN" are listed here.
-// All the AggFunc implementations for "GROUP_CONCAT" are listed here.
-// All the AggFunc implementations for "BIT_OR" are listed here.
+	// All the AggFunc implementations for "FIRSTROW" are listed here.
+	// All the AggFunc implementations for "MAX" are listed here.
+	// All the AggFunc implementations for "MIN" are listed here.
+	// All the AggFunc implementations for "GROUP_CONCAT" are listed here.
+	// All the AggFunc implementations for "BIT_OR" are listed here.
+	_ AggFunc = (*bitOrUint64)(nil)
+
 // All the AggFunc implementations for "BIT_XOR" are listed here.
 // All the AggFunc implementations for "BIT_AND" are listed here.
 )

executor/aggfuncs/builder.go

@@ -17,6 +17,7 @@ import (
 	"github.com/pingcap/tidb/ast"
 	"github.com/pingcap/tidb/expression/aggregation"
 	"github.com/pingcap/tidb/mysql"
+	"github.com/pingcap/tidb/types"
 )
 
 // Build is used to build a specific AggFunc implementation according to the
@@ -120,6 +121,15 @@ func buildGroupConcat(aggFuncDesc *aggregation.AggFuncDesc, ordinal int) AggFunc
 
 // buildCount builds the AggFunc implementation for function "BIT_OR".
 func buildBitOr(aggFuncDesc *aggregation.AggFuncDesc, ordinal int) AggFunc {
+	// BIT_OR doesn't need to handle the distinct property.
+	switch aggFuncDesc.Args[0].GetType().EvalType() {
+	case types.ETInt:
+		base := baseAggFunc{
+			args:    aggFuncDesc.Args,
+			ordinal: ordinal,
+		}
+		return &bitOrUint64{baseBitAggFunc{base}}
+	}
 	return nil
 }
 

executor/aggfuncs/func_bitfuncs.go

@@ -0,0 +1,60 @@
+// Copyright 2018 PingCAP, Inc.
+//
+// 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,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package aggfuncs
+
+import (
+	"github.com/juju/errors"
+	"github.com/pingcap/tidb/sessionctx"
+	"github.com/pingcap/tidb/util/chunk"
+)
+
+type baseBitAggFunc struct {
+	baseAggFunc
+}
+
+type bitOrUint64 struct {
+	baseBitAggFunc
+}
+
+type partialResult4BitFunc = uint64
+
+func (e *baseBitAggFunc) AllocPartialResult() PartialResult {
+	return PartialResult(new(partialResult4BitFunc))
+}
+
+func (e *baseBitAggFunc) ResetPartialResult(pr PartialResult) {
+	p := (*partialResult4BitFunc)(pr)
+	*p = 0
+}
+
+func (e *baseBitAggFunc) AppendFinalResult2Chunk(sctx sessionctx.Context, pr PartialResult, chk *chunk.Chunk) error {
+	p := (*partialResult4BitFunc)(pr)
+	chk.AppendUint64(e.ordinal, *p)
+	return nil
+}
+
+func (e *bitOrUint64) UpdatePartialResult(sctx sessionctx.Context, rowsInGroup []chunk.Row, pr PartialResult) error {
+	p := (*partialResult4BitFunc)(pr)
+	for _, row := range rowsInGroup {
+		inputValue, isNull, err := e.args[0].EvalInt(sctx, row)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		if isNull {
+			continue
+		}
+		*p |= uint64(inputValue)
+	}
+	return nil
+}