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 b02cc011 perf(arrow/compute): avoid getter wrappers in chunked binary
take (#1359)
b02cc011 is described below
commit b02cc0111095ac8c1659300c2a276354ddd29294
Author: Minh Vu <[email protected]>
AuthorDate: Tue Oct 6 21:54:13 2026 +0200
perf(arrow/compute): avoid getter wrappers in chunked binary take (#1359)
### What
- Faster chunked string and binary `take`.
### Why
Each lookup goes through an interface and a generic method wrapper.
Apple M1 Pro, Go 1.26.3. 64 chunks x 4,096 rows, 8-byte values, 10%
selected, sequential indices. Medians from 12 x 100ms samples per
revision.
| String values | Before | After | Time |
| --- | ---: | ---: | ---: |
| No nulls | 504.8 us | 484.8 us | -4.0% |
| 1% nulls | 661.9 us | 632.4 us | -4.5% |
| 50% nulls | 463.8 us | 440.7 us | -5.0% |
| All nulls | 309.3 us | 287.1 us | -7.2% |
Random access is mostly unchanged in 20 x 500ms samples. Allocations are
unchanged.
### Implementation
- Pass the concrete getter through dispatch.
- Call value and validity helpers directly.
---
arrow/compute/chunked_take_nulls_benchmark_test.go | 97 +++++++++++++++
arrow/compute/chunked_take_nulls_test.go | 132 +++++++++++++++++++++
arrow/compute/internal/kernels/vector_selection.go | 12 +-
3 files changed, 231 insertions(+), 10 deletions(-)
diff --git a/arrow/compute/chunked_take_nulls_benchmark_test.go
b/arrow/compute/chunked_take_nulls_benchmark_test.go
new file mode 100644
index 00000000..bdc53d44
--- /dev/null
+++ b/arrow/compute/chunked_take_nulls_benchmark_test.go
@@ -0,0 +1,97 @@
+// 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"
+ "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 BenchmarkTakeChunkedBinaryNulls(b *testing.B) {
+ for _, typ := range chunkedTakeBinaryTypes {
+ for _, numChunks := range []int{8, 64} {
+ for _, nullPattern := range []string{"none", "sparse",
"alternating", "all"} {
+ for _, order := range []string{"sequential",
"random"} {
+ name :=
fmt.Sprintf("%s/chunks=%d/nulls=%s/%s", typ, numChunks, nullPattern, order)
+ b.Run(name, func(b *testing.B) {
+ mem := memory.DefaultAllocator
+ ctx :=
compute.WithAllocator(context.Background(), mem)
+ const rowsPerChunk = 4096
+ chunks := make([]arrow.Array,
numChunks)
+ for i := range chunks {
+ bldr :=
array.NewBinaryBuilder(mem, typ.(arrow.BinaryDataType))
+
bldr.Reserve(rowsPerChunk)
+
bldr.ReserveData(rowsPerChunk * 8)
+ for j := 0; j <
rowsPerChunk; j++ {
+ isNull :=
nullPattern == "all" ||
+
(nullPattern == "sparse" && j%100 == 0) ||
+
(nullPattern == "alternating" && j%2 == 0)
+ if isNull {
+
bldr.AppendNull()
+ } else {
+
bldr.Append([]byte("value123"))
+ }
+ }
+ chunks[i] =
bldr.NewArray()
+ bldr.Release()
+ }
+ values := arrow.NewChunked(typ,
chunks)
+ defer values.Release()
+ for _, chunk := range chunks {
+ chunk.Release()
+ }
+
+ totalRows := numChunks *
rowsPerChunk
+ indicesBldr :=
array.NewInt64Builder(mem)
+ indicesBldr.Reserve(totalRows /
10)
+ for i := 0; i < totalRows/10;
i++ {
+ idx := int64(i)
+ if order == "random" {
+ idx =
(idx*1103515245 + 12345) % int64(totalRows)
+ }
+ indicesBldr.Append(idx)
+ }
+ indices :=
indicesBldr.NewArray()
+ indicesBldr.Release()
+ defer indices.Release()
+ valuesDatum :=
&compute.ChunkedDatum{Value: values}
+ indicesDatum :=
&compute.ArrayDatum{Value: indices.Data()}
+ opts :=
*compute.DefaultTakeOptions()
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ result, err :=
compute.Take(ctx, opts, valuesDatum, indicesDatum)
+ if err != nil {
+ b.Fatal(err)
+ }
+ result.Release()
+ }
+ })
+ }
+ }
+ }
+ }
+}
diff --git a/arrow/compute/chunked_take_nulls_test.go
b/arrow/compute/chunked_take_nulls_test.go
new file mode 100644
index 00000000..ae8d5a92
--- /dev/null
+++ b/arrow/compute/chunked_take_nulls_test.go
@@ -0,0 +1,132 @@
+// 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"
+ "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"
+ "github.com/stretchr/testify/require"
+)
+
+func TestChunkedBinaryTakeNullPatterns(t *testing.T) {
+ for _, typ := range chunkedTakeBinaryTypes {
+ for _, nullPattern := range []string{"none", "sparse",
"alternating", "all", "mixed_chunks"} {
+ for _, prefix := range []int{1, 3, 7, 8, 15} {
+ t.Run(fmt.Sprintf("%s/%s/offset=%d", typ,
nullPattern, prefix), func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx :=
compute.WithAllocator(context.Background(), mem)
+ const rowsPerChunk = 9
+ logicalValues := make([][]byte,
2*rowsPerChunk)
+ logicalValid := make([]bool,
len(logicalValues))
+ chunks := make([]arrow.Array, 0, 5)
+ empty := newTestBinaryArray(mem, typ,
nil, nil)
+ defer empty.Release()
+ for chunkIndex := 0; chunkIndex < 2;
chunkIndex++ {
+ values := make([][]byte,
prefix+rowsPerChunk+1)
+ valid := make([]bool,
len(values))
+ for i := range values {
+ values[i] =
[]byte("padding")
+ valid[i] = true
+ }
+ for i := 0; i < rowsPerChunk;
i++ {
+ globalIndex :=
chunkIndex*rowsPerChunk + i
+ value :=
[]byte(fmt.Sprintf("row-%d", globalIndex))
+ if i == 0 || i ==
rowsPerChunk-1 {
+ value = nil
+ }
+ isNull := nullPattern
== "all" ||
+ (nullPattern ==
"sparse" && i == 1) ||
+ (nullPattern ==
"alternating" && i%2 == 0) ||
+ (nullPattern ==
"mixed_chunks" && chunkIndex == 1 && i%2 == 0)
+ values[prefix+i] = value
+ valid[prefix+i] =
!isNull
+
logicalValues[globalIndex] = value
+
logicalValid[globalIndex] = !isNull
+ }
+ full := newTestBinaryArray(mem,
typ, values, valid)
+ if nullPattern ==
"mixed_chunks" && chunkIndex == 0 {
+ buffers :=
append([]*memory.Buffer(nil), full.Data().Buffers()...)
+ buffers[0] = nil
+ data :=
array.NewData(typ, full.Len(), buffers, nil, 0, 0)
+ full.Release()
+ full =
array.MakeFromData(data)
+ data.Release()
+ }
+ sliced := array.NewSlice(full,
int64(prefix), int64(prefix+rowsPerChunk))
+ full.Release()
+ chunks = append(chunks, empty,
sliced)
+ }
+ chunks = append(chunks, empty)
+ values := arrow.NewChunked(typ, chunks)
+ defer values.Release()
+ for _, chunk := range chunks {
+ if chunk != empty {
+ chunk.Release()
+ }
+ }
+
+ for _, indexType := range
chunkedTakeIndexTypes {
+ t.Run(indexType.String(),
func(t *testing.T) {
+ selected := []int{17,
0, 9, 8, 10, 9, 1, 16, 3, 17}
+ indicesValues :=
make([]string, len(selected)+2)
+ indicesValid :=
make([]bool, len(indicesValues))
+ for i := range
indicesValues {
+
indicesValues[i] = "99"
+ indicesValid[i]
= true
+ }
+ wantValues :=
make([][]byte, len(selected))
+ wantValid :=
make([]bool, len(selected))
+ for i, idx := range
selected {
+ if i == 4 {
+
indicesValid[i+1] = false
+ continue
+ }
+
indicesValues[i+1] = fmt.Sprint(idx)
+ wantValues[i] =
logicalValues[idx]
+ wantValid[i] =
logicalValid[idx]
+ }
+ indicesFull :=
makeChunkedTakeIndexArray(t, mem, indexType, indicesValues, indicesValid)
+ defer
indicesFull.Release()
+ indices :=
array.NewSlice(indicesFull, 1, int64(1+len(selected)))
+ defer indices.Release()
+ result, err :=
compute.Take(ctx, *compute.DefaultTakeOptions(),
+
&compute.ChunkedDatum{Value: values}, &compute.ArrayDatum{Value:
indices.Data()})
+ require.NoError(t, err)
+ defer result.Release()
+ expectedArray :=
newTestBinaryArray(mem, typ, wantValues, wantValid)
+ defer
expectedArray.Release()
+ expected :=
arrow.NewChunked(typ, []arrow.Array{expectedArray})
+ defer expected.Release()
+ actual :=
result.(*compute.ChunkedDatum).Value
+ require.True(t,
array.ChunkedEqual(expected, actual))
+ require.Equal(t,
expected.NullN(), actual.NullN())
+ })
+ }
+ })
+ }
+ }
+ }
+}
diff --git a/arrow/compute/internal/kernels/vector_selection.go
b/arrow/compute/internal/kernels/vector_selection.go
index 174148f2..855551ca 100644
--- a/arrow/compute/internal/kernels/vector_selection.go
+++ b/arrow/compute/internal/kernels/vector_selection.go
@@ -700,14 +700,6 @@ func (c *chunkedPrimitiveGetter[T]) GetValue(i int64) T {
func (c *chunkedPrimitiveGetter[T]) NullCount() int64 { return c.nulls }
func (c *chunkedPrimitiveGetter[T]) Len() int64 { return c.len }
-type binaryGetter interface {
- IsValid(int64) bool
- GetValue(int64) []byte
- NullCount() int64
- Len() int64
- DataLen() int64
-}
-
type chunkedBinaryGetter[OffsetT int32 | int64] struct {
resolver *exec.ChunkResolver
offsets [][]OffsetT
@@ -1285,7 +1277,7 @@ func checkBinaryTakeOffset[OffsetT int32 | int64](offset
OffsetT, valueLen int64
return nil
}
-func takeChunkedBinaryImpl[IdxT arrow.UintType, OffsetT int32 | int64](ctx
*exec.KernelCtx, indices *exec.ArraySpan, values binaryGetter, out
*exec.ExecResult) error {
+func takeChunkedBinaryImpl[IdxT arrow.UintType, OffsetT int32 | int64](ctx
*exec.KernelCtx, indices *exec.ArraySpan, values *chunkedBinaryGetter[OffsetT],
out *exec.ExecResult) error {
var (
indicesValues = exec.GetSpanValues[IdxT](indices, 1)
indicesIsValid = bitutil.OptionalBitIndexer{Bitmap:
indices.Buffers[0].Buf, Offset: int(indices.Offset)}
@@ -1422,7 +1414,7 @@ func takeChunkedBinaryImpl[IdxT arrow.UintType, OffsetT
int32 | int64](ctx *exec
return nil
}
-func takeChunkedBinaryDispatch[IdxT arrow.UintType, OffsetT int32 | int64](ctx
*exec.KernelCtx, values binaryGetter, indices *arrow.Chunked, out
[]*exec.ExecResult) error {
+func takeChunkedBinaryDispatch[IdxT arrow.UintType, OffsetT int32 | int64](ctx
*exec.KernelCtx, values *chunkedBinaryGetter[OffsetT], indices *arrow.Chunked,
out []*exec.ExecResult) error {
var span exec.ArraySpan
for i, chunk := range indices.Chunks() {
span.SetMembers(chunk.Data())
