perf(arrow/compute): avoid goroutines for serial record filtering

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

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
+}

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)
+}