diff --git a/arrow/compute/selection.go b/arrow/compute/selection.go index 38d43728c..ea69f5754 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 000000000..0bb210d68 --- /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 000000000..3f22e8398 --- /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) +}