in arrow/compute/arithmetic.go [577:1006]
func RegisterScalarArithmetic(reg FunctionRegistry) {
ops := []struct {
funcName string
op kernels.ArithmeticOp
decPromote decimalPromotion
doc FunctionDoc
}{
{"add_unchecked", kernels.OpAdd, decPromoteAdd, addUncheckedDoc},
{"add", kernels.OpAddChecked, decPromoteAdd, addDoc},
}
for _, o := range ops {
fn := &arithmeticFunction{*NewScalarFunction(o.funcName, Binary(), o.doc), o.decPromote}
kns := append(kernels.GetArithmeticBinaryKernels(o.op), kernels.GetDecimalBinaryKernels(o.op)...)
kns = append(kns, kernels.GetArithmeticFunctionTimeDuration(o.op)...)
for _, k := range kns {
if err := fn.AddKernel(k); err != nil {
panic(err)
}
}
for _, unit := range arrow.TimeUnitValues {
inType := exec.NewMatchedInput(exec.TimestampTypeUnit(unit))
inDuration := exec.NewExactInput(&arrow.DurationType{Unit: unit})
ex := kernels.ArithmeticExecSameType(arrow.TIMESTAMP, o.op)
err := fn.AddNewKernel([]exec.InputType{inType, inDuration}, kernels.OutputFirstType, ex, nil)
if err != nil {
panic(err)
}
err = fn.AddNewKernel([]exec.InputType{inDuration, inType}, kernels.OutputLastType, ex, nil)
if err != nil {
panic(err)
}
matchDur := exec.NewMatchedInput(exec.DurationTypeUnit(unit))
ex = kernels.ArithmeticExecSameType(arrow.DURATION, o.op)
err = fn.AddNewKernel([]exec.InputType{matchDur, matchDur}, exec.NewOutputType(&arrow.DurationType{Unit: unit}), ex, nil)
if err != nil {
panic(err)
}
}
reg.AddFunction(fn, false)
}
ops = []struct {
funcName string
op kernels.ArithmeticOp
decPromote decimalPromotion
doc FunctionDoc
}{
{"sub_unchecked", kernels.OpSub, decPromoteAdd, subUncheckedDoc},
{"sub", kernels.OpSubChecked, decPromoteAdd, subDoc},
{"subtract_unchecked", kernels.OpSub, decPromoteAdd, subUncheckedDoc},
{"subtract", kernels.OpSubChecked, decPromoteAdd, subDoc},
}
for _, o := range ops {
fn := &arithmeticFunction{*NewScalarFunction(o.funcName, Binary(), o.doc), o.decPromote}
kns := append(kernels.GetArithmeticBinaryKernels(o.op), kernels.GetDecimalBinaryKernels(o.op)...)
kns = append(kns, kernels.GetArithmeticFunctionTimeDuration(o.op)...)
for _, k := range kns {
if err := fn.AddKernel(k); err != nil {
panic(err)
}
}
for _, unit := range arrow.TimeUnitValues {
// timestamp - timestamp => duration
inType := exec.NewMatchedInput(exec.TimestampTypeUnit(unit))
ex := kernels.ArithmeticExecSameType(arrow.TIMESTAMP, o.op)
err := fn.AddNewKernel([]exec.InputType{inType, inType}, kernels.OutputResolveTemporal, ex, nil)
if err != nil {
panic(err)
}
// timestamp - duration => timestamp
inDuration := exec.NewExactInput(&arrow.DurationType{Unit: unit})
ex = kernels.ArithmeticExecSameType(arrow.TIMESTAMP, o.op)
err = fn.AddNewKernel([]exec.InputType{inType, inDuration}, kernels.OutputFirstType, ex, nil)
if err != nil {
panic(err)
}
// duration - duration = duration
matchDur := exec.NewMatchedInput(exec.DurationTypeUnit(unit))
ex = kernels.ArithmeticExecSameType(arrow.DURATION, o.op)
err = fn.AddNewKernel([]exec.InputType{matchDur, matchDur}, exec.NewOutputType(&arrow.DurationType{Unit: unit}), ex, nil)
if err != nil {
panic(err)
}
}
// time32 - time32 = duration
for _, unit := range []arrow.TimeUnit{arrow.Second, arrow.Millisecond} {
inType := exec.NewMatchedInput(exec.Time32TypeUnit(unit))
internalEx := kernels.ArithmeticExecSameType(arrow.TIME32, o.op)
ex := func(ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.ExecResult) error {
if err := internalEx(ctx, batch, out); err != nil {
return err
}
// the allocated space is for duration (an int64) but we
// wrote the time32 - time32 as if the output was time32
// so a quick copy in reverse expands the int32s to int64.
rawData := arrow.GetData[int32](out.Buffers[1].Buf)
outData := arrow.GetData[int64](out.Buffers[1].Buf)
for i := out.Len - 1; i >= 0; i-- {
outData[i] = int64(rawData[i])
}
return nil
}
err := fn.AddNewKernel([]exec.InputType{inType, inType},
exec.NewOutputType(&arrow.DurationType{Unit: unit}), ex, nil)
if err != nil {
panic(err)
}
}
// time64 - time64 = duration
for _, unit := range []arrow.TimeUnit{arrow.Microsecond, arrow.Nanosecond} {
inType := exec.NewMatchedInput(exec.Time64TypeUnit(unit))
ex := kernels.ArithmeticExecSameType(arrow.TIME64, o.op)
err := fn.AddNewKernel([]exec.InputType{inType, inType}, exec.NewOutputType(&arrow.DurationType{Unit: unit}), ex, nil)
if err != nil {
panic(err)
}
}
inDate32 := exec.NewExactInput(arrow.FixedWidthTypes.Date32)
ex := kernels.SubtractDate32(o.op)
err := fn.AddNewKernel([]exec.InputType{inDate32, inDate32}, exec.NewOutputType(arrow.FixedWidthTypes.Duration_s), ex, nil)
if err != nil {
panic(err)
}
inDate64 := exec.NewExactInput(arrow.FixedWidthTypes.Date64)
ex = kernels.ArithmeticExecSameType(arrow.DATE64, o.op)
err = fn.AddNewKernel([]exec.InputType{inDate64, inDate64}, exec.NewOutputType(arrow.FixedWidthTypes.Duration_ms), ex, nil)
if err != nil {
panic(err)
}
reg.AddFunction(fn, false)
}
oplist := []struct {
funcName string
op kernels.ArithmeticOp
decPromote decimalPromotion
doc FunctionDoc
commutative bool
}{
{"multiply_unchecked", kernels.OpMul, decPromoteMultiply, mulUncheckedDoc, true},
{"multiply", kernels.OpMulChecked, decPromoteMultiply, mulDoc, true},
{"divide_unchecked", kernels.OpDiv, decPromoteDivide, divUncheckedDoc, false},
{"divide", kernels.OpDivChecked, decPromoteDivide, divDoc, false},
}
for _, o := range oplist {
fn := &arithmeticFunction{*NewScalarFunction(o.funcName, Binary(), o.doc), o.decPromote}
for _, k := range append(kernels.GetArithmeticBinaryKernels(o.op), kernels.GetDecimalBinaryKernels(o.op)...) {
if err := fn.AddKernel(k); err != nil {
panic(err)
}
}
for _, unit := range arrow.TimeUnitValues {
durInput := exec.NewExactInput(&arrow.DurationType{Unit: unit})
i64Input := exec.NewExactInput(arrow.PrimitiveTypes.Int64)
durOutput := exec.NewOutputType(&arrow.DurationType{Unit: unit})
ex := kernels.ArithmeticExecSameType(arrow.DURATION, o.op)
err := fn.AddNewKernel([]exec.InputType{durInput, i64Input}, durOutput, ex, nil)
if err != nil {
panic(err)
}
if o.commutative {
err = fn.AddNewKernel([]exec.InputType{i64Input, durInput}, durOutput, ex, nil)
if err != nil {
panic(err)
}
}
}
reg.AddFunction(fn, false)
}
ops = []struct {
funcName string
op kernels.ArithmeticOp
decPromote decimalPromotion
doc FunctionDoc
}{
{"abs_unchecked", kernels.OpAbsoluteValue, decPromoteNone, absoluteValueUncheckedDoc},
{"abs", kernels.OpAbsoluteValueChecked, decPromoteNone, absoluteValueDoc},
{"negate_unchecked", kernels.OpNegate, decPromoteNone, negateUncheckedDoc},
}
for _, o := range ops {
fn := &arithmeticFunction{*NewScalarFunction(o.funcName, Unary(), o.doc), decPromoteNone}
kns := append(kernels.GetArithmeticUnaryKernels(o.op), kernels.GetDecimalUnaryKernels(o.op)...)
for _, k := range kns {
if err := fn.AddKernel(k); err != nil {
panic(err)
}
}
reg.AddFunction(fn, false)
}
fn := &arithmeticFunction{*NewScalarFunction("negate", Unary(), negateDoc), decPromoteNone}
kns := append(kernels.GetArithmeticUnarySignedKernels(kernels.OpNegateChecked), kernels.GetDecimalUnaryKernels(kernels.OpNegateChecked)...)
for _, k := range kns {
if err := fn.AddKernel(k); err != nil {
panic(err)
}
}
reg.AddFunction(fn, false)
ops = []struct {
funcName string
op kernels.ArithmeticOp
decPromote decimalPromotion
doc FunctionDoc
}{
{"sqrt_unchecked", kernels.OpSqrt, decPromoteNone, sqrtUncheckedDoc},
{"sqrt", kernels.OpSqrtChecked, decPromoteNone, sqrtDoc},
{"sin_unchecked", kernels.OpSin, decPromoteNone, sinUncheckedDoc},
{"sin", kernels.OpSinChecked, decPromoteNone, sinDoc},
{"cos_unchecked", kernels.OpCos, decPromoteNone, cosUncheckedDoc},
{"cos", kernels.OpCosChecked, decPromoteNone, cosDoc},
{"tan_unchecked", kernels.OpTan, decPromoteNone, tanUncheckedDoc},
{"tan", kernels.OpTanChecked, decPromoteNone, tanDoc},
{"asin_unchecked", kernels.OpAsin, decPromoteNone, asinUncheckedDoc},
{"asin", kernels.OpAsinChecked, decPromoteNone, asinDoc},
{"acos_unchecked", kernels.OpAcos, decPromoteNone, acosUncheckedDoc},
{"acos", kernels.OpAcosChecked, decPromoteNone, acosDoc},
{"atan", kernels.OpAtan, decPromoteNone, atanDoc},
{"ln_unchecked", kernels.OpLn, decPromoteNone, lnUncheckedDoc},
{"ln", kernels.OpLnChecked, decPromoteNone, lnDoc},
{"log10_unchecked", kernels.OpLog10, decPromoteNone, log10UncheckedDoc},
{"log10", kernels.OpLog10Checked, decPromoteNone, log10Doc},
{"log2_unchecked", kernels.OpLog2, decPromoteNone, log2UncheckedDoc},
{"log2", kernels.OpLog2Checked, decPromoteNone, log2Doc},
{"log1p_unchecked", kernels.OpLog1p, decPromoteNone, log1pUncheckedDoc},
{"log1p", kernels.OpLog1pChecked, decPromoteNone, log1pDoc},
}
for _, o := range ops {
fn := &arithmeticFloatingPointFunc{arithmeticFunction{*NewScalarFunction(o.funcName, Unary(), o.doc), decPromoteNone}}
kns := kernels.GetArithmeticUnaryFloatingPointKernels(o.op)
for _, k := range kns {
if err := fn.AddKernel(k); err != nil {
panic(err)
}
}
reg.AddFunction(fn, false)
}
ops = []struct {
funcName string
op kernels.ArithmeticOp
decPromote decimalPromotion
doc FunctionDoc
}{
{"atan2", kernels.OpAtan2, decPromoteNone, atan2Doc},
{"logb_unchecked", kernels.OpLogb, decPromoteNone, logbUncheckedDoc},
{"logb", kernels.OpLogbChecked, decPromoteNone, logbDoc},
}
for _, o := range ops {
fn := &arithmeticFloatingPointFunc{arithmeticFunction{*NewScalarFunction(o.funcName, Binary(), addDoc), decPromoteNone}}
kns := kernels.GetArithmeticFloatingPointKernels(o.op)
for _, k := range kns {
if err := fn.AddKernel(k); err != nil {
panic(err)
}
}
reg.AddFunction(fn, false)
}
fn = &arithmeticFunction{*NewScalarFunction("sign", Unary(), signDoc), decPromoteNone}
kns = kernels.GetArithmeticUnaryFixedIntOutKernels(arrow.PrimitiveTypes.Int8, kernels.OpSign)
for _, k := range kns {
if err := fn.AddKernel(k); err != nil {
panic(err)
}
}
reg.AddFunction(fn, false)
ops = []struct {
funcName string
op kernels.ArithmeticOp
decPromote decimalPromotion
doc FunctionDoc
}{
{"power_unchecked", kernels.OpPower, decPromoteNone, powUncheckedDoc},
{"power", kernels.OpPowerChecked, decPromoteNone, powDoc},
}
for _, o := range ops {
fn := &arithmeticDecimalToFloatingPointFunc{arithmeticFunction{*NewScalarFunction(o.funcName, Binary(), o.doc), o.decPromote}}
kns := kernels.GetArithmeticBinaryKernels(o.op)
for _, k := range kns {
if err := fn.AddKernel(k); err != nil {
panic(err)
}
}
reg.AddFunction(fn, false)
}
bitWiseOps := []struct {
funcName string
op kernels.BitwiseOp
doc FunctionDoc
}{
{"bit_wise_and", kernels.OpBitAnd, bitWiseAndDoc},
{"bit_wise_or", kernels.OpBitOr, bitWiseOrDoc},
{"bit_wise_xor", kernels.OpBitXor, bitWiseXorDoc},
}
for _, o := range bitWiseOps {
fn := &arithmeticFunction{*NewScalarFunction(o.funcName, Binary(), o.doc), decPromoteNone}
kns := kernels.GetBitwiseBinaryKernels(o.op)
for _, k := range kns {
if err := fn.AddKernel(k); err != nil {
panic(err)
}
}
reg.AddFunction(fn, false)
}
fn = &arithmeticFunction{*NewScalarFunction("bit_wise_not", Unary(), bitWiseNotDoc), decPromoteNone}
for _, k := range kernels.GetBitwiseUnaryKernels() {
if err := fn.AddKernel(k); err != nil {
panic(err)
}
}
reg.AddFunction(fn, false)
shiftOps := []struct {
funcName string
dir kernels.ShiftDir
checked bool
doc FunctionDoc
}{
{"shift_left", kernels.ShiftLeft, true, shiftLeftDoc},
{"shift_left_unchecked", kernels.ShiftLeft, false, shiftLeftUncheckedDoc},
{"shift_right", kernels.ShiftRight, true, shiftRightDoc},
{"shift_right_unchecked", kernels.ShiftRight, false, shiftRightUncheckedDoc},
}
for _, o := range shiftOps {
fn := &arithmeticFunction{*NewScalarFunction(o.funcName, Binary(), o.doc), decPromoteNone}
kns := kernels.GetShiftKernels(o.dir, o.checked)
for _, k := range kns {
if err := fn.AddKernel(k); err != nil {
panic(err)
}
}
reg.AddFunction(fn, false)
}
floorFn := &arithmeticIntegerToFloatingPointFunc{arithmeticFunction{*NewScalarFunction("floor", Unary(), floorDoc), decPromoteNone}}
kns = kernels.GetSimpleRoundKernels(kernels.RoundDown)
for _, k := range kns {
if err := floorFn.AddKernel(k); err != nil {
panic(err)
}
}
floorFn.AddNewKernel([]exec.InputType{exec.NewIDInput(arrow.DECIMAL128)},
kernels.OutputFirstType, kernels.FixedRoundDecimalExec[decimal128.Num](kernels.RoundDown), nil)
floorFn.AddNewKernel([]exec.InputType{exec.NewIDInput(arrow.DECIMAL256)},
kernels.OutputFirstType, kernels.FixedRoundDecimalExec[decimal256.Num](kernels.RoundDown), nil)
reg.AddFunction(floorFn, false)
ceilFn := &arithmeticIntegerToFloatingPointFunc{arithmeticFunction{*NewScalarFunction("ceil", Unary(), ceilDoc), decPromoteNone}}
kns = kernels.GetSimpleRoundKernels(kernels.RoundUp)
for _, k := range kns {
if err := ceilFn.AddKernel(k); err != nil {
panic(err)
}
}
ceilFn.AddNewKernel([]exec.InputType{exec.NewIDInput(arrow.DECIMAL128)},
kernels.OutputFirstType, kernels.FixedRoundDecimalExec[decimal128.Num](kernels.RoundUp), nil)
ceilFn.AddNewKernel([]exec.InputType{exec.NewIDInput(arrow.DECIMAL256)},
kernels.OutputFirstType, kernels.FixedRoundDecimalExec[decimal256.Num](kernels.RoundUp), nil)
reg.AddFunction(ceilFn, false)
truncFn := &arithmeticIntegerToFloatingPointFunc{arithmeticFunction{*NewScalarFunction("trunc", Unary(), truncDoc), decPromoteNone}}
kns = kernels.GetSimpleRoundKernels(kernels.TowardsZero)
for _, k := range kns {
if err := truncFn.AddKernel(k); err != nil {
panic(err)
}
}
truncFn.AddNewKernel([]exec.InputType{exec.NewIDInput(arrow.DECIMAL128)},
kernels.OutputFirstType, kernels.FixedRoundDecimalExec[decimal128.Num](kernels.TowardsZero), nil)
truncFn.AddNewKernel([]exec.InputType{exec.NewIDInput(arrow.DECIMAL256)},
kernels.OutputFirstType, kernels.FixedRoundDecimalExec[decimal256.Num](kernels.TowardsZero), nil)
reg.AddFunction(truncFn, false)
roundFn := &arithmeticIntegerToFloatingPointFunc{arithmeticFunction{*NewScalarFunction("round", Unary(), roundDoc), decPromoteNone}}
kns = kernels.GetRoundUnaryKernels(kernels.InitRoundState, kernels.UnaryRoundExec)
for _, k := range kns {
if err := roundFn.AddKernel(k); err != nil {
panic(err)
}
}
roundFn.defaultOpts = DefaultRoundOptions
reg.AddFunction(roundFn, false)
roundToMultipleFn := &arithmeticIntegerToFloatingPointFunc{arithmeticFunction{*NewScalarFunction("round_to_multiple", Unary(), roundToMultipleDoc), decPromoteNone}}
kns = kernels.GetRoundUnaryKernels(kernels.InitRoundToMultipleState, kernels.UnaryRoundToMultipleExec)
for _, k := range kns {
if err := roundToMultipleFn.AddKernel(k); err != nil {
panic(err)
}
}
roundToMultipleFn.defaultOpts = DefaultRoundToMultipleOptions
reg.AddFunction(roundToMultipleFn, false)
}