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 4dfd78d2 perf(compute): fast-path single-span execution (#1233)
4dfd78d2 is described below
commit 4dfd78d220f814adb3a32a5ad9ddf27f0fa034d9
Author: Minh Vu <[email protected]>
AuthorDate: Fri Aug 28 18:55:53 2026 +0200
perf(compute): fast-path single-span execution (#1233)
### Rationale for this change
`iterateExecSpans` allocates three tracking slices before it knows
whether the
batch can be processed as one contiguous span. Most non-chunked batches
do not
need that state.
### What changes are included in this PR?
- Detect non-chunked batches that fit in one execution span.
- Return that span directly without allocating chunk iteration state.
- Keep the existing path for chunked batches and batches split by
`maxChunkSize`.
- Preserve sliced array offsets, scalar promotion, and iterator
completion
behavior.
- Add focused tests and a benchmark.
### Benchmark
The benchmark covers one array, two arrays, scalar plus array, all
scalars, and
a chunked control case at lengths 1, 16, 256, and 4096.
Representative results from one run on an Apple M1 Pro with Go 1.26.3:
| Case | Before | After | Allocations |
| --- | ---: | ---: | ---: |
| Array, length 1 | 342.0 ns/op, 528 B/op | 189.0 ns/op, 304 B/op | 8 ->
5 |
| Two arrays, length 1 | 451.4 ns/op, 776 B/op | 286.6 ns/op, 528 B/op |
8 -> 5 |
| Scalar plus array, length 1 | 433.8 ns/op, 776 B/op | 260.8 ns/op, 528
B/op | 8 -> 5 |
| All scalars | 468.1 ns/op, 784 B/op | 281.4 ns/op, 528 B/op | 10 -> 7
|
| Chunked control | 356.6 ns/op, 528 B/op | 339.2 ns/op, 528 B/op | 8 ->
8 |
Command:
```bash
go test ./arrow/compute -run '^$' -bench '^BenchmarkIterateExecSpans$'
-benchmem -benchtime=100ms -count=1 -cpu=1
```
### Are these changes tested?
- `go test ./arrow/compute -count=1`
- `go test -race ./arrow/compute -count=1`
---
arrow/compute/exec_spans_benchmark_test.go | 109 +++++++++++++++++++++++++++++
arrow/compute/exec_spans_test.go | 85 ++++++++++++++++++++++
arrow/compute/executor.go | 37 +++++++---
3 files changed, 222 insertions(+), 9 deletions(-)
diff --git a/arrow/compute/exec_spans_benchmark_test.go
b/arrow/compute/exec_spans_benchmark_test.go
new file mode 100644
index 00000000..070c8fd9
--- /dev/null
+++ b/arrow/compute/exec_spans_benchmark_test.go
@@ -0,0 +1,109 @@
+// 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
+
+import (
+ "fmt"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+)
+
+func BenchmarkIterateExecSpans(b *testing.B) {
+ for _, length := range []int64{1, 16, 256, 4096} {
+ b.Run(fmt.Sprintf("array/%d", length), func(b *testing.B) {
+ batch := newExecSpanBenchmarkBatch(b, length, 1, false)
+ benchmarkIterateExecSpans(b, batch, DefaultMaxChunkSize)
+ })
+
+ b.Run(fmt.Sprintf("binary/%d", length), func(b *testing.B) {
+ batch := newExecSpanBenchmarkBatch(b, length, 2, false)
+ benchmarkIterateExecSpans(b, batch, DefaultMaxChunkSize)
+ })
+
+ b.Run(fmt.Sprintf("scalar-array/%d", length), func(b
*testing.B) {
+ batch := newExecSpanBenchmarkBatch(b, length, 1, true)
+ benchmarkIterateExecSpans(b, batch, DefaultMaxChunkSize)
+ })
+ }
+
+ b.Run("all-scalars", func(b *testing.B) {
+ batch := &ExecBatch{
+ Values: []Datum{NewDatum(int32(1)), NewDatum(int32(2))},
+ Len: 1,
+ }
+ benchmarkIterateExecSpans(b, batch, DefaultMaxChunkSize)
+ })
+
+ b.Run("chunked-control", func(b *testing.B) {
+ batch := newExecSpanBenchmarkBatch(b, 4096, 1, false)
+ arr := batch.Values[0].(*ArrayDatum).Value
+ chunk := array.MakeFromData(arr)
+ chunked := arrow.NewChunked(arrow.PrimitiveTypes.Int32,
[]arrow.Array{chunk})
+ chunk.Release()
+ b.Cleanup(chunked.Release)
+ batch.Values[0] = &ChunkedDatum{Value: chunked}
+ benchmarkIterateExecSpans(b, batch, DefaultMaxChunkSize)
+ })
+}
+
+func benchmarkIterateExecSpans(b *testing.B, batch *ExecBatch, maxChunkSize
int64) {
+ b.Helper()
+ b.ReportAllocs()
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ _, iter, err := iterateExecSpans(batch, maxChunkSize, true)
+ if err != nil {
+ b.Fatal(err)
+ }
+
+ span, pos, ok := iter()
+ if !ok || pos != batch.Len || span.Len != batch.Len {
+ b.Fatalf("unexpected first span: len=%d pos=%d ok=%t",
span.Len, pos, ok)
+ }
+
+ _, pos, ok = iter()
+ if ok || pos != batch.Len {
+ b.Fatalf("unexpected second span: pos=%d ok=%t", pos,
ok)
+ }
+ }
+}
+
+func newExecSpanBenchmarkBatch(tb testing.TB, length int64, arrayCount int,
addScalar bool) *ExecBatch {
+ tb.Helper()
+ values := make([]Datum, 0, arrayCount+1)
+ for i := 0; i < arrayCount; i++ {
+ builder := array.NewInt32Builder(memory.DefaultAllocator)
+ builder.Reserve(int(length))
+ for j := int64(0); j < length; j++ {
+ builder.Append(int32(i) + int32(j))
+ }
+ arr := builder.NewInt32Array()
+ builder.Release()
+ tb.Cleanup(arr.Release)
+ values = append(values, &ArrayDatum{Value: arr.Data()})
+ }
+
+ if addScalar {
+ values = append(values, NewDatum(int32(1)))
+ }
+ return &ExecBatch{Values: values, Len: length}
+}
diff --git a/arrow/compute/exec_spans_test.go b/arrow/compute/exec_spans_test.go
new file mode 100644
index 00000000..3361a7bb
--- /dev/null
+++ b/arrow/compute/exec_spans_test.go
@@ -0,0 +1,85 @@
+// 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
+
+import (
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/stretchr/testify/require"
+)
+
+func TestIterateExecSpansSingleSpan(t *testing.T) {
+ const length = 16
+
+ builder := array.NewInt32Builder(memory.DefaultAllocator)
+ builder.AppendValues(make([]int32, length+2), nil)
+ full := builder.NewInt32Array()
+ builder.Release()
+ input := array.NewSlice(full, 1, length+1)
+ full.Release()
+ defer input.Release()
+
+ batch := &ExecBatch{
+ Values: []Datum{
+ &ArrayDatum{Value: input.Data()},
+ NewDatum(int32(5)),
+ },
+ Len: int64(input.Len()),
+ }
+
+ allScalars, iter, err := iterateExecSpans(batch, DefaultMaxChunkSize,
true)
+ require.NoError(t, err)
+ require.False(t, allScalars)
+
+ span, pos, ok := iter()
+ require.True(t, ok)
+ require.EqualValues(t, length, span.Len)
+ require.EqualValues(t, length, pos)
+ require.EqualValues(t, input.Data().Offset(),
span.Values[0].Array.Offset)
+ require.EqualValues(t, input.Len(), span.Values[0].Array.Len)
+ require.True(t, span.Values[1].IsScalar())
+
+ _, pos, ok = iter()
+ require.False(t, ok)
+ require.EqualValues(t, length, pos)
+}
+
+func TestIterateExecSpansPromotesAllScalars(t *testing.T) {
+ batch := &ExecBatch{
+ Values: []Datum{NewDatum(int32(1)), NewDatum(int32(2))},
+ Len: 1,
+ }
+
+ allScalars, iter, err := iterateExecSpans(batch, DefaultMaxChunkSize,
true)
+ require.NoError(t, err)
+ require.True(t, allScalars)
+
+ span, pos, ok := iter()
+ require.True(t, ok)
+ require.EqualValues(t, 1, span.Len)
+ require.EqualValues(t, 1, pos)
+ require.True(t, span.Values[0].IsArray())
+ require.True(t, span.Values[1].IsArray())
+
+ _, pos, ok = iter()
+ require.False(t, ok)
+ require.EqualValues(t, 1, pos)
+}
diff --git a/arrow/compute/executor.go b/arrow/compute/executor.go
index febbec08..c32c9768 100644
--- a/arrow/compute/executor.go
+++ b/arrow/compute/executor.go
@@ -781,36 +781,40 @@ func iterateExecSpans(batch *ExecBatch, maxChunkSize
int64, promoteIfAllScalar b
}
var (
- args = batch.Values
- haveChunked bool
- chunkIdxes = make([]int, len(args))
- valuePositions = make([]int64, len(args))
- valueOffsets = make([]int64, len(args))
- pos, length int64 = 0, batch.Len
+ args = batch.Values
+ pos, length int64 = 0, batch.Len
)
haveAllScalars = checkIfAllScalar(batch)
maxChunkSize = exec.Min(length, maxChunkSize)
span := exec.ExecSpan{Values: make([]exec.ExecValue, len(args)), Len: 0}
+ haveChunked := false
for i, a := range args {
switch arg := a.(type) {
case *ScalarDatum:
span.Values[i].Scalar = arg.Value
case *ArrayDatum:
span.Values[i].Array.SetMembers(arg.Value)
- valueOffsets[i] = int64(arg.Value.Offset())
case *ChunkedDatum:
+ haveChunked = true
// populate from first chunk
carr := arg.Value
if len(carr.Chunks()) > 0 {
arr := carr.Chunk(0).Data()
span.Values[i].Array.SetMembers(arr)
- valueOffsets[i] = int64(arr.Offset())
} else {
// fill as zero len
exec.FillZeroLength(carr.DataType(),
&span.Values[i].Array)
}
- haveChunked = true
+ }
+ }
+
+ singleSpan := !haveChunked && length > 0 && length <= maxChunkSize
+ var valueOffsets []int64
+ if !singleSpan {
+ valueOffsets = make([]int64, len(args))
+ for i := range args {
+ valueOffsets[i] = span.Values[i].Array.Offset
}
}
@@ -818,6 +822,21 @@ func iterateExecSpans(batch *ExecBatch, maxChunkSize
int64, promoteIfAllScalar b
exec.PromoteExecSpanScalars(span)
}
+ if singleSpan {
+ span.Len = length
+ returned := false
+ return haveAllScalars, func() (exec.ExecSpan, int64, bool) {
+ if returned {
+ return exec.ExecSpan{}, length, false
+ }
+ returned = true
+ return span, length, true
+ }, nil
+ }
+
+ chunkIdxes := make([]int, len(args))
+ valuePositions := make([]int64, len(args))
+
nextChunkSpan := func(iterSz int64, span exec.ExecSpan) int64 {
for i := 0; i < len(args) && iterSz > 0; i++ {
// if the argument is not chunked, it's either a scalar
or an array
