This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new 497a462d perf(compute): vectorize numeric to boolean casts (#1282)
497a462d is described below
commit 497a462dea0fe412c3f4fee6a47be5ab62dc0df2
Author: Minh Vu <[email protected]>
AuthorDate: Thu Sep 3 18:06:46 2026 +0200
perf(compute): vectorize numeric to boolean casts (#1282)
## What does this PR do?
- Adds an ARM64 NEON fast path for numeric to boolean casts.
- Reuses the existing NEON `array != scalar` comparison kernel to
produce the output bitmap.
- Covers int32, uint32, int64, uint64, float32, and float64.
- Keeps the scalar path for narrow integers, noasm builds, unsupported
CPUs, and appengine.
- Adds coverage for NaN, signed zero, infinities, sliced input, and
bitmap tails.
- Adds a benchmark matrix for the supported types, sizes, and zero
distributions.
## Benchmark
Apple M1 Pro, `GOMAXPROCS=1`, compared with `-tags noasm`.
- 32-bit types: about 9x faster at 65K values and about 14x faster at 1M
values.
- 64-bit types: about 6x faster at 65K values and about 8x faster at 1M
values.
- Small arrays keep the same allocation count and retain the scalar tail
path.
## Tests
- `go test ./arrow/compute ./arrow/compute/internal/kernels`
- `go test -tags noasm ./arrow/compute ./arrow/compute/internal/kernels`
- amd64 compile check for the compute package noasm fallback
---
arrow/compute/boolean_cast_bench_test.go | 84 ++++++++++++++++++++++
arrow/compute/cast_test.go | 40 +++++++++++
arrow/compute/internal/kernels/boolean_cast.go | 18 +++--
.../internal/kernels/boolean_cast_neon_arm64.go | 47 ++++++++++++
.../compute/internal/kernels/boolean_cast_noasm.go | 28 ++++++++
5 files changed, 211 insertions(+), 6 deletions(-)
diff --git a/arrow/compute/boolean_cast_bench_test.go
b/arrow/compute/boolean_cast_bench_test.go
new file mode 100644
index 00000000..7c4e2873
--- /dev/null
+++ b/arrow/compute/boolean_cast_bench_test.go
@@ -0,0 +1,84 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build go1.18
+
+package compute_test
+
+import (
+ "context"
+ "fmt"
+ "strconv"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/compute"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+)
+
+func newBooleanCastBenchmarkArray(mem memory.Allocator, typ arrow.DataType,
size int, zeroFraction float64) arrow.Array {
+ builder := array.NewBuilder(mem, typ)
+ builder.Reserve(size)
+ for i := 0; i < size; i++ {
+ if float64(i)/float64(size) < zeroFraction {
+ if err := builder.AppendValueFromString("0"); err !=
nil {
+ panic(err)
+ }
+ } else if err := builder.AppendValueFromString("1"); err != nil
{
+ panic(err)
+ }
+ }
+ result := builder.NewArray()
+ builder.Release()
+ return result
+}
+
+func BenchmarkNumericToBoolCast(b *testing.B) {
+ for _, typ := range []arrow.DataType{
+ arrow.PrimitiveTypes.Int32,
+ arrow.PrimitiveTypes.Uint32,
+ arrow.PrimitiveTypes.Int64,
+ arrow.PrimitiveTypes.Uint64,
+ arrow.PrimitiveTypes.Float32,
+ arrow.PrimitiveTypes.Float64,
+ } {
+ width := int64(typ.(arrow.FixedWidthDataType).Bytes())
+ for _, size := range []int{64, 1024, 65536, 1_000_000} {
+ for _, zeroFraction := range []float64{0, 0.5, 1} {
+ b.Run(fmt.Sprintf("type=%s/size=%d/zeros=%s",
typ, size, strconv.FormatFloat(zeroFraction, 'f', -1, 64)), func(b *testing.B) {
+ mem := memory.NewGoAllocator()
+ input :=
newBooleanCastBenchmarkArray(mem, typ, size, zeroFraction)
+ defer input.Release()
+ opts := compute.DefaultCastOptions(true)
+ opts.ToType =
arrow.FixedWidthTypes.Boolean
+ ctx := context.Background()
+
+ b.ReportAllocs()
+ b.SetBytes(int64(size) * width)
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ output, err :=
compute.CastArray(ctx, input, opts)
+ if err != nil {
+ b.Fatal(err)
+ }
+ output.Release()
+ }
+ })
+ }
+ }
+ }
+}
diff --git a/arrow/compute/cast_test.go b/arrow/compute/cast_test.go
index ae0229a3..94805db3 100644
--- a/arrow/compute/cast_test.go
+++ b/arrow/compute/cast_test.go
@@ -466,6 +466,46 @@ func (c *CastSuite) TestNumericToBool() {
}
}
+func (c *CastSuite) TestNumericToBoolSpecialValues() {
+ builder := array.NewFloat64Builder(c.mem)
+ builder.AppendValues([]float64{
+ 1,
+ math.Copysign(0, -1),
+ 0,
+ math.NaN(),
+ math.Inf(1),
+ math.Inf(-1),
+ 1,
+ -1,
+ 0,
+ 1,
+ math.Copysign(0, -1),
+ math.NaN(),
+ math.Inf(1),
+ math.Inf(-1),
+ 0,
+ -1,
+ 1,
+ }, nil)
+ input := builder.NewArray()
+ builder.Release()
+ defer input.Release()
+
+ sliced := array.NewSlice(input, 1, 16)
+ defer sliced.Release()
+
+ expectedBuilder := array.NewBooleanBuilder(c.mem)
+ expectedBuilder.AppendValues([]bool{
+ false, false, true, true, true, true, true, false,
+ true, false, true, true, true, false, true,
+ }, nil)
+ expected := expectedBuilder.NewArray()
+ expectedBuilder.Release()
+ defer expected.Release()
+
+ checkCast(c.T(), sliced, expected, *compute.DefaultCastOptions(true))
+}
+
func (c *CastSuite) StringToBool() {
for _, dt := range []arrow.DataType{arrow.BinaryTypes.String,
arrow.BinaryTypes.LargeString} {
c.checkCast(dt, arrow.FixedWidthTypes.Boolean,
diff --git a/arrow/compute/internal/kernels/boolean_cast.go
b/arrow/compute/internal/kernels/boolean_cast.go
index dbe96f10..4232b0df 100644
--- a/arrow/compute/internal/kernels/boolean_cast.go
+++ b/arrow/compute/internal/kernels/boolean_cast.go
@@ -35,6 +35,12 @@ func isNonZero[T arrow.FixedWidthType](ctx *exec.KernelCtx,
in []T, out []byte)
return nil
}
+func numericToBoolKernel[T arrow.NumericType](typ arrow.Type)
func(*exec.KernelCtx, []T, []byte) error {
+ return func(ctx *exec.KernelCtx, in []T, out []byte) error {
+ return numericToBoolNeon(typ, ctx, in, out)
+ }
+}
+
// GetBooleanCastKernels returns the slice of scalar kernels for casting
// values *to* a boolean type.
func GetBooleanCastKernels() []exec.ScalarKernel {
@@ -55,17 +61,17 @@ func GetBooleanCastKernels() []exec.ScalarKernel {
case arrow.UINT16:
ex = ScalarUnaryBoolOutput(isNonZero[uint16])
case arrow.INT32:
- ex = ScalarUnaryBoolOutput(isNonZero[int32])
+ ex =
ScalarUnaryBoolOutput(numericToBoolKernel[int32](arrow.INT32))
case arrow.UINT32:
- ex = ScalarUnaryBoolOutput(isNonZero[uint32])
+ ex =
ScalarUnaryBoolOutput(numericToBoolKernel[uint32](arrow.UINT32))
case arrow.INT64:
- ex = ScalarUnaryBoolOutput(isNonZero[int64])
+ ex =
ScalarUnaryBoolOutput(numericToBoolKernel[int64](arrow.INT64))
case arrow.UINT64:
- ex = ScalarUnaryBoolOutput(isNonZero[uint64])
+ ex =
ScalarUnaryBoolOutput(numericToBoolKernel[uint64](arrow.UINT64))
case arrow.FLOAT32:
- ex = ScalarUnaryBoolOutput(isNonZero[float32])
+ ex =
ScalarUnaryBoolOutput(numericToBoolKernel[float32](arrow.FLOAT32))
case arrow.FLOAT64:
- ex = ScalarUnaryBoolOutput(isNonZero[float64])
+ ex =
ScalarUnaryBoolOutput(numericToBoolKernel[float64](arrow.FLOAT64))
}
k := exec.NewScalarKernel(
[]exec.InputType{exec.NewExactInput(ty)}, out, ex, nil)
diff --git a/arrow/compute/internal/kernels/boolean_cast_neon_arm64.go
b/arrow/compute/internal/kernels/boolean_cast_neon_arm64.go
new file mode 100644
index 00000000..78211f29
--- /dev/null
+++ b/arrow/compute/internal/kernels/boolean_cast_neon_arm64.go
@@ -0,0 +1,47 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build go1.18 && arm64 && !noasm && !appengine
+
+package kernels
+
+import (
+ "unsafe"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/bitutil"
+ "github.com/apache/arrow-go/v18/arrow/compute/exec"
+ "golang.org/x/sys/cpu"
+)
+
+func numericToBoolNeon[T arrow.NumericType](typ arrow.Type, ctx
*exec.KernelCtx, in []T, out []byte) error {
+ if !cpu.ARM64.HasASIMD {
+ return isNonZero(ctx, in, out)
+ }
+
+ var zero T
+ bulk := len(in) &^ 7
+ if bulk != 0 {
+ left := arrow.GetBytes(in[:bulk])
+ right := unsafe.Slice((*byte)(unsafe.Pointer(&zero)),
int(unsafe.Sizeof(zero)))
+ _comparison_neon(int(typ), int(CmpNE), neonCompareArrayScalar,
+ unsafe.Pointer(&left[0]), unsafe.Pointer(&right[0]),
unsafe.Pointer(&out[0]), int64(bulk/8))
+ }
+ for i, v := range in[bulk:] {
+ bitutil.SetBitTo(out, bulk+i, v != zero)
+ }
+ return nil
+}
diff --git a/arrow/compute/internal/kernels/boolean_cast_noasm.go
b/arrow/compute/internal/kernels/boolean_cast_noasm.go
new file mode 100644
index 00000000..2bdd800c
--- /dev/null
+++ b/arrow/compute/internal/kernels/boolean_cast_noasm.go
@@ -0,0 +1,28 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build go1.18 && (noasm || !arm64 || appengine)
+
+package kernels
+
+import (
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/compute/exec"
+)
+
+func numericToBoolNeon[T arrow.NumericType](_ arrow.Type, ctx *exec.KernelCtx,
in []T, out []byte) error {
+ return isNonZero(ctx, in, out)
+}
