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 fbb40237 perf(arrow/compute): avoid goroutines for serial take (#1226)
fbb40237 is described below
commit fbb40237981c9b30043456c120716f5e1a76756b
Author: Minh Vu <[email protected]>
AuthorDate: Mon Aug 31 20:20:55 2026 +0200
perf(arrow/compute): avoid goroutines for serial take (#1226)
## What changed
- Added direct serial execution for Take on record batches and tables
when there is one column or NumParallel is 1 or less.
- Added the same path for chunked index arrays when there is one index
chunk or NumParallel is 1 or less.
- Kept the existing errgroup path for actual parallel work.
## Why
The old path created errgroup and goroutine setup even when the work was
already serial. This adds overhead for small record batches, one-column
tables, and chunked indices.
## Benchmarks
Command:
go test ./arrow/compute -run "^$" -bench
"^BenchmarkTake(RecordSingleColumn|RecordMultiColumnSerial|TableSingleColumn|ArraySingleIndexChunk|ArrayMultiChunkSerial)$"
-benchmem -benchtime=300ms -count=3
Apple M1 Pro, Go 1.26.3, darwin/arm64.
| Benchmark | Before | After |
| --- | ---: | ---: |
| Record, one column | 45 allocs/op, 3403 B/op | 36 allocs/op, 2642 B/op
|
| Record, 16 columns, serial | 480 allocs/op, 41945 B/op | 441
allocs/op, 39380 B/op |
| Table, one column | 52 allocs/op, 3819 B/op | 43 allocs/op, 3042 B/op
|
| Array, one index chunk | 44 allocs/op, 3370 B/op | 35 allocs/op, 2578
B/op |
| Array, 8 index chunks, serial | 247 allocs/op, 17995 B/op | 224
allocs/op, 16136 B/op |
## Tests
- go test ./arrow/compute -count=1
- go test -race ./arrow/compute -count=1
- go test ./arrow/... -count=1
- Full go test ./... -count=1 with the Parquet test data directory
---
arrow/compute/selection.go | 155 +++++++++------
.../selection_take_serial_benchmark_test.go | 210 +++++++++++++++++++++
arrow/compute/selection_take_serial_test.go | 143 ++++++++++++++
3 files changed, 453 insertions(+), 55 deletions(-)
diff --git a/arrow/compute/selection.go b/arrow/compute/selection.go
index ea69f575..6b748daa 100644
--- a/arrow/compute/selection.go
+++ b/arrow/compute/selection.go
@@ -114,6 +114,26 @@ given by "indices". Nulls in "indices" emit null in the
output`,
})
)
+func takeTableColumn(ctx context.Context, opts FunctionOptions, tbl
arrow.Table, indices Datum, cols []arrow.Column, i int) error {
+ inCol := tbl.Column(i)
+ result, err := CallFunction(ctx, "take", opts,
+ &ChunkedDatum{Value: inCol.Data()},
+ indices)
+ if err != nil {
+ return err
+ }
+ defer result.Release()
+ out := result.(ArrayLikeDatum)
+ chunks := out.Chunks()
+ if out.Kind() == KindArray {
+ defer chunks[0].Release()
+ }
+ chk := arrow.NewChunked(out.Type(), chunks)
+ defer chk.Release()
+ cols[i] = *arrow.NewColumn(inCol.Field(), chk)
+ return nil
+}
+
func takeTableImpl(ctx context.Context, opts FunctionOptions, args ...Datum)
(Datum, error) {
tbl := args[0].(*TableDatum).Value
ncols := int(tbl.NumCols())
@@ -124,39 +144,42 @@ func takeTableImpl(ctx context.Context, opts
FunctionOptions, args ...Datum) (Da
}
}()
- eg, cctx := errgroup.WithContext(ctx)
- eg.SetLimit(GetExecCtx(ctx).NumParallel)
- for i := 0; i < ncols; i++ {
- i := i
- eg.Go(func() error {
- inCol := tbl.Column(i)
- result, err := CallFunction(cctx, "take", opts,
- &ChunkedDatum{Value: inCol.Data()},
- args[1])
- if err != nil {
- return err
- }
- defer result.Release()
- out := result.(ArrayLikeDatum)
- chunks := out.Chunks()
- if out.Kind() == KindArray {
- defer chunks[0].Release()
+ numParallel := GetExecCtx(ctx).NumParallel
+ if ncols <= 1 || numParallel <= 1 {
+ for i := 0; i < ncols; i++ {
+ if err := takeTableColumn(ctx, opts, tbl, args[1],
cols, i); err != nil {
+ return nil, err
}
- chk := arrow.NewChunked(out.Type(), chunks)
- defer chk.Release()
- cols[i] = *arrow.NewColumn(inCol.Field(), chk)
- return nil
- })
- }
+ }
+ } else {
+ eg, cctx := errgroup.WithContext(ctx)
+ eg.SetLimit(numParallel)
+ for i := 0; i < ncols; i++ {
+ i := i
+ eg.Go(func() error {
+ return takeTableColumn(cctx, opts, tbl,
args[1], cols, i)
+ })
+ }
- if err := eg.Wait(); err != nil {
- return nil, err
+ if err := eg.Wait(); err != nil {
+ return nil, err
+ }
}
final := array.NewTable(tbl.Schema(), cols, -1)
return &TableDatum{Value: final}, nil
}
+func takeRecordColumn(ctx context.Context, opts FunctionOptions, rb
arrow.RecordBatch, indices Datum, cols []arrow.Array, i int) error {
+ out, err := CallFunction(ctx, "array_take", opts, &ArrayDatum{Value:
rb.Column(i).Data()}, indices)
+ if err != nil {
+ return err
+ }
+ defer out.Release()
+ cols[i] = out.(*ArrayDatum).MakeArray()
+ return nil
+}
+
func takeRecordImpl(ctx context.Context, opts FunctionOptions, args ...Datum)
(Datum, error) {
indices := args[1]
if indices.Kind() == KindChunked {
@@ -169,7 +192,7 @@ func takeRecordImpl(ctx context.Context, opts
FunctionOptions, args ...Datum) (D
}
rb := args[0].(*RecordDatum).Value
- ncols := rb.NumCols()
+ ncols := int(rb.NumCols())
nrows := args[1].(ArrayLikeDatum).Len()
cols := make([]arrow.Array, ncols)
defer func() {
@@ -180,29 +203,41 @@ func takeRecordImpl(ctx context.Context, opts
FunctionOptions, args ...Datum) (D
}
}()
- eg, cctx := errgroup.WithContext(ctx)
- eg.SetLimit(GetExecCtx(ctx).NumParallel)
- for i := range rb.Columns() {
- i := i
- eg.Go(func() error {
- out, err := CallFunction(cctx, "array_take", opts,
&ArrayDatum{Value: rb.Column(i).Data()}, indices)
- if err != nil {
- return err
+ numParallel := GetExecCtx(ctx).NumParallel
+ if ncols <= 1 || numParallel <= 1 {
+ for i := 0; i < ncols; i++ {
+ if err := takeRecordColumn(ctx, opts, rb, indices,
cols, i); err != nil {
+ return nil, err
}
- defer out.Release()
- cols[i] = out.(*ArrayDatum).MakeArray()
- return nil
- })
- }
+ }
+ } else {
+ eg, cctx := errgroup.WithContext(ctx)
+ eg.SetLimit(numParallel)
+ for i := 0; i < ncols; i++ {
+ i := i
+ eg.Go(func() error {
+ return takeRecordColumn(cctx, opts, rb,
indices, cols, i)
+ })
+ }
- if err := eg.Wait(); err != nil {
- return nil, err
+ if err := eg.Wait(); err != nil {
+ return nil, err
+ }
}
outRec := array.NewRecordBatch(rb.Schema(), cols, nrows)
return &RecordDatum{Value: outRec}, nil
}
+func takeArrayChunk(ctx context.Context, opts FunctionOptions, values Datum,
chunk arrow.Array) (arrow.Array, error) {
+ result, err := CallFunction(ctx, "array_take", opts, values,
&ArrayDatum{Value: chunk.Data()})
+ if err != nil {
+ return nil, err
+ }
+ defer result.Release()
+ return result.(*ArrayDatum).MakeArray(), nil
+}
+
func takeArrayImpl(ctx context.Context, opts FunctionOptions, args ...Datum)
(Datum, error) {
switch args[1].Kind() {
case KindArray:
@@ -218,22 +253,32 @@ func takeArrayImpl(ctx context.Context, opts
FunctionOptions, args ...Datum) (Da
}
}()
- eg, cctx := errgroup.WithContext(ctx)
- eg.SetLimit(GetExecCtx(ctx).NumParallel)
- for i := range chunks {
- i := i
- eg.Go(func() error {
- result, err := CallFunction(cctx, "array_take",
opts, args[0], &ArrayDatum{Value: chunks[i].Data()})
+ numParallel := GetExecCtx(ctx).NumParallel
+ if len(chunks) <= 1 || numParallel <= 1 {
+ for i, chunk := range chunks {
+ result, err := takeArrayChunk(ctx, opts,
args[0], chunk)
if err != nil {
- return err
+ return nil, err
}
- defer result.Release()
- out[i] = result.(*ArrayDatum).MakeArray()
- return nil
- })
- }
- if err := eg.Wait(); err != nil {
- return nil, err
+ out[i] = result
+ }
+ } else {
+ eg, cctx := errgroup.WithContext(ctx)
+ eg.SetLimit(numParallel)
+ for i := range chunks {
+ i := i
+ eg.Go(func() error {
+ result, err := takeArrayChunk(cctx,
opts, args[0], chunks[i])
+ if err != nil {
+ return err
+ }
+ out[i] = result
+ return nil
+ })
+ }
+ if err := eg.Wait(); err != nil {
+ return nil, err
+ }
}
return &ChunkedDatum{
Value: arrow.NewChunked(args[0].(*ArrayDatum).Type(),
out)}, nil
diff --git a/arrow/compute/selection_take_serial_benchmark_test.go
b/arrow/compute/selection_take_serial_benchmark_test.go
new file mode 100644
index 00000000..3edec311
--- /dev/null
+++ b/arrow/compute/selection_take_serial_benchmark_test.go
@@ -0,0 +1,210 @@
+// 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.24
+
+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 newTakeInt32Array(mem memory.Allocator, n int) arrow.Array {
+ values := make([]int32, n)
+ for i := range values {
+ values[i] = int32(i)
+ }
+ bldr := array.NewInt32Builder(mem)
+ bldr.AppendValues(values, nil)
+ result := bldr.NewInt32Array()
+ bldr.Release()
+ return result
+}
+
+func newTakeIndices(mem memory.Allocator, n int) arrow.Array {
+ values := make([]int32, n)
+ for i := range values {
+ values[i] = int32(n - i - 1)
+ }
+ bldr := array.NewInt32Builder(mem)
+ bldr.AppendValues(values, nil)
+ result := bldr.NewInt32Array()
+ bldr.Release()
+ return result
+}
+
+func BenchmarkTakeRecordSingleColumn(b *testing.B) {
+ const (
+ nrows = 256
+ nselect = 128
+ )
+
+ mem := memory.NewGoAllocator()
+ ctx := compute.WithAllocator(context.Background(), mem)
+ field := arrow.Field{Name: "value", Type: arrow.PrimitiveTypes.Int32}
+ schema := arrow.NewSchema([]arrow.Field{field}, nil)
+ values := newTakeInt32Array(mem, nrows)
+ defer values.Release()
+ indices := newTakeIndices(mem, nselect)
+ defer indices.Release()
+ batch := array.NewRecordBatch(schema, []arrow.Array{values}, nrows)
+ defer batch.Release()
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for b.Loop() {
+ result, err := compute.Take(ctx, *compute.DefaultTakeOptions(),
+ &compute.RecordDatum{Value: batch},
&compute.ArrayDatum{Value: indices.Data()})
+ if err != nil {
+ b.Fatal(err)
+ }
+ result.Release()
+ }
+}
+
+func BenchmarkTakeRecordMultiColumnSerial(b *testing.B) {
+ const (
+ nrows = 256
+ ncols = 16
+ nselect = 128
+ )
+
+ mem := memory.NewGoAllocator()
+ ctx := serialTakeContext(mem, 1)
+ fields := make([]arrow.Field, ncols)
+ values := make([]arrow.Array, ncols)
+ for i := range fields {
+ fields[i] = arrow.Field{Name: fmt.Sprintf("value_%d", i), Type:
arrow.PrimitiveTypes.Int32}
+ values[i] = newTakeInt32Array(mem, nrows)
+ }
+ schema := arrow.NewSchema(fields, nil)
+ batch := array.NewRecordBatch(schema, values, nrows)
+ for _, value := range values {
+ value.Release()
+ }
+ defer batch.Release()
+ indices := newTakeIndices(mem, nselect)
+ defer indices.Release()
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for b.Loop() {
+ result, err := compute.Take(ctx, *compute.DefaultTakeOptions(),
+ &compute.RecordDatum{Value: batch},
&compute.ArrayDatum{Value: indices.Data()})
+ if err != nil {
+ b.Fatal(err)
+ }
+ result.Release()
+ }
+}
+
+func BenchmarkTakeTableSingleColumn(b *testing.B) {
+ const (
+ nrows = 256
+ nselect = 128
+ )
+
+ mem := memory.NewGoAllocator()
+ ctx := compute.WithAllocator(context.Background(), mem)
+ field := arrow.Field{Name: "value", Type: arrow.PrimitiveTypes.Int32}
+ schema := arrow.NewSchema([]arrow.Field{field}, nil)
+ values := newTakeInt32Array(mem, nrows)
+ defer values.Release()
+ column := arrow.NewColumnFromArr(field, values)
+ table := array.NewTable(schema, []arrow.Column{column}, nrows)
+ column.Release()
+ defer table.Release()
+ indices := newTakeIndices(mem, nselect)
+ defer indices.Release()
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for b.Loop() {
+ result, err := compute.Take(ctx, *compute.DefaultTakeOptions(),
+ &compute.TableDatum{Value: table},
&compute.ArrayDatum{Value: indices.Data()})
+ if err != nil {
+ b.Fatal(err)
+ }
+ result.Release()
+ }
+}
+
+func BenchmarkTakeArraySingleIndexChunk(b *testing.B) {
+ const (
+ nrows = 256
+ nselect = 128
+ )
+
+ mem := memory.NewGoAllocator()
+ ctx := compute.WithAllocator(context.Background(), mem)
+ values := newTakeInt32Array(mem, nrows)
+ defer values.Release()
+ indices := newTakeIndices(mem, nselect)
+ defer indices.Release()
+ chunkedIndices := arrow.NewChunked(indices.DataType(),
[]arrow.Array{indices})
+ defer chunkedIndices.Release()
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for b.Loop() {
+ result, err := compute.Take(ctx, *compute.DefaultTakeOptions(),
+ &compute.ArrayDatum{Value: values.Data()},
&compute.ChunkedDatum{Value: chunkedIndices})
+ if err != nil {
+ b.Fatal(err)
+ }
+ result.Release()
+ }
+}
+
+func BenchmarkTakeArrayMultiChunkSerial(b *testing.B) {
+ const (
+ nrows = 256
+ chunksCount = 8
+ rowsPerChunk = 16
+ )
+
+ mem := memory.NewGoAllocator()
+ ctx := serialTakeContext(mem, 1)
+ values := newTakeInt32Array(mem, nrows)
+ defer values.Release()
+ chunks := make([]arrow.Array, chunksCount)
+ for i := range chunks {
+ chunks[i] = newTakeIndices(mem, rowsPerChunk)
+ }
+ chunkedIndices := arrow.NewChunked(arrow.PrimitiveTypes.Int32, chunks)
+ for _, chunk := range chunks {
+ chunk.Release()
+ }
+ defer chunkedIndices.Release()
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for b.Loop() {
+ result, err := compute.Take(ctx, *compute.DefaultTakeOptions(),
+ &compute.ArrayDatum{Value: values.Data()},
&compute.ChunkedDatum{Value: chunkedIndices})
+ if err != nil {
+ b.Fatal(err)
+ }
+ result.Release()
+ }
+}
diff --git a/arrow/compute/selection_take_serial_test.go
b/arrow/compute/selection_take_serial_test.go
new file mode 100644
index 00000000..9f57c149
--- /dev/null
+++ b/arrow/compute/selection_take_serial_test.go
@@ -0,0 +1,143 @@
+// 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/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func serialTakeContext(mem memory.Allocator, numParallel int) context.Context {
+ execCtx := compute.DefaultExecCtx()
+ execCtx.NumParallel = numParallel
+ return compute.SetExecCtx(compute.WithAllocator(context.Background(),
mem), execCtx)
+}
+
+func TestTakeRecordBatchSerialExecution(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "a", Type: arrow.PrimitiveTypes.Int32, Nullable: true},
+ {Name: "b", Type: arrow.BinaryTypes.String, Nullable: true},
+ }, nil)
+ values, _, err := array.RecordFromJSON(mem, schema, strings.NewReader(`[
+ {"a": 0, "b": "zero"},
+ {"a": 1, "b": "one"},
+ {"a": null, "b": "null"},
+ {"a": 3, "b": "three"}
+ ]`))
+ require.NoError(t, err)
+ defer values.Release()
+ indices, _, err := array.FromJSON(mem, arrow.PrimitiveTypes.Int32,
strings.NewReader(`[3, 1, null, 0]`))
+ require.NoError(t, err)
+ defer indices.Release()
+ expected, _, err := array.RecordFromJSON(mem, schema,
strings.NewReader(`[
+ {"a": 3, "b": "three"},
+ {"a": 1, "b": "one"},
+ {"a": null, "b": null},
+ {"a": 0, "b": "zero"}
+ ]`))
+ require.NoError(t, err)
+ defer expected.Release()
+
+ for _, numParallel := range []int{1, 0, -1} {
+ numParallel := numParallel
+ t.Run(fmt.Sprintf("parallelism-%d", numParallel), func(t
*testing.T) {
+ result, err := compute.Take(serialTakeContext(mem,
numParallel), *compute.DefaultTakeOptions(),
+ &compute.RecordDatum{Value: values},
&compute.ArrayDatum{Value: indices.Data()})
+ require.NoError(t, err)
+ defer result.Release()
+
+ assert.True(t, array.RecordEqual(expected,
result.(*compute.RecordDatum).Value))
+ })
+ }
+}
+
+func TestTakeTableSerialExecution(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "a", Type: arrow.PrimitiveTypes.Int32, Nullable: true},
+ {Name: "b", Type: arrow.BinaryTypes.String, Nullable: true},
+ }, nil)
+ values, err := array.TableFromJSON(mem, schema, []string{
+ `[{"a": 0, "b": "zero"}, {"a": 1, "b": "one"}]`,
+ `[{"a": 2, "b": "two"}, {"a": 3, "b": "three"}]`,
+ })
+ require.NoError(t, err)
+ defer values.Release()
+ indices, _, err := array.FromJSON(mem, arrow.PrimitiveTypes.Int32,
strings.NewReader(`[3, 1, 0]`))
+ require.NoError(t, err)
+ defer indices.Release()
+ expected, err := array.TableFromJSON(mem, schema, []string{
+ `[{"a": 3, "b": "three"}, {"a": 1, "b": "one"}, {"a": 0, "b":
"zero"}]`,
+ })
+ require.NoError(t, err)
+ defer expected.Release()
+
+ for _, numParallel := range []int{1, 0, -1} {
+ numParallel := numParallel
+ t.Run(fmt.Sprintf("parallelism-%d", numParallel), func(t
*testing.T) {
+ result, err := compute.Take(serialTakeContext(mem,
numParallel), *compute.DefaultTakeOptions(),
+ &compute.TableDatum{Value: values},
&compute.ArrayDatum{Value: indices.Data()})
+ require.NoError(t, err)
+ defer result.Release()
+
+ assert.True(t, array.TableEqual(expected,
result.(*compute.TableDatum).Value))
+ })
+ }
+}
+
+func TestTakeArraySerialExecution(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ values, _, err := array.FromJSON(mem, arrow.PrimitiveTypes.Int32,
strings.NewReader(`[10, 20, 30, 40]`))
+ require.NoError(t, err)
+ defer values.Release()
+ indices, err := array.ChunkedFromJSON(mem, arrow.PrimitiveTypes.Int32,
[]string{`[3, 1]`, `[null, 0]`})
+ require.NoError(t, err)
+ defer indices.Release()
+ expected, err := array.ChunkedFromJSON(mem, arrow.PrimitiveTypes.Int32,
[]string{`[40, 20]`, `[null, 10]`})
+ require.NoError(t, err)
+ defer expected.Release()
+
+ for _, numParallel := range []int{1, 0, -1} {
+ numParallel := numParallel
+ t.Run(fmt.Sprintf("parallelism-%d", numParallel), func(t
*testing.T) {
+ result, err := compute.Take(serialTakeContext(mem,
numParallel), *compute.DefaultTakeOptions(),
+ &compute.ArrayDatum{Value: values.Data()},
&compute.ChunkedDatum{Value: indices})
+ require.NoError(t, err)
+ defer result.Release()
+
+ assert.True(t, array.ChunkedEqual(expected,
result.(*compute.ChunkedDatum).Value))
+ })
+ }
+}
