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 4a8148fb perf(compute): batch contiguous filter take indices (#1197)
4a8148fb is described below
commit 4a8148fb7d00b8cdc9ec123ea99a8443ce081e9e
Author: Minh Vu <[email protected]>
AuthorDate: Wed Aug 26 18:43:34 2026 +0200
perf(compute): batch contiguous filter take indices (#1197)
## Summary
* Add a range writer for uint16 and uint32 take indices.
* Use it when a filter block or bitmap run is fully selected.
* Keep the scalar path for mixed blocks.
* Add benchmarks for dense, clustered, alternating, short, and nullable
filters.
## Benchmark
On an Apple M1 Pro, the 1M-row all-selected case improved from about 3.0
ms to 0.5 ms locally. Fragmented filters stayed close to the baseline.
## Tests
* `go test ./arrow/compute/internal/kernels ./arrow/compute`
* `go test -race ./arrow/compute/internal/kernels ./arrow/compute`
* Full package test run passed with the local Parquet test data
checkout, excluding the CSV example that needs a separate Arrow CSV
fixture.
---
arrow/compute/internal/kernels/vector_selection.go | 27 ++-
.../kernels/vector_selection_bench_test.go | 80 ++++++++
.../internal/kernels/vector_selection_test.go | 203 +++++++++++++++++++++
3 files changed, 301 insertions(+), 9 deletions(-)
diff --git a/arrow/compute/internal/kernels/vector_selection.go
b/arrow/compute/internal/kernels/vector_selection.go
index 00bc141a..cbe0c220 100644
--- a/arrow/compute/internal/kernels/vector_selection.go
+++ b/arrow/compute/internal/kernels/vector_selection.go
@@ -99,6 +99,19 @@ type builder[T any] interface {
UnsafeAppendBoolToBitmap(bool)
}
+func unsafeAppendRange[T arrow.IntType | arrow.UintType](b *bufferBuilder[T],
start T, n int) {
+ if n == 0 {
+ return
+ }
+
+ b.reserve(n)
+ values := arrow.GetData[T](b.data[b.sz:])[:n]
+ for i := range values {
+ values[i] = start + T(i)
+ }
+ b.sz += len(arrow.GetBytes(values))
+}
+
func getTakeIndices[T arrow.IntType | arrow.UintType](mem memory.Allocator,
filter *exec.ArraySpan, nullSelect NullSelectionBehavior) arrow.ArrayData {
var (
filterData = filter.Buffers[1].Buf
@@ -129,6 +142,7 @@ func getTakeIndices[T arrow.IntType | arrow.UintType](mem
memory.Allocator, filt
// true OR NOT valid
selectedOrNullBlock := filterCounter.NextOrNotWord()
if selectedOrNullBlock.NoneSet() {
+ isValidCounter.NextWord()
pos += T(selectedOrNullBlock.Len)
posWithOffset += int64(selectedOrNullBlock.Len)
continue
@@ -183,16 +197,14 @@ func getTakeIndices[T arrow.IntType | arrow.UintType](mem
memory.Allocator, filt
filterCounter := bitutils.NewBinaryBitBlockCounter(filterData,
filterIsValid, filter.Offset, filter.Offset, filter.Len)
for int64(pos) < filter.Len {
andBlock := filterCounter.NextAndWord()
- bldr.reserve(int(andBlock.Popcnt))
if andBlock.AllSet() {
// all the values are selected and non-null
- for i := 0; i < int(andBlock.Len); i++ {
- bldr.unsafeAppend(pos)
- pos++
- }
+ unsafeAppendRange(bldr, pos, int(andBlock.Len))
+ pos += T(andBlock.Len)
posWithOffset += int64(andBlock.Len)
} else if !andBlock.NoneSet() {
// some values are false or null
+ bldr.reserve(int(andBlock.Popcnt))
for i := 0; i < int(andBlock.Len); i++ {
if bitutil.BitIsSet(filterIsValid,
int(posWithOffset)) && bitutil.BitIsSet(filterData, int(posWithOffset)) {
bldr.unsafeAppend(pos)
@@ -210,10 +222,7 @@ func getTakeIndices[T arrow.IntType | arrow.UintType](mem
memory.Allocator, filt
bitutils.VisitSetBitRuns(filterData, filter.Offset, filter.Len,
func(pos, length int64) error {
// append consecutive run of indices
- bldr.reserve(int(length))
- for i := int64(0); i < length; i++ {
- bldr.unsafeAppend(T(pos + i))
- }
+ unsafeAppendRange(bldr, T(pos), int(length))
return nil
})
}
diff --git a/arrow/compute/internal/kernels/vector_selection_bench_test.go
b/arrow/compute/internal/kernels/vector_selection_bench_test.go
new file mode 100644
index 00000000..e904e881
--- /dev/null
+++ b/arrow/compute/internal/kernels/vector_selection_bench_test.go
@@ -0,0 +1,80 @@
+// 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 kernels
+
+import (
+ "fmt"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/compute/exec"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+)
+
+type filterBenchmarkPattern struct {
+ name string
+ selected func(int) bool
+ nullable bool
+}
+
+var filterBenchmarkPatterns = []filterBenchmarkPattern{
+ {name: "all", selected: func(int) bool { return true }},
+ {name: "long-runs", selected: func(i int) bool { return i%1024 < 900 }},
+ {name: "alternating", selected: func(i int) bool { return i%2 == 0 }},
+ {name: "short-runs", selected: func(i int) bool { return i%100 < 10 }},
+ {name: "nullable-long-runs", selected: func(i int) bool { return i%1024
< 900 }, nullable: true},
+}
+
+func makeFilterBenchmarkSpan(tb testing.TB, n int, pattern
filterBenchmarkPattern) *exec.ArraySpan {
+ tb.Helper()
+ bldr := array.NewBooleanBuilder(memory.DefaultAllocator)
+ for i := 0; i < n; i++ {
+ if pattern.nullable && i%257 == 0 {
+ bldr.AppendNull()
+ } else {
+ bldr.Append(pattern.selected(i))
+ }
+ }
+ filter := bldr.NewArray()
+ bldr.Release()
+ tb.Cleanup(filter.Release)
+
+ span := &exec.ArraySpan{}
+ span.SetMembers(filter.Data())
+ return span
+}
+
+func BenchmarkGetTakeIndices(b *testing.B) {
+ for _, n := range []int{64 * 1024, 1024 * 1024} {
+ for _, pattern := range filterBenchmarkPatterns {
+ b.Run(fmt.Sprintf("%s/%d", pattern.name, n), func(b
*testing.B) {
+ filter := makeFilterBenchmarkSpan(b, n, pattern)
+ b.ReportAllocs()
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ result, err :=
GetTakeIndices(memory.DefaultAllocator, filter, DropNulls)
+ if err != nil {
+ b.Fatal(err)
+ }
+ result.Release()
+ }
+ })
+ }
+ }
+}
diff --git a/arrow/compute/internal/kernels/vector_selection_test.go
b/arrow/compute/internal/kernels/vector_selection_test.go
new file mode 100644
index 00000000..46bf8143
--- /dev/null
+++ b/arrow/compute/internal/kernels/vector_selection_test.go
@@ -0,0 +1,203 @@
+// 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 kernels
+
+import (
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/bitutil"
+ "github.com/apache/arrow-go/v18/arrow/compute/exec"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func makeBooleanFilter(t *testing.T, values []bool, valid []bool, mem
memory.Allocator) arrow.Array {
+ t.Helper()
+ if valid != nil {
+ require.Len(t, valid, len(values))
+ }
+
+ bldr := array.NewBooleanBuilder(mem)
+ defer bldr.Release()
+ for i, value := range values {
+ if valid != nil && !valid[i] {
+ bldr.AppendNull()
+ } else {
+ bldr.Append(value)
+ }
+ }
+ return bldr.NewArray()
+}
+
+func makeSlicedBooleanFilter(t *testing.T, values []bool, valid []bool, prefix
int, mem memory.Allocator) arrow.Array {
+ t.Helper()
+ allValues := make([]bool, len(values)+2*prefix)
+ for i := range allValues {
+ allValues[i] = true
+ }
+ copy(allValues[prefix:], values)
+
+ var allValid []bool
+ if valid != nil {
+ allValid = make([]bool, len(allValues))
+ for i := range allValid {
+ allValid[i] = true
+ }
+ copy(allValid[prefix:], valid)
+ }
+
+ base := makeBooleanFilter(t, allValues, allValid, mem)
+ sliced := array.NewSlice(base, int64(prefix), int64(prefix+len(values)))
+ base.Release()
+ return sliced
+}
+
+func makeIndexValues[T arrow.IntType | arrow.UintType](selected []int) []T {
+ values := make([]T, len(selected))
+ for i, value := range selected {
+ values[i] = T(value)
+ }
+ return values
+}
+
+func assertTakeIndices[T arrow.IntType | arrow.UintType](t *testing.T, data
arrow.ArrayData, wantValues []T, wantValid []bool) {
+ t.Helper()
+ require.Equal(t, arrow.GetDataType[T]().ID(), data.DataType().ID())
+ require.Equal(t, len(wantValues), data.Len())
+ require.NotNil(t, data.Buffers()[1])
+
+ values := arrow.GetData[T](data.Buffers()[1].Bytes())
+ values = values[data.Offset() : data.Offset()+data.Len()]
+ for i, want := range wantValues {
+ if wantValid == nil || wantValid[i] {
+ assert.Equal(t, want, values[i], "value at index %d", i)
+ }
+ }
+
+ if wantValid == nil {
+ require.Nil(t, data.Buffers()[0])
+ require.Zero(t, data.NullN())
+ return
+ }
+
+ require.NotNil(t, data.Buffers()[0])
+ nulls := 0
+ for i, want := range wantValid {
+ if !want {
+ nulls++
+ }
+ got := bitutil.BitIsSet(data.Buffers()[0].Bytes(),
data.Offset()+i)
+ assert.Equal(t, want, got, "validity at index %d", i)
+ }
+ require.Equal(t, nulls, data.NullN())
+}
+
+func getTakeIndicesForTest[T arrow.IntType | arrow.UintType](mem
memory.Allocator, filter arrow.Array, nullSelect NullSelectionBehavior)
arrow.ArrayData {
+ var span exec.ArraySpan
+ span.SetMembers(filter.Data())
+ return getTakeIndices[T](mem, &span, nullSelect)
+}
+
+func TestGetTakeIndicesBatchedRanges(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ const (
+ length = 192
+ prefix = 3
+ )
+
+ values := make([]bool, length)
+ for i := range values {
+ values[i] = i < 64 || i >= 128
+ }
+
+ selected := make([]int, 0, 128)
+ for i, value := range values {
+ if value {
+ selected = append(selected, i)
+ }
+ }
+
+ nullableValid := make([]bool, length)
+ for i := range nullableValid {
+ nullableValid[i] = true
+ }
+ nullableValid[140] = false
+
+ dropSelected := make([]int, 0, len(selected)-1)
+ for _, value := range selected {
+ if value != 140 {
+ dropSelected = append(dropSelected, value)
+ }
+ }
+ emitValid := make([]bool, len(selected))
+ for i, value := range selected {
+ emitValid[i] = value != 140
+ }
+
+ tests := []struct {
+ name string
+ valid []bool
+ nullSelect NullSelectionBehavior
+ selected []int
+ validOut []bool
+ }{
+ {
+ name: "non_nullable_runs",
+ nullSelect: DropNulls,
+ selected: selected,
+ },
+ {
+ name: "nullable_drop_nulls",
+ valid: nullableValid,
+ nullSelect: DropNulls,
+ selected: dropSelected,
+ },
+ {
+ name: "nullable_emit_nulls",
+ valid: nullableValid,
+ nullSelect: EmitNulls,
+ selected: selected,
+ validOut: emitValid,
+ },
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ filter := makeSlicedBooleanFilter(t, values, tc.valid,
prefix, mem)
+ defer filter.Release()
+
+ t.Run("uint16", func(t *testing.T) {
+ result := getTakeIndicesForTest[uint16](mem,
filter, tc.nullSelect)
+ defer result.Release()
+ assertTakeIndices(t, result,
makeIndexValues[uint16](tc.selected), tc.validOut)
+ })
+ t.Run("uint32", func(t *testing.T) {
+ result := getTakeIndicesForTest[uint32](mem,
filter, tc.nullSelect)
+ defer result.Release()
+ assertTakeIndices(t, result,
makeIndexValues[uint32](tc.selected), tc.validOut)
+ })
+ })
+ }
+}
