package compute
import (
"context"
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/compute/exec"
"github.com/apache/arrow-go/v18/arrow/compute/internal/kernels"
)
type compareFunction struct {
ScalarFunction
}
func (fn *compareFunction ) Execute (ctx context .Context , opt FunctionOptions , args ...Datum ) (Datum , error ) {
return execInternal (ctx , fn , opt , -1 , args ...)
}
func (fn *compareFunction ) DispatchBest (vals ...arrow .DataType ) (exec .Kernel , error ) {
if err := fn .checkArity (len (vals )); err != nil {
return nil , err
}
if hasDecimal (vals ...) {
if err := castBinaryDecimalArgs (decPromoteAdd , vals ...); err != nil {
return nil , err
}
}
if kn , err := fn .DispatchExact (vals ...); err == nil {
return kn , nil
}
ensureDictionaryDecoded (vals ...)
ensureNoExtensionType (vals ...)
replaceNullWithOtherType (vals ...)
if dt := commonNumeric (vals ...); dt != nil {
replaceTypes (dt , vals ...)
} else if dt := commonTemporal (vals ...); dt != nil {
replaceTypes (dt , vals ...)
} else if dt := commonBinary (vals ...); dt != nil {
replaceTypes (dt , vals ...)
}
return fn .DispatchExact (vals ...)
}
type flippedData struct {
*kernels .CompareData
unflippedExec exec .ArrayKernelExec
}
func flippedCompare(ctx *exec .KernelCtx , batch *exec .ExecSpan , out *exec .ExecResult ) error {
kn := ctx .Kernel .(*exec .ScalarKernel )
knData := kn .Data .(*flippedData )
flippedBatch := exec .ExecSpan {
Len : batch .Len ,
Values : []exec .ExecValue {batch .Values [1 ], batch .Values [0 ]},
}
return knData .unflippedExec (ctx , &flippedBatch , out )
}
func makeFlippedCompare(name string , fn *compareFunction , doc FunctionDoc ) *compareFunction {
flipped := &compareFunction {*NewScalarFunction (name , Binary (), doc )}
for _ , k := range fn .kernels {
flippedKernel := k
if k .Data != nil {
cmpData := k .Data .(*kernels .CompareData )
flippedKernel .Data = &flippedData {CompareData : cmpData ,
unflippedExec : k .ExecFn }
} else {
flippedKernel .Data = &flippedData {unflippedExec : k .ExecFn }
}
flippedKernel .ExecFn = flippedCompare
flipped .AddKernel (flippedKernel )
}
return flipped
}
func RegisterScalarComparisons (reg FunctionRegistry ) {
eqFn := &compareFunction {*NewScalarFunction ("equal" , Binary (), EmptyFuncDoc )}
for _ , k := range kernels .CompareKernels (kernels .CmpEQ ) {
if err := eqFn .AddKernel (k ); err != nil {
panic (err )
}
}
reg .AddFunction (eqFn , false )
neqFn := &compareFunction {*NewScalarFunction ("not_equal" , Binary (), EmptyFuncDoc )}
for _ , k := range kernels .CompareKernels (kernels .CmpNE ) {
if err := neqFn .AddKernel (k ); err != nil {
panic (err )
}
}
reg .AddFunction (neqFn , false )
gtFn := &compareFunction {*NewScalarFunction ("greater" , Binary (), EmptyFuncDoc )}
for _ , k := range kernels .CompareKernels (kernels .CmpGT ) {
if err := gtFn .AddKernel (k ); err != nil {
panic (err )
}
}
reg .AddFunction (gtFn , false )
gteFn := &compareFunction {*NewScalarFunction ("greater_equal" , Binary (), EmptyFuncDoc )}
for _ , k := range kernels .CompareKernels (kernels .CmpGE ) {
if err := gteFn .AddKernel (k ); err != nil {
panic (err )
}
}
reg .AddFunction (gteFn , false )
ltFn := makeFlippedCompare ("less" , gtFn , EmptyFuncDoc )
reg .AddFunction (ltFn , false )
lteFn := makeFlippedCompare ("less_equal" , gteFn , EmptyFuncDoc )
reg .AddFunction (lteFn , false )
isOrNotNullKns := kernels .IsNullNotNullKernels ()
isNullFn := &compareFunction {*NewScalarFunction ("is_null" , Unary (), EmptyFuncDoc )}
if err := isNullFn .AddKernel (isOrNotNullKns [0 ]); err != nil {
panic (err )
}
isNotNullFn := &compareFunction {*NewScalarFunction ("is_not_null" , Unary (), EmptyFuncDoc )}
if err := isNotNullFn .AddKernel (isOrNotNullKns [1 ]); err != nil {
panic (err )
}
reg .AddFunction (isNullFn , false )
reg .AddFunction (isNotNullFn , false )
isNaNFn := &compareFunction {*NewScalarFunction ("is_nan" , Unary (), EmptyFuncDoc )}
for _ , k := range kernels .IsNaNKernels () {
if err := isNaNFn .AddKernel (k ); err != nil {
panic (err )
}
}
reg .AddFunction (isNaNFn , false )
}
The pages are generated with Golds v0.8.2 . (GOOS=linux GOARCH=amd64)
Golds is a Go 101 project developed by Tapir Liu .
PR and bug reports are welcome and can be submitted to the issue list .
Please follow @zigo_101 (reachable from the left QR code) to get the latest news of Golds .