From 6a95b4878ee8a02e5e4cdd7e9933ff48d40ec7ca Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Fri, 25 Oct 2024 17:23:04 -0400 Subject: [PATCH 1/2] feat(arrow/compute): Implement kernel for "not" function --- arrow/compute/exprs/types.go | 4 +-- .../internal/kernels/scalar_boolean.go | 15 +++++++++- arrow/compute/scalar_bool.go | 2 ++ arrow/compute/scalar_bool_test.go | 28 +++++++++++++++++++ 4 files changed, 46 insertions(+), 3 deletions(-) diff --git a/arrow/compute/exprs/types.go b/arrow/compute/exprs/types.go index 0a42e6cd..d1adbdaa 100644 --- a/arrow/compute/exprs/types.go +++ b/arrow/compute/exprs/types.go @@ -111,7 +111,7 @@ func init() { } } - for _, fn := range []string{"and", "or"} { + for _, fn := range []string{"and", "or", "not"} { err := DefaultExtensionIDRegistry.AddSubstraitScalarToArrow( extensions.ID{URI: SubstraitBooleanFuncsURI, Name: fn}, simpleMapSubstraitToArrowFunc) @@ -120,7 +120,7 @@ func init() { } } - for _, fn := range []string{"and_kleene", "or_kleene"} { + for _, fn := range []string{"and_kleene", "or_kleene", "not"} { err := DefaultExtensionIDRegistry.AddArrowToSubstrait(fn, simpleMapArrowToSubstraitFunc(SubstraitBooleanFuncsURI)) if err != nil { diff --git a/arrow/compute/internal/kernels/scalar_boolean.go b/arrow/compute/internal/kernels/scalar_boolean.go index cc5d0aec..36598a09 100644 --- a/arrow/compute/internal/kernels/scalar_boolean.go +++ b/arrow/compute/internal/kernels/scalar_boolean.go @@ -26,7 +26,7 @@ import ( type computeWordFN func(leftTrue, leftFalse, rightTrue, rightFalse uint64) (outValid, outData uint64) -func computeKleene(computeWord computeWordFN, ctx *exec.KernelCtx, left, right *exec.ArraySpan, out *exec.ExecResult) error { +func computeKleene(computeWord computeWordFN, _ *exec.KernelCtx, left, right *exec.ArraySpan, out *exec.ExecResult) error { var ( inBMs = [4]bitutil.Bitmap{ {Data: left.Buffers[0].Buf, Offset: left.Offset, Len: left.Len}, @@ -332,3 +332,16 @@ func (KleeneAndNotOpKernel) CallScalarLeft(ctx *exec.KernelCtx, left scalar.Scal func (KleeneAndNotOpKernel) CallScalarRight(ctx *exec.KernelCtx, left *exec.ArraySpan, right scalar.Scalar, out *exec.ExecResult) error { return (KleeneAndOpKernel{}).CallScalarRight(ctx, left, invertScalar(right), out) } + +func NotExecKernel(ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.ExecResult) error { + bitutil.InvertBitmap(batch.Values[0].Array.Buffers[1].Buf, int(batch.Values[0].Array.Offset), + int(batch.Values[0].Array.Len), out.Buffers[1].Buf, int(out.Offset)) + + out.Buffers[0] = batch.Values[0].Array.Buffers[0] + if out.Buffers[0].SelfAlloc { + out.Buffers[0].SelfAlloc = false + } + out.Nulls = batch.Values[0].Array.Nulls + + return nil +} diff --git a/arrow/compute/scalar_bool.go b/arrow/compute/scalar_bool.go index a46d221b..d01675d9 100644 --- a/arrow/compute/scalar_bool.go +++ b/arrow/compute/scalar_bool.go @@ -130,4 +130,6 @@ func RegisterScalarBoolean(reg FunctionRegistry) { andNotKleeneDoc, exec.NullComputedPrealloc) makeFunction(reg, "or_kleene", 2, kernels.SimpleBinary[kernels.KleeneOrOpKernel], orKleeneDoc, exec.NullComputedPrealloc) + makeFunction(reg, "not", 1, kernels.NotExecKernel, EmptyFuncDoc, + exec.NullComputedNoPrealloc) } diff --git a/arrow/compute/scalar_bool_test.go b/arrow/compute/scalar_bool_test.go index 63b5ab2d..78db5de6 100644 --- a/arrow/compute/scalar_bool_test.go +++ b/arrow/compute/scalar_bool_test.go @@ -152,3 +152,31 @@ func TestBooleanKleeneKernels(t *testing.T) { }) } } + +func TestBooleanNot(t *testing.T) { + tests := []struct { + inputJSON, expectedJSON string + }{ + {"[true, true, false, false]", "[false, false, true, true]"}, + {"[null, true, null, false]", "[null, false, null, true]"}, + {"[null, null, null, null]", "[null, null, null, null]"}, + } + + for _, tt := range tests { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + input, _, err := array.FromJSON(mem, arrow.FixedWidthTypes.Boolean, + strings.NewReader(tt.inputJSON)) + require.NoError(t, err) + defer input.Release() + + expected, _, err := array.FromJSON(mem, arrow.FixedWidthTypes.Boolean, + strings.NewReader(tt.expectedJSON)) + require.NoError(t, err) + defer expected.Release() + + checkScalarUnary(t, "not", compute.NewDatumWithoutOwning(input), + compute.NewDatumWithoutOwning(expected), nil) + } +} From 23b194b3fe3eb9ca9e3031396bd72b0bf36e0340 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Mon, 28 Oct 2024 11:45:45 -0400 Subject: [PATCH 2/2] add function doc for 'not' --- arrow/compute/scalar_bool.go | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/arrow/compute/scalar_bool.go b/arrow/compute/scalar_bool.go index d01675d9..9a64382a 100644 --- a/arrow/compute/scalar_bool.go +++ b/arrow/compute/scalar_bool.go @@ -93,6 +93,11 @@ var ( For a different null behavior, see function "and".`, ArgNames: []string{"x", "y"}, } + notDoc = FunctionDoc{ + Summary: "Logical 'not' boolean values", + Description: "Negates the input boolean value", + ArgNames: []string{"x"}, + } ) func makeFunction(reg FunctionRegistry, name string, arity int, ex exec.ArrayKernelExec, doc FunctionDoc, nulls exec.NullHandling) { @@ -130,6 +135,6 @@ func RegisterScalarBoolean(reg FunctionRegistry) { andNotKleeneDoc, exec.NullComputedPrealloc) makeFunction(reg, "or_kleene", 2, kernels.SimpleBinary[kernels.KleeneOrOpKernel], orKleeneDoc, exec.NullComputedPrealloc) - makeFunction(reg, "not", 1, kernels.NotExecKernel, EmptyFuncDoc, + makeFunction(reg, "not", 1, kernels.NotExecKernel, notDoc, exec.NullComputedNoPrealloc) }