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 8bb9e075 perf(arrow/compute): avoid goroutines for serial record
filtering (#1245)
8bb9e075 is described below
commit 8bb9e0750ece00a142c16b4a38099e78daa27884
Author: Minh Vu <[email protected]>
AuthorDate: Mon Aug 31 17:27:33 2026 +0200
perf(arrow/compute): avoid goroutines for serial record filtering (#1245)
## Summary
- **Run FilterRecordBatch columns synchronously for serial execution.**
- One-column batches use the direct path even when parallelism is
configured.
- `NumParallel <= 1` uses the direct path for all columns.
- Keep the errgroup path for multi-column parallel execution.
- Add coverage for one-column, zero/one worker, parallel, null
selection, and length errors.
## Benchmark
Local Apple M1 Pro run with `NumParallel=1`:
- 8 columns / 16 rows: **53.3 us -> 26.2 us**, **263 -> 240 allocs/op**
- 128 columns / 16 rows: **709.8 us -> 420.9 us**, **3,983 -> 3,720
allocs/op**
- 128 columns / 4,096 rows: **1.07 ms -> 0.77 ms**, **4,046 -> 3,783
allocs/op**
## Tests
- `go test ./...`
- `GOOS=linux GOARCH=386 go build ./arrow/compute/...`
---
arrow/compute/selection.go | 40 +++++--
...selection_filter_record_batch_benchmark_test.go | 92 +++++++++++++++
.../compute/selection_filter_record_batch_test.go | 129 +++++++++++++++++++++
3 files changed, 249 insertions(+), 12 deletions(-)
diff --git a/arrow/compute/selection.go b/arrow/compute/selection.go
index 38d43728..ea69f575 100644
--- a/arrow/compute/selection.go
+++ b/arrow/compute/selection.go
@@ -676,6 +676,10 @@ func FilterArray(ctx context.Context, values, filter
arrow.Array, options Filter
return outDatum.(*ArrayDatum).MakeArray(), nil
}
+func filterRecordBatchColumn(ctx context.Context, col, indices arrow.Array)
(arrow.Array, error) {
+ return TakeArrayOpts(ctx, col, indices,
kernels.TakeOptions{BoundsCheck: false})
+}
+
func FilterRecordBatch(ctx context.Context, batch arrow.RecordBatch, filter
arrow.Array, opts *FilterOptions) (arrow.RecordBatch, error) {
if batch.NumRows() != int64(filter.Len()) {
return nil, fmt.Errorf("%w: filter inputs must all be the same
length", arrow.ErrInvalid)
@@ -701,22 +705,34 @@ func FilterRecordBatch(ctx context.Context, batch
arrow.RecordBatch, filter arro
}
}
}()
- eg, cctx := errgroup.WithContext(ctx)
- eg.SetLimit(GetExecCtx(ctx).NumParallel)
- for i, col := range batch.Columns() {
- i, col := i, col
- eg.Go(func() error {
- out, err := TakeArrayOpts(cctx, col, indicesArr,
kernels.TakeOptions{BoundsCheck: false})
+
+ numParallel := GetExecCtx(ctx).NumParallel
+ if batch.NumCols() == 1 || numParallel <= 1 {
+ for i, col := range batch.Columns() {
+ out, err := filterRecordBatchColumn(ctx, col,
indicesArr)
if err != nil {
- return err
+ return nil, err
}
cols[i] = out
- return nil
- })
- }
+ }
+ } else {
+ eg, cctx := errgroup.WithContext(ctx)
+ eg.SetLimit(numParallel)
+ for i, col := range batch.Columns() {
+ i, col := i, col
+ eg.Go(func() error {
+ out, err := filterRecordBatchColumn(cctx, col,
indicesArr)
+ if err != nil {
+ return err
+ }
+ cols[i] = out
+ return nil
+ })
+ }
- if err := eg.Wait(); err != nil {
- return nil, err
+ if err := eg.Wait(); err != nil {
+ return nil, err
+ }
}
return array.NewRecordBatch(batch.Schema(), cols,
int64(indicesArr.Len())), nil
diff --git a/arrow/compute/selection_filter_record_batch_benchmark_test.go
b/arrow/compute/selection_filter_record_batch_benchmark_test.go
new file mode 100644
index 00000000..0bb210d6
--- /dev/null
+++ b/arrow/compute/selection_filter_record_batch_benchmark_test.go
@@ -0,0 +1,92 @@
+// 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"
+)
+
+var benchmarkFilterRecordBatchRows int64
+
+func BenchmarkFilterRecordBatchSerial(b *testing.B) {
+ for _, numCols := range []int{1, 8, 32, 128} {
+ for _, numRows := range []int{16, 256, 4096} {
+ b.Run(fmt.Sprintf("columns=%d/rows=%d", numCols,
numRows), func(b *testing.B) {
+ batch, filter :=
makeFilterRecordBatchBenchmarkInput(b, numCols, numRows)
+ defer batch.Release()
+ defer filter.Release()
+
+ execCtx := compute.DefaultExecCtx()
+ execCtx.NumParallel = 1
+ ctx := compute.SetExecCtx(context.Background(),
execCtx)
+
+ b.ReportAllocs()
+ b.SetBytes(int64(numCols * numRows * 8))
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ result, err :=
compute.FilterRecordBatch(ctx, batch, filter, compute.DefaultFilterOptions())
+ if err != nil {
+ b.Fatal(err)
+ }
+ benchmarkFilterRecordBatchRows =
result.NumRows()
+ result.Release()
+ }
+ })
+ }
+ }
+}
+
+func makeFilterRecordBatchBenchmarkInput(b *testing.B, numCols, numRows int)
(arrow.RecordBatch, arrow.Array) {
+ b.Helper()
+ mem := memory.DefaultAllocator
+ fields := make([]arrow.Field, numCols)
+ cols := make([]arrow.Array, numCols)
+ for col := 0; col < numCols; col++ {
+ fields[col] = arrow.Field{Name: fmt.Sprintf("col_%d", col),
Type: arrow.PrimitiveTypes.Int64}
+ builder := array.NewInt64Builder(mem)
+ builder.Reserve(numRows)
+ for row := 0; row < numRows; row++ {
+ builder.Append(int64(col*numRows + row))
+ }
+ cols[col] = builder.NewInt64Array()
+ builder.Release()
+ }
+
+ schema := arrow.NewSchema(fields, nil)
+ batch := array.NewRecordBatch(schema, cols, int64(numRows))
+ for _, col := range cols {
+ col.Release()
+ }
+
+ filterBuilder := array.NewBooleanBuilder(mem)
+ filterBuilder.Reserve(numRows)
+ for row := 0; row < numRows; row++ {
+ filterBuilder.Append(row%2 == 0)
+ }
+ filter := filterBuilder.NewBooleanArray()
+ filterBuilder.Release()
+ return batch, filter
+}
diff --git a/arrow/compute/selection_filter_record_batch_test.go
b/arrow/compute/selection_filter_record_batch_test.go
new file mode 100644
index 00000000..3f22e839
--- /dev/null
+++ b/arrow/compute/selection_filter_record_batch_test.go
@@ -0,0 +1,129 @@
+// 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"
+ "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 TestFilterRecordBatchSerialPaths(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ fields := []arrow.Field{
+ {Name: "a", Type: arrow.PrimitiveTypes.Int32, Nullable: true},
+ {Name: "b", Type: arrow.BinaryTypes.String, Nullable: true},
+ }
+ schema := arrow.NewSchema(fields, nil)
+ batch, _, err := array.RecordFromJSON(mem, schema, strings.NewReader(`[
+ {"a": null, "b": "yo"},
+ {"a": 1, "b": ""},
+ {"a": 2, "b": "hello"},
+ {"a": 4, "b": "eh"}
+ ]`))
+ require.NoError(t, err)
+ defer batch.Release()
+
+ filter, _, err := array.FromJSON(mem, arrow.FixedWidthTypes.Boolean,
strings.NewReader(`[true, null, false, true]`))
+ require.NoError(t, err)
+ defer filter.Release()
+
+ oneColumnSchema := arrow.NewSchema(fields[:1], nil)
+ oneColumnBatch := array.NewRecordBatch(oneColumnSchema,
[]arrow.Array{batch.Column(0)}, batch.NumRows())
+ defer oneColumnBatch.Release()
+
+ tests := []struct {
+ name string
+ batch arrow.RecordBatch
+ numParallel int
+ nullSelection compute.NullSelectionBehavior
+ expected string
+ }{
+ {
+ name: "one column",
+ batch: oneColumnBatch,
+ numParallel: 2,
+ nullSelection: compute.SelectionEmitNulls,
+ expected: `[{"a": null}, {"a": null}, {"a": 4}]`,
+ },
+ {
+ name: "one parallel worker",
+ batch: batch,
+ numParallel: 1,
+ nullSelection: compute.SelectionEmitNulls,
+ expected: `[
+ {"a": null, "b": "yo"},
+ {"a": null, "b": null},
+ {"a": 4, "b": "eh"}
+ ]`,
+ },
+ {
+ name: "zero parallel workers",
+ batch: batch,
+ numParallel: 0,
+ nullSelection: compute.SelectionDropNulls,
+ expected: `[
+ {"a": null, "b": "yo"},
+ {"a": 4, "b": "eh"}
+ ]`,
+ },
+ {
+ name: "parallel workers",
+ batch: batch,
+ numParallel: 2,
+ nullSelection: compute.SelectionDropNulls,
+ expected: `[
+ {"a": null, "b": "yo"},
+ {"a": 4, "b": "eh"}
+ ]`,
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ execCtx := compute.DefaultExecCtx()
+ execCtx.NumParallel = tt.numParallel
+ ctx := compute.SetExecCtx(context.Background(), execCtx)
+
+ actual, err := compute.FilterRecordBatch(ctx, tt.batch,
filter, &compute.FilterOptions{NullSelection: tt.nullSelection})
+ require.NoError(t, err)
+ defer actual.Release()
+
+ expected, _, err := array.RecordFromJSON(mem,
tt.batch.Schema(), strings.NewReader(tt.expected))
+ require.NoError(t, err)
+ defer expected.Release()
+ assert.Truef(t, array.RecordEqual(expected, actual),
"expected: %s\ngot: %s", expected, actual)
+ })
+ }
+
+ shortFilter, _, err := array.FromJSON(mem,
arrow.FixedWidthTypes.Boolean, strings.NewReader(`[true]`))
+ require.NoError(t, err)
+ defer shortFilter.Release()
+ _, err = compute.FilterRecordBatch(context.Background(), batch,
shortFilter, compute.DefaultFilterOptions())
+ require.ErrorIs(t, err, arrow.ErrInvalid)
+}
