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 cde116fa perf(arrow/compute): read flat take values directly (#1360)
cde116fa is described below
commit cde116fae951b080e51f2bac4fafef47185a3334
Author: Minh Vu <[email protected]>
AuthorDate: Tue Oct 13 16:31:41 2026 +0200
perf(arrow/compute): read flat take values directly (#1360)
### What
- Faster null-free primitive `take` with random and reverse indices.
### Why
The sorted path already reads values directly. Other orders call a
getter for every row.
Bounds-checked `compute.Take`, int64 values and indices. Apple M1 Pro,
Go 1.26.3, GOMAXPROCS=1. Medians from 12 x 200ms interleaved samples per
revision. Both revisions use the corrected random workload.
| Rows / order | Before | After | Time |
| --- | ---: | ---: | ---: |
| 50K / random | 298.2 us | 232.1 us | -22.2% |
| 50K / reverse | 299.4 us | 211.6 us | -29.3% |
| 100K / random | 542.8 us | 377.8 us | -30.4% |
| 100K / reverse | 565.4 us | 399.6 us | -29.3% |
Sorted and nullable controls are similar or faster. Allocation counts
are unchanged.
### Implementation
- Read the flat value slice directly after the existing null checks.
- Share the existing unrolled copy loop.
- Keep reverse indices sequential.
- Fix the random benchmark to use seeded PCG indices.
---
arrow/compute/internal/kernels/vector_selection.go | 42 +++--
arrow/compute/take_primitive_test.go | 209 +++++++++++++++++++++
arrow/compute/vector_selection_test.go | 31 +--
3 files changed, 257 insertions(+), 25 deletions(-)
diff --git a/arrow/compute/internal/kernels/vector_selection.go
b/arrow/compute/internal/kernels/vector_selection.go
index 855551ca..20a8379a 100644
--- a/arrow/compute/internal/kernels/vector_selection.go
+++ b/arrow/compute/internal/kernels/vector_selection.go
@@ -838,6 +838,19 @@ func isReverseSorted[IdxT arrow.UintType](indices []IdxT)
bool {
return true
}
+func takePrimitiveValues[IdxT arrow.UintType, ValT arrow.IntType](values
[]ValT, indices []IdxT, out []ValT) {
+ i := 0
+ for ; i+4 <= len(indices); i += 4 {
+ out[i] = values[indices[i]]
+ out[i+1] = values[indices[i+1]]
+ out[i+2] = values[indices[i+2]]
+ out[i+3] = values[indices[i+3]]
+ }
+ for ; i < len(indices); i++ {
+ out[i] = values[indices[i]]
+ }
+}
+
// primitiveTakeImplSorted is optimized for sorted (monotonically increasing)
indices
// This enables better CPU cache utilization and branch prediction
func primitiveTakeImplSorted[IdxT arrow.UintType, ValT arrow.IntType](values
primitiveGetter[ValT], indices *exec.ArraySpan, out *exec.ExecResult) {
@@ -851,18 +864,7 @@ func primitiveTakeImplSorted[IdxT arrow.UintType, ValT
arrow.IntType](values pri
// Try to access underlying values directly for better
performance
if valImpl, ok := values.(*primitiveGetterImpl[ValT]); ok {
// Direct memory access for primitiveGetterImpl
- valData := valImpl.values
- // Unroll loop for better performance
- i := 0
- for ; i+4 <= len(indicesData); i += 4 {
- outData[i] = valData[indicesData[i]]
- outData[i+1] = valData[indicesData[i+1]]
- outData[i+2] = valData[indicesData[i+2]]
- outData[i+3] = valData[indicesData[i+3]]
- }
- for ; i < len(indicesData); i++ {
- outData[i] = valData[indicesData[i]]
- }
+ takePrimitiveValues(valImpl.values, indicesData,
outData)
} else {
// Fallback to GetValue interface
for i, idx := range indicesData {
@@ -932,14 +934,26 @@ func primitiveTakeImpl[IdxT arrow.UintType, ValT
arrow.IntType](values primitive
// Check for reverse sorted - use sequential loop to
avoid cache penalties
// Loop unrolling amplifies cache miss penalties in
reverse access patterns
if isReverseSorted(indicesData) {
- for i := 0; i < len(indicesData); i++ {
- outData[i] =
values.GetValue(int64(indicesData[i]))
+ if valImpl, ok :=
values.(*primitiveGetterImpl[ValT]); ok {
+ for i, idx := range indicesData {
+ outData[i] = valImpl.values[idx]
+ }
+ } else {
+ for i := 0; i < len(indicesData); i++ {
+ outData[i] =
values.GetValue(int64(indicesData[i]))
+ }
}
out.Nulls = 0
return
}
}
+ if valImpl, ok := values.(*primitiveGetterImpl[ValT]); ok {
+ takePrimitiveValues(valImpl.values, indicesData,
outData)
+ out.Nulls = 0
+ return
+ }
+
// Unroll loop for better performance (random access patterns)
i := 0
for ; i+4 <= len(indicesData); i += 4 {
diff --git a/arrow/compute/take_primitive_test.go
b/arrow/compute/take_primitive_test.go
new file mode 100644
index 00000000..7e7d54d7
--- /dev/null
+++ b/arrow/compute/take_primitive_test.go
@@ -0,0 +1,209 @@
+// 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"
+ "strings"
+ "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 TestPrimitiveTakeAccessPatterns(t *testing.T) {
+ const nvalues = 97
+ for _, valueType := range numericTypes {
+ t.Run(valueType.String(), func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+ valueStrings := make([]string, nvalues+6)
+ for i := range valueStrings {
+ valueStrings[i] = "null"
+ }
+ for i := 0; i < nvalues; i++ {
+ valueStrings[i+3] = fmt.Sprint((i*13 + 7) %
nvalues)
+ }
+ full, _, err := array.FromJSON(mem, valueType,
strings.NewReader("["+strings.Join(valueStrings, ",")+"]"))
+ require.NoError(t, err)
+ values := array.NewSlice(full, 3, nvalues+3)
+ full.Release()
+ defer values.Release()
+ require.Zero(t, values.NullN())
+
+ for _, size := range []int{0, 1, 3, 4, 5, 31, 32, 33,
255, 256, 257} {
+ for _, order := range []string{"sorted",
"reverse", "random", "repeated"} {
+ t.Run(fmt.Sprintf("size=%d/%s", size,
order), func(t *testing.T) {
+ indexStrings := make([]string,
size+2)
+ indexStrings[0],
indexStrings[size+1] = "99", "99"
+ wantStrings := make([]string,
size)
+ for i := 0; i < size; i++ {
+ idx := 0
+ switch order {
+ case "sorted":
+ idx = i *
(nvalues - 1) / size
+ case "reverse":
+ idx = (nvalues
- 1) - i*(nvalues-1)/size
+ case "random":
+ idx = (i*37 +
19) % nvalues
+ case "repeated":
+ if (i/4)%2 == 0
{
+ idx =
nvalues - 1
+ }
+ }
+ indexStrings[i+1] =
fmt.Sprint(idx)
+ wantStrings[i] =
valueStrings[idx+3]
+ }
+ expected, _, err :=
array.FromJSON(mem, valueType, strings.NewReader("["+strings.Join(wantStrings,
",")+"]"))
+ require.NoError(t, err)
+ defer expected.Release()
+
+ for _, indexType := range
chunkedTakeIndexTypes {
+
t.Run(indexType.String(), func(t *testing.T) {
+ indicesFull :=
makeChunkedTakeIndexArray(t, mem, indexType, indexStrings, nil)
+ indices :=
array.NewSlice(indicesFull, 1, int64(size+1))
+
indicesFull.Release()
+ defer
indices.Release()
+ actual, err :=
compute.TakeArray(ctx, values, indices)
+
require.NoError(t, err)
+ defer
actual.Release()
+ require.True(t,
array.Equal(expected, actual))
+ require.Zero(t,
actual.NullN())
+
require.Empty(t, actual.NullBitmapBytes())
+ })
+ }
+ })
+ }
+ }
+ })
+ }
+}
+
+func TestPrimitiveTakeFloatBits(t *testing.T) {
+ for _, tc := range []struct {
+ typ arrow.DataType
+ raw []byte
+ }{
+ {arrow.PrimitiveTypes.Float32,
arrow.Uint32Traits.CastToBytes([]uint32{
+ 0, 0, 0x80000000, 0x7f800000, 0xff800000, 0x7fc12345,
0x7f812345, 1, 0xffc12345, 0,
+ })},
+ {arrow.PrimitiveTypes.Float64,
arrow.Uint64Traits.CastToBytes([]uint64{
+ 0, 0, 0x8000000000000000, 0x7ff0000000000000,
0xfff0000000000000,
+ 0x7ff8123456789abc, 0x7ff0123456789abc, 1,
0xfff8123456789abc, 0,
+ })},
+ } {
+ t.Run(tc.typ.String(), func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+ width := tc.typ.(arrow.FixedWidthDataType).Bytes()
+ buffer := memory.NewResizableBuffer(mem)
+ buffer.Resize(len(tc.raw))
+ copy(buffer.Bytes(), tc.raw)
+ data := array.NewData(tc.typ, len(tc.raw)/width,
[]*memory.Buffer{nil, buffer}, nil, 0, 0)
+ buffer.Release()
+ full := array.MakeFromData(data)
+ data.Release()
+ values := array.NewSlice(full, 1, 9)
+ full.Release()
+ defer values.Release()
+
+ for _, order := range []string{"random", "reverse"} {
+ t.Run(order, func(t *testing.T) {
+ bldr := array.NewInt32Builder(mem)
+ bldr.Append(99)
+ var expected []byte
+ for i := 0; i < 33; i++ {
+ idx := (i*5 + 3) % values.Len()
+ if order == "reverse" {
+ idx = values.Len() - 1
- i*values.Len()/33
+ }
+ bldr.Append(int32(idx))
+ expected = append(expected,
tc.raw[(idx+1)*width:(idx+2)*width]...)
+ }
+ bldr.Append(99)
+ indicesFull := bldr.NewArray()
+ bldr.Release()
+ indices := array.NewSlice(indicesFull,
1, 34)
+ indicesFull.Release()
+ defer indices.Release()
+ actual, err := compute.TakeArray(ctx,
values, indices)
+ require.NoError(t, err)
+ defer actual.Release()
+ start := actual.Data().Offset() * width
+ require.Equal(t, expected,
actual.Data().Buffers()[1].Bytes()[start:start+len(expected)])
+ })
+ }
+ })
+ }
+}
+
+func TestPrimitiveTakeNullIndexPayloads(t *testing.T) {
+ for _, size := range []int{31, 32, 33, 257} {
+ t.Run(fmt.Sprint(size), func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+ fullValues, _, err := array.FromJSON(mem,
arrow.PrimitiveTypes.Int64, strings.NewReader(`[99, 101, 202, 303, 99]`))
+ require.NoError(t, err)
+ values := array.NewSlice(fullValues, 1, 4)
+ fullValues.Release()
+ defer values.Release()
+
+ rawIndices := make([]int64, size+2)
+ valid := make([]bool, size+2)
+ expectedBuilder := array.NewInt64Builder(mem)
+ defer expectedBuilder.Release()
+ for i := 0; i < size; i++ {
+ if i == 0 || i == size/2 || i == size-1 {
+ // Null slots may contain arbitrary
bits, including invalid indices.
+ rawIndices[i+1] = -1 << 63
+ expectedBuilder.AppendNull()
+ } else {
+ rawIndices[i+1] = int64(i % 3)
+ valid[i+1] = true
+ expectedBuilder.Append(int64((i%3 + 1)
* 101))
+ }
+ }
+ indicesBuilder := array.NewInt64Builder(mem)
+ indicesBuilder.AppendValues(rawIndices, valid)
+ fullIndices := indicesBuilder.NewArray()
+ indicesBuilder.Release()
+ indices := array.NewSlice(fullIndices, 1, int64(size+1))
+ fullIndices.Release()
+ defer indices.Release()
+ expected := expectedBuilder.NewArray()
+ defer expected.Release()
+
+ for _, boundsCheck := range []bool{false, true} {
+ t.Run(fmt.Sprintf("bounds-check=%t",
boundsCheck), func(t *testing.T) {
+ actual, err :=
compute.TakeArrayOpts(ctx, values, indices, compute.TakeOptions{BoundsCheck:
boundsCheck})
+ require.NoError(t, err)
+ defer actual.Release()
+ require.True(t, array.Equal(expected,
actual))
+ })
+ }
+ })
+ }
+}
diff --git a/arrow/compute/vector_selection_test.go
b/arrow/compute/vector_selection_test.go
index 1e9233e6..6d18b89b 100644
--- a/arrow/compute/vector_selection_test.go
+++ b/arrow/compute/vector_selection_test.go
@@ -14,13 +14,14 @@
// See the License for the specific language governing permissions and
// limitations under the License.
-//go:build go1.18
+//go:build go1.24
package compute_test
import (
"context"
"fmt"
+ "math/rand/v2"
"strings"
"testing"
@@ -2183,22 +2184,26 @@ func BenchmarkTakePrimitive(b *testing.B) {
}
default: // "random"
// Random indices
+ rng := rand.New(rand.NewPCG(randomSeed, 0))
for i := int64(0); i < bm.numRows; i++ {
- indicesBldr.Append(i % bm.numRows)
+
indicesBldr.Append(rng.Int64N(bm.numRows))
}
}
indices := indicesBldr.NewArray()
defer indices.Release()
- b.ReportMetric(float64(bm.numRows), "rows/sec")
- b.ResetTimer()
- for i := 0; i < b.N; i++ {
- result, err := compute.Take(ctx,
*compute.DefaultTakeOptions(), &compute.ArrayDatum{values.Data()},
&compute.ArrayDatum{indices.Data()})
+ opts := *compute.DefaultTakeOptions()
+ valuesDatum := &compute.ArrayDatum{Value: values.Data()}
+ indicesDatum := &compute.ArrayDatum{Value:
indices.Data()}
+ b.ReportAllocs()
+ for b.Loop() {
+ result, err := compute.Take(ctx, opts,
valuesDatum, indicesDatum)
if err != nil {
b.Fatal(err)
}
result.Release()
}
+ b.ReportMetric(float64(bm.numRows), "rows/op")
})
}
}
@@ -2246,22 +2251,26 @@ func BenchmarkTakePrimitiveWithNulls(b *testing.B) {
indicesBldr.Append(i)
}
default: // "random"
+ rng := rand.New(rand.NewPCG(randomSeed, 0))
for i := int64(0); i < bm.numRows; i++ {
- indicesBldr.Append((i * 1103515245) %
bm.numRows)
+
indicesBldr.Append(rng.Int64N(bm.numRows))
}
}
indices := indicesBldr.NewArray()
defer indices.Release()
- b.ReportMetric(float64(bm.numRows), "rows/sec")
- b.ResetTimer()
- for i := 0; i < b.N; i++ {
- result, err := compute.Take(ctx,
*compute.DefaultTakeOptions(), &compute.ArrayDatum{values.Data()},
&compute.ArrayDatum{indices.Data()})
+ opts := *compute.DefaultTakeOptions()
+ valuesDatum := &compute.ArrayDatum{Value: values.Data()}
+ indicesDatum := &compute.ArrayDatum{Value:
indices.Data()}
+ b.ReportAllocs()
+ for b.Loop() {
+ result, err := compute.Take(ctx, opts,
valuesDatum, indicesDatum)
if err != nil {
b.Fatal(err)
}
result.Release()
}
+ b.ReportMetric(float64(bm.numRows), "rows/op")
})
}
}
