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

     1  // Copyright 2024 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 agg
    16  
    17  import (
    18  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    19  	"github.com/matrixorigin/matrixone/pkg/container/types"
    20  	"github.com/matrixorigin/matrixone/pkg/sql/colexec/aggexec"
    21  )
    22  
    23  func RegisterSum1(id int64) {
    24  	aggexec.RegisterSingleAggFromFixedToFixed(
    25  		aggexec.MakeSingleAgg1RegisteredInfo(
    26  			aggexec.MakeSingleColumnAggInformation(id, types.T_bit.ToType(), SumReturnType, false, true),
    27  			aggexec.GenerateEmptyContextFromFixedToFixed[uint64, uint64],
    28  			nil,
    29  			FillAggSum1[uint64, uint64], nil, FillsAggSum1[uint64, uint64],
    30  			MergeAggSum1[uint64, uint64],
    31  			nil,
    32  		))
    33  
    34  	aggexec.RegisterSingleAggFromFixedToFixed(
    35  		aggexec.MakeSingleAgg1RegisteredInfo(
    36  			aggexec.MakeSingleColumnAggInformation(id, types.T_int8.ToType(), SumReturnType, false, true),
    37  			aggexec.GenerateEmptyContextFromFixedToFixed[int8, int64],
    38  			nil,
    39  			FillAggSum1[int8, int64], nil, FillsAggSum1[int8, int64],
    40  			MergeAggSum1[int8, int64],
    41  			nil,
    42  		))
    43  
    44  	aggexec.RegisterSingleAggFromFixedToFixed(
    45  		aggexec.MakeSingleAgg1RegisteredInfo(
    46  			aggexec.MakeSingleColumnAggInformation(id, types.T_int16.ToType(), SumReturnType, false, true),
    47  			aggexec.GenerateEmptyContextFromFixedToFixed[int16, int64],
    48  			nil,
    49  			FillAggSum1[int16, int64], nil, FillsAggSum1[int16, int64],
    50  			MergeAggSum1[int16, int64],
    51  			nil,
    52  		))
    53  
    54  	aggexec.RegisterSingleAggFromFixedToFixed(
    55  		aggexec.MakeSingleAgg1RegisteredInfo(
    56  			aggexec.MakeSingleColumnAggInformation(id, types.T_int32.ToType(), SumReturnType, false, true),
    57  			aggexec.GenerateEmptyContextFromFixedToFixed[int32, int64],
    58  			nil,
    59  			FillAggSum1[int32, int64], nil, FillsAggSum1[int32, int64],
    60  			MergeAggSum1[int32, int64],
    61  			nil,
    62  		))
    63  
    64  	aggexec.RegisterSingleAggFromFixedToFixed(
    65  		aggexec.MakeSingleAgg1RegisteredInfo(
    66  			aggexec.MakeSingleColumnAggInformation(id, types.T_int64.ToType(), SumReturnType, false, true),
    67  			aggexec.GenerateEmptyContextFromFixedToFixed[int64, int64],
    68  			nil,
    69  			FillAggSum1[int64, int64], nil, FillsAggSum1[int64, int64],
    70  			MergeAggSum1[int64, int64],
    71  			nil,
    72  		))
    73  
    74  	aggexec.RegisterSingleAggFromFixedToFixed(
    75  		aggexec.MakeSingleAgg1RegisteredInfo(
    76  			aggexec.MakeSingleColumnAggInformation(id, types.T_uint8.ToType(), SumReturnType, false, true),
    77  			aggexec.GenerateEmptyContextFromFixedToFixed[uint8, uint64],
    78  			nil,
    79  			FillAggSum1[uint8, uint64], nil, FillsAggSum1[uint8, uint64],
    80  			MergeAggSum1[uint8, uint64],
    81  			nil,
    82  		))
    83  
    84  	aggexec.RegisterSingleAggFromFixedToFixed(
    85  		aggexec.MakeSingleAgg1RegisteredInfo(
    86  			aggexec.MakeSingleColumnAggInformation(id, types.T_uint16.ToType(), SumReturnType, false, true),
    87  			aggexec.GenerateEmptyContextFromFixedToFixed[uint16, uint64],
    88  			nil,
    89  			FillAggSum1[uint16, uint64], nil, FillsAggSum1[uint16, uint64],
    90  			MergeAggSum1[uint16, uint64],
    91  			nil,
    92  		))
    93  
    94  	aggexec.RegisterSingleAggFromFixedToFixed(
    95  		aggexec.MakeSingleAgg1RegisteredInfo(
    96  			aggexec.MakeSingleColumnAggInformation(id, types.T_uint32.ToType(), SumReturnType, false, true),
    97  			aggexec.GenerateEmptyContextFromFixedToFixed[uint32, uint64],
    98  			nil,
    99  			FillAggSum1[uint32, uint64], nil, FillsAggSum1[uint32, uint64],
   100  			MergeAggSum1[uint32, uint64],
   101  			nil,
   102  		))
   103  
   104  	aggexec.RegisterSingleAggFromFixedToFixed(
   105  		aggexec.MakeSingleAgg1RegisteredInfo(
   106  			aggexec.MakeSingleColumnAggInformation(id, types.T_uint64.ToType(), SumReturnType, false, true),
   107  			aggexec.GenerateEmptyContextFromFixedToFixed[uint64, uint64],
   108  			nil,
   109  			FillAggSum1[uint64, uint64], nil, FillsAggSum1[uint64, uint64],
   110  			MergeAggSum1[uint64, uint64],
   111  			nil,
   112  		))
   113  
   114  	aggexec.RegisterSingleAggFromFixedToFixed(
   115  		aggexec.MakeSingleAgg1RegisteredInfo(
   116  			aggexec.MakeSingleColumnAggInformation(id, types.T_float32.ToType(), SumReturnType, false, true),
   117  			aggexec.GenerateEmptyContextFromFixedToFixed[float32, float64],
   118  			nil,
   119  			FillAggSum1[float32, float64], nil, FillsAggSum1[float32, float64],
   120  			MergeAggSum1[float32, float64],
   121  			nil,
   122  		))
   123  
   124  	aggexec.RegisterSingleAggFromFixedToFixed(
   125  		aggexec.MakeSingleAgg1RegisteredInfo(
   126  			aggexec.MakeSingleColumnAggInformation(id, types.T_float64.ToType(), SumReturnType, false, true),
   127  			aggexec.GenerateEmptyContextFromFixedToFixed[float64, float64],
   128  			nil,
   129  			FillAggSum1[float64, float64], nil, FillsAggSum1[float64, float64],
   130  			MergeAggSum1[float64, float64],
   131  			nil,
   132  		))
   133  
   134  	aggexec.RegisterSingleAggFromFixedToFixed(
   135  		aggexec.MakeSingleAgg1RegisteredInfo(
   136  			aggexec.MakeSingleColumnAggInformation(id, types.T_decimal64.ToType(), SumReturnType, false, true),
   137  			newAggSumDecimal64,
   138  			InitAggSumDecimal64,
   139  			FillAggSumDecimal64, nil, FillsAggSumDecimal64,
   140  			MergeAggSumDecimal64,
   141  			nil,
   142  		))
   143  
   144  	aggexec.RegisterSingleAggFromFixedToFixed(
   145  		aggexec.MakeSingleAgg1RegisteredInfo(
   146  			aggexec.MakeSingleColumnAggInformation(id, types.T_decimal128.ToType(), SumReturnType, false, true),
   147  			newAggSumDecimal128,
   148  			InitAggSumDecimal128,
   149  			FillAggSumDecimal128, nil, FillsAggSumDecimal128,
   150  			MergeAggSumDecimal128,
   151  			nil,
   152  		))
   153  }
   154  
   155  var (
   156  	SumSupportedTypes = []types.T{
   157  		types.T_bit,
   158  		types.T_uint8, types.T_uint16, types.T_uint32, types.T_uint64,
   159  		types.T_int8, types.T_int16, types.T_int32, types.T_int64,
   160  		types.T_float32, types.T_float64,
   161  		types.T_decimal64, types.T_decimal128,
   162  	}
   163  	SumReturnType = func(typs []types.Type) types.Type {
   164  		switch typs[0].Oid {
   165  		case types.T_float32, types.T_float64:
   166  			return types.T_float64.ToType()
   167  		case types.T_int8, types.T_int16, types.T_int32, types.T_int64:
   168  			return types.T_int64.ToType()
   169  		case types.T_uint8, types.T_uint16, types.T_uint32, types.T_uint64:
   170  			return types.T_uint64.ToType()
   171  		case types.T_bit:
   172  			return types.T_uint64.ToType()
   173  		case types.T_decimal64:
   174  			return types.New(types.T_decimal128, 38, typs[0].Scale)
   175  		case types.T_decimal128:
   176  			return types.New(types.T_decimal128, 38, typs[0].Scale)
   177  		}
   178  		panic(moerr.NewInternalErrorNoCtx("unsupported type '%v' for sum", typs[0]))
   179  	}
   180  )
   181  
   182  func FillAggSum1[from numeric, to numericWithMaxScale](
   183  	exec aggexec.SingleAggFromFixedRetFixed[from, to], value from, getter aggexec.AggGetter[to], setter aggexec.AggSetter[to]) error {
   184  	setter(getter() + to(value))
   185  	return nil
   186  }
   187  func FillsAggSum1[from numeric, to numericWithMaxScale](
   188  	exec aggexec.SingleAggFromFixedRetFixed[from, to],
   189  	value from, isNull bool, count int, getter aggexec.AggGetter[to], setter aggexec.AggSetter[to]) error {
   190  	if !isNull {
   191  		setter(getter() + to(value)*to(count))
   192  	}
   193  	return nil
   194  }
   195  func MergeAggSum1[from numeric, to numericWithMaxScale](
   196  	exec1, exec2 aggexec.SingleAggFromFixedRetFixed[from, to], getter1, getter2 aggexec.AggGetter[to], setter aggexec.AggSetter[to]) error {
   197  	setter(getter1() + getter2())
   198  	return nil
   199  }
   200  
   201  type aggSumDecimal64 struct {
   202  	argScale int32
   203  }
   204  
   205  func newAggSumDecimal64() aggexec.SingleAggFromFixedRetFixed[types.Decimal64, types.Decimal128] {
   206  	return &aggSumDecimal64{}
   207  }
   208  
   209  func (a *aggSumDecimal64) Marshal() []byte     { return types.EncodeInt32(&a.argScale) }
   210  func (a *aggSumDecimal64) Unmarshal(bs []byte) { a.argScale = types.DecodeInt32(bs) }
   211  
   212  func InitAggSumDecimal64(
   213  	exec aggexec.SingleAggFromFixedRetFixed[types.Decimal64, types.Decimal128], setter aggexec.AggSetter[types.Decimal128], arg, ret types.Type) error {
   214  	setter(types.Decimal128{B0_63: 0, B64_127: 0})
   215  	a := exec.(*aggSumDecimal64)
   216  	a.argScale = arg.Scale
   217  	return nil
   218  }
   219  func FillAggSumDecimal64(
   220  	exec aggexec.SingleAggFromFixedRetFixed[types.Decimal64, types.Decimal128], value types.Decimal64, getter aggexec.AggGetter[types.Decimal128], setter aggexec.AggSetter[types.Decimal128]) error {
   221  	r, err := getter().Add64(value)
   222  	setter(r)
   223  	return err
   224  }
   225  func FillsAggSumDecimal64(
   226  	exec aggexec.SingleAggFromFixedRetFixed[types.Decimal64, types.Decimal128],
   227  	value types.Decimal64, isNull bool, count int, getter aggexec.AggGetter[types.Decimal128], setter aggexec.AggSetter[types.Decimal128]) error {
   228  	if !isNull {
   229  		a := exec.(*aggSumDecimal64)
   230  
   231  		v := types.Decimal128{B0_63: uint64(value), B64_127: 0}
   232  		if value.Sign() {
   233  			v.B64_127 = ^v.B64_127
   234  		}
   235  		r, _, err := v.Mul(types.Decimal128{B0_63: uint64(count), B64_127: 0}, a.argScale, 0)
   236  		if err != nil {
   237  			return err
   238  		}
   239  		r, err = getter().Add128(r)
   240  		setter(r)
   241  		return err
   242  	}
   243  	return nil
   244  }
   245  func MergeAggSumDecimal64(
   246  	exec1, exec2 aggexec.SingleAggFromFixedRetFixed[types.Decimal64, types.Decimal128], getter1, getter2 aggexec.AggGetter[types.Decimal128], setter aggexec.AggSetter[types.Decimal128]) error {
   247  	r, err := getter1().Add128(getter2())
   248  	setter(r)
   249  	return err
   250  }
   251  
   252  type aggSumDecimal128 struct {
   253  	argScale int32
   254  }
   255  
   256  func newAggSumDecimal128() aggexec.SingleAggFromFixedRetFixed[types.Decimal128, types.Decimal128] {
   257  	return &aggSumDecimal128{}
   258  }
   259  
   260  func (a *aggSumDecimal128) Marshal() []byte     { return types.EncodeInt32(&a.argScale) }
   261  func (a *aggSumDecimal128) Unmarshal(bs []byte) { a.argScale = types.DecodeInt32(bs) }
   262  
   263  func InitAggSumDecimal128(
   264  	exec aggexec.SingleAggFromFixedRetFixed[types.Decimal128, types.Decimal128], setter aggexec.AggSetter[types.Decimal128], arg, ret types.Type) error {
   265  	setter(types.Decimal128{B0_63: 0, B64_127: 0})
   266  	a := exec.(*aggSumDecimal128)
   267  	a.argScale = arg.Scale
   268  	return nil
   269  }
   270  func FillAggSumDecimal128(
   271  	exec aggexec.SingleAggFromFixedRetFixed[types.Decimal128, types.Decimal128], value types.Decimal128, getter aggexec.AggGetter[types.Decimal128], setter aggexec.AggSetter[types.Decimal128]) error {
   272  	r, err := getter().Add128(value)
   273  	setter(r)
   274  	return err
   275  }
   276  func FillsAggSumDecimal128(
   277  	exec aggexec.SingleAggFromFixedRetFixed[types.Decimal128, types.Decimal128],
   278  	value types.Decimal128, isNull bool, count int, getter aggexec.AggGetter[types.Decimal128], setter aggexec.AggSetter[types.Decimal128]) error {
   279  	if !isNull {
   280  		a := exec.(*aggSumDecimal128)
   281  		r, _, err := value.Mul(types.Decimal128{B0_63: uint64(count), B64_127: 0}, a.argScale, 0)
   282  		if err != nil {
   283  			return err
   284  		}
   285  		r, err = getter().Add128(r)
   286  		setter(r)
   287  		return err
   288  	}
   289  	return nil
   290  }
   291  
   292  func MergeAggSumDecimal128(
   293  	exec1, exec2 aggexec.SingleAggFromFixedRetFixed[types.Decimal128, types.Decimal128], getter1, getter2 aggexec.AggGetter[types.Decimal128], setter aggexec.AggSetter[types.Decimal128]) error {
   294  	r, err := getter1().Add128(getter2())
   295  	setter(r)
   296  	return err
   297  }