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 4caad589 perf(compute): SIMD-compact fragmented filters into take
indices (#1287)
4caad589 is described below
commit 4caad589dde22a0313c6792013671e5d782c7532
Author: Minh Vu <[email protected]>
AuthorDate: Fri Sep 4 18:28:17 2026 +0200
perf(compute): SIMD-compact fragmented filters into take indices (#1287)
### Rationale for this change
`GetTakeIndices` currently grows a buffer once per selected run.
Fragmented filters such as alternating and short random runs create many
allocations and are much slower than dense filters.
### What changes are included in this PR?
- Add an ARM64 NEON path for large, byte-aligned, non-null `uint32`
filters.
- Compact 4-bit mask lookups into exact-size take-index output.
- Keep `uint16`, nullable, unaligned, small, dense, long-run, and
`EmitNulls` paths unchanged.
- Add coverage for all 4-bit masks, offsets, tails, padding bits, null
behavior, and width boundaries.
- Add GetTakeIndices and large FilterRecordBatch benchmarks.
**Benchmark hardware:** Apple M1 Pro, macOS arm64
**Commands:**
```text
go test ./arrow/compute/internal/kernels -run '^$'
-bench='^BenchmarkGetTakeIndices/(all-selected|all-clear|long-runs|alternating|short-runs|random)/'
-benchmem -benchtime=3x -count=3
go test ./arrow/compute -run '^$'
-bench='^BenchmarkFilterRecordBatchGetTakeIndices/(rows=65536|rows=1048576)$'
-benchmem -benchtime=1x -count=3
```
Median of the three reported benchmark samples. Values are ns/op per
call.
| Workload | Selection pattern | Existing | NEON path | Speedup |
| --- | --- | ---: | ---: | ---: |
| GetTakeIndices 64K | all-selected | 41.8 us | 41.2 us | 1.0x |
| GetTakeIndices 64K | all-clear | 2.54 us | 3.03 us | 0.84x |
| GetTakeIndices 64K | long-runs | 826 us | 923 us | 0.89x |
| GetTakeIndices 64K | alternating | 48.1 ms | 21.9 us | 2,198x |
| GetTakeIndices 64K | short-runs | 4.87 ms | 22.6 us | 215x |
| GetTakeIndices 64K | random | 73.5 ms | 90.0 us | 817x |
| GetTakeIndices 1M | all-selected | 1.80 ms | 0.80 ms | 2.3x* |
| GetTakeIndices 1M | all-clear | 37.8 us | 39.0 us | 0.97x |
| GetTakeIndices 1M | long-runs | 705 ms | 176 ms | 4.0x* |
| GetTakeIndices 1M | alternating | 4.31 s | 515 us | 8,369x |
| GetTakeIndices 1M | short-runs | 201 ms | 230 us | 873x |
| GetTakeIndices 1M | random | 3.56 s | 1.69 ms | 2,111x |
`*` These patterns use the existing fallback. The measured difference is
run-to-run allocator and GC noise, not a new fast path.
Downstream serial `FilterRecordBatch` with one Int64 column and an
alternating filter measured **17.7 ms to 152 us (116x)** at 64K rows and
**2.06 s to 831 us (2,477x)** at 1M rows. This benchmark includes both
take-index generation and the downstream Take call.
### Are these changes tested?
- `go test ./...` with the Arrow and Parquet test-data submodules
initialized and `PARQUET_TEST_DATA` / `ARROW_TEST_DATA` set.
- `go test ./arrow/compute/...`
- `go test -tags noasm ./arrow/compute/internal/kernels -run
'TestGetTakeIndices(BatchedRanges|Uint32Coverage)$' -count=1`
- `GOOS=linux GOARCH=amd64 go test -c ./arrow/compute/internal/kernels
-o /dev/null`
### Are there any user-facing changes?
No API changes. The optimization is ARM64-only and only applies to the
narrow `uint32`, non-null, `DropNulls` case.
---
.../kernels/get_take_indices_neon_arm64.go | 107 ++++++++++++++
.../internal/kernels/get_take_indices_neon_arm64.s | 107 ++++++++++++++
.../kernels/get_take_indices_neon_noasm.go | 30 ++++
arrow/compute/internal/kernels/vector_selection.go | 55 ++++----
.../kernels/vector_selection_bench_test.go | 11 +-
.../internal/kernels/vector_selection_test.go | 154 +++++++++++++++++++++
...selection_filter_record_batch_benchmark_test.go | 28 ++++
7 files changed, 466 insertions(+), 26 deletions(-)
diff --git a/arrow/compute/internal/kernels/get_take_indices_neon_arm64.go
b/arrow/compute/internal/kernels/get_take_indices_neon_arm64.go
new file mode 100644
index 00000000..17f6f7ac
--- /dev/null
+++ b/arrow/compute/internal/kernels/get_take_indices_neon_arm64.go
@@ -0,0 +1,107 @@
+// 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 && arm64 && !noasm && !appengine
+
+package kernels
+
+import (
+ "unsafe"
+
+ "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"
+ "golang.org/x/sys/cpu"
+)
+
+var takeIndicesUint32NeonPositions = makeTakeIndicesUint32NeonPositions()
+
+var takeIndicesUint32NeonCounts = [16]uint8{
+ 0, 1, 1, 2, 1, 2, 2, 3,
+ 1, 2, 2, 3, 2, 3, 3, 4,
+}
+
+func makeTakeIndicesUint32NeonPositions() (positions [16][4]uint32) {
+ for mask := 0; mask < len(positions); mask++ {
+ n := 0
+ for bit := 0; bit < 4; bit++ {
+ if mask&(1<<uint(bit)) != 0 {
+ positions[mask][n] = uint32(bit)
+ n++
+ }
+ }
+ }
+ return
+}
+
+//go:noescape
+func _getTakeIndicesUint32NEON(filter, output, positions, counts
unsafe.Pointer, nbytes, tailMask int64)
+
+func getTakeIndicesUint32NEON(mem memory.Allocator, filter *exec.ArraySpan)
(arrow.ArrayData, bool) {
+ if !cpu.ARM64.HasASIMD || filter.MayHaveNulls() || filter.Offset%8 != 0
|| filter.Len < 64 {
+ return nil, false
+ }
+
+ filterData := filter.Buffers[1].Buf
+ byteOffset := filter.Offset / 8
+ nbytes := (filter.Len + 7) / 8
+ if byteOffset < 0 || byteOffset+nbytes > int64(len(filterData)) {
+ return nil, false
+ }
+
+ // VisitSetBitRuns is especially effective for long runs, so only use
the
+ // compactor when a short sample shows enough fragmented bytes to
amortize
+ // its setup cost.
+ const (
+ sampleBytes = 256
+ minMixed = 4
+ )
+ mixed := 0
+ for i := int64(0); i < nbytes && i < sampleBytes; i++ {
+ mask := filterData[byteOffset+i]
+ if mask != 0 && mask != 0xff {
+ mixed++
+ }
+ }
+ if mixed < minMixed {
+ return nil, false
+ }
+
+ length := int64(bitutil.CountSetBits(filterData, int(filter.Offset),
int(filter.Len)))
+ if length == 0 {
+ return array.NewData(arrow.PrimitiveTypes.Uint32, 0,
[]*memory.Buffer{nil, memory.NewBufferBytes(nil)}, nil, 0, 0), true
+ }
+
+ outputBuf := memory.NewBufferWithAllocator(mem.Allocate(int(length*4)),
mem)
+ defer outputBuf.Release()
+ output := arrow.GetData[uint32](outputBuf.Bytes())
+ tailMask := int64(0xff)
+ if tailBits := filter.Len & 7; tailBits != 0 {
+ tailMask = int64((uint64(1) << uint(tailBits)) - 1)
+ }
+ _getTakeIndicesUint32NEON(
+ unsafe.Pointer(unsafe.SliceData(filterData[byteOffset:])),
+ unsafe.Pointer(unsafe.SliceData(output)),
+
unsafe.Pointer(unsafe.SliceData(takeIndicesUint32NeonPositions[:])),
+
unsafe.Pointer(unsafe.SliceData(takeIndicesUint32NeonCounts[:])),
+ nbytes,
+ tailMask,
+ )
+ return array.NewData(arrow.PrimitiveTypes.Uint32, int(length),
[]*memory.Buffer{nil, outputBuf}, nil, 0, 0), true
+}
diff --git a/arrow/compute/internal/kernels/get_take_indices_neon_arm64.s
b/arrow/compute/internal/kernels/get_take_indices_neon_arm64.s
new file mode 100644
index 00000000..fe9f5cee
--- /dev/null
+++ b/arrow/compute/internal/kernels/get_take_indices_neon_arm64.s
@@ -0,0 +1,107 @@
+// 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 && arm64 && !noasm && !appengine
+
+#include "textflag.h"
+
+// func _getTakeIndicesUint32NEON(filter, output, positions, counts
unsafe.Pointer, nbytes, tailMask int64)
+TEXT ยท_getTakeIndicesUint32NEON(SB), NOSPLIT|NOFRAME, $0-48
+ MOVD filter+0(FP), R0
+ MOVD output+8(FP), R1
+ MOVD positions+16(FP), R2
+ MOVD counts+24(FP), R3
+ MOVD nbytes+32(FP), R4
+ MOVD tailMask+40(FP), R5
+ MOVD $0, R6 // input bit position
+
+byte_loop:
+ CBZ R4, done
+ MOVBU (R0), R7
+ SUB $1, R4, R8
+ CBNZ R8, byte_not_tail
+ AND R5, R7, R7
+
+byte_not_tail:
+ AND $15, R7, R8
+ MOVBU (R3)(R8), R9
+ CBZ R9, high_nibble
+ LSL $4, R8, R10
+ ADD R2, R10, R10
+ VLD1 (R10), [V0.S4]
+ VDUP R6, V1.S4
+ VADD V1.S4, V0.S4, V0.S4
+ CMP $4, R9
+ BNE low_scalar
+ VST1 [V0.S4], (R1)
+ ADD $16, R1, R1
+ JMP high_nibble
+
+low_scalar:
+ VMOV V0.S[0], R10
+ MOVW R10, (R1)
+ ADD $4, R1, R1
+ CMP $1, R9
+ BEQ high_nibble
+ VMOV V0.S[1], R10
+ MOVW R10, (R1)
+ ADD $4, R1, R1
+ CMP $2, R9
+ BEQ high_nibble
+ VMOV V0.S[2], R10
+ MOVW R10, (R1)
+ ADD $4, R1, R1
+
+high_nibble:
+ LSR $4, R7, R8
+ MOVBU (R3)(R8), R9
+ CBZ R9, next_byte
+ LSL $4, R8, R10
+ ADD R2, R10, R10
+ VLD1 (R10), [V0.S4]
+ ADD $4, R6, R10
+ VDUP R10, V1.S4
+ VADD V1.S4, V0.S4, V0.S4
+ CMP $4, R9
+ BNE high_scalar
+ VST1 [V0.S4], (R1)
+ ADD $16, R1, R1
+ JMP next_byte
+
+high_scalar:
+ VMOV V0.S[0], R10
+ MOVW R10, (R1)
+ ADD $4, R1, R1
+ CMP $1, R9
+ BEQ next_byte
+ VMOV V0.S[1], R10
+ MOVW R10, (R1)
+ ADD $4, R1, R1
+ CMP $2, R9
+ BEQ next_byte
+ VMOV V0.S[2], R10
+ MOVW R10, (R1)
+ ADD $4, R1, R1
+
+next_byte:
+ ADD $8, R6, R6
+ ADD $1, R0, R0
+ SUB $1, R4, R4
+ JMP byte_loop
+
+done:
+ RET
diff --git a/arrow/compute/internal/kernels/get_take_indices_neon_noasm.go
b/arrow/compute/internal/kernels/get_take_indices_neon_noasm.go
new file mode 100644
index 00000000..89b836c5
--- /dev/null
+++ b/arrow/compute/internal/kernels/get_take_indices_neon_noasm.go
@@ -0,0 +1,30 @@
+// 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 && (noasm || !arm64 || appengine)
+
+package kernels
+
+import (
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/compute/exec"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+)
+
+func getTakeIndicesUint32NEON(memory.Allocator, *exec.ArraySpan)
(arrow.ArrayData, bool) {
+ return nil, false
+}
diff --git a/arrow/compute/internal/kernels/vector_selection.go
b/arrow/compute/internal/kernels/vector_selection.go
index 68198628..09454116 100644
--- a/arrow/compute/internal/kernels/vector_selection.go
+++ b/arrow/compute/internal/kernels/vector_selection.go
@@ -21,6 +21,7 @@ package kernels
import (
"fmt"
"math"
+ "unsafe"
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/array"
@@ -99,19 +100,6 @@ 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
@@ -183,11 +171,15 @@ func getTakeIndices[T arrow.IntType | arrow.UintType](mem
memory.Allocator, filt
return result.Data()
}
- bldr := newBufferBuilder[T](mem)
if haveFilterNulls {
// the filter may have nulls, so we scan the validity bitmap
// and the filter data bitmap together
debug.Assert(nullSelect == DropNulls, "incorrect nullselect
logic")
+ length := getFilterOutputSize(filter, DropNulls)
+ outBuf :=
memory.NewBufferWithAllocator(mem.Allocate(int(length)*int(unsafe.Sizeof(*new(T)))),
mem)
+ defer outBuf.Release()
+ out := arrow.GetData[T](outBuf.Bytes())
+ outPos := 0
// position relative to start of the filter
var pos T
@@ -199,15 +191,18 @@ func getTakeIndices[T arrow.IntType | arrow.UintType](mem
memory.Allocator, filt
andBlock := filterCounter.NextAndWord()
if andBlock.AllSet() {
// all the values are selected and non-null
- unsafeAppendRange(bldr, pos, int(andBlock.Len))
- pos += T(andBlock.Len)
+ for i := 0; i < int(andBlock.Len); i++ {
+ out[outPos] = pos
+ outPos++
+ pos++
+ }
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)
+ out[outPos] = pos
+ outPos++
}
pos++
posWithOffset++
@@ -217,20 +212,25 @@ func getTakeIndices[T arrow.IntType | arrow.UintType](mem
memory.Allocator, filt
posWithOffset += int64(andBlock.Len)
}
}
+ return array.NewData(idxType, int(length),
[]*memory.Buffer{nil, outBuf}, nil, 0, 0)
} else {
// filter has no nulls, so we only need to look for true values
+ length := int64(bitutil.CountSetBits(filterData,
int(filter.Offset), int(filter.Len)))
+ outBuf :=
memory.NewBufferWithAllocator(mem.Allocate(int(length)*int(unsafe.Sizeof(*new(T)))),
mem)
+ defer outBuf.Release()
+ out := arrow.GetData[T](outBuf.Bytes())
+ outPos := 0
bitutils.VisitSetBitRuns(filterData, filter.Offset, filter.Len,
- func(pos, length int64) error {
+ func(pos, runLength int64) error {
// append consecutive run of indices
- unsafeAppendRange(bldr, T(pos), int(length))
+ for i := int64(0); i < runLength; i++ {
+ out[outPos] = T(pos + i)
+ outPos++
+ }
return nil
})
+ return array.NewData(idxType, int(length),
[]*memory.Buffer{nil, outBuf}, nil, 0, 0)
}
-
- length := bldr.len()
- outBuf := bldr.finish()
- defer outBuf.Release()
- return array.NewData(idxType, length, []*memory.Buffer{nil, outBuf},
nil, 0, 0)
}
func GetTakeIndices(mem memory.Allocator, filter *exec.ArraySpan, nullSelect
NullSelectionBehavior) (arrow.ArrayData, error) {
@@ -238,6 +238,11 @@ func GetTakeIndices(mem memory.Allocator, filter
*exec.ArraySpan, nullSelect Nul
if filter.Len < math.MaxUint16 {
return getTakeIndices[uint16](mem, filter, nullSelect), nil
} else if filter.Len < math.MaxUint32 {
+ if nullSelect == DropNulls {
+ if result, ok := getTakeIndicesUint32NEON(mem, filter);
ok {
+ return result, nil
+ }
+ }
return getTakeIndices[uint32](mem, filter, nullSelect), nil
}
return nil, fmt.Errorf("%w: filter length exceeds UINT32_MAX, consider
a different strategy for selecting elements",
diff --git a/arrow/compute/internal/kernels/vector_selection_bench_test.go
b/arrow/compute/internal/kernels/vector_selection_bench_test.go
index e904e881..e9000ff5 100644
--- a/arrow/compute/internal/kernels/vector_selection_bench_test.go
+++ b/arrow/compute/internal/kernels/vector_selection_bench_test.go
@@ -34,13 +34,22 @@ type filterBenchmarkPattern struct {
}
var filterBenchmarkPatterns = []filterBenchmarkPattern{
- {name: "all", selected: func(int) bool { return true }},
+ {name: "all-selected", selected: func(int) bool { return true }},
+ {name: "all-clear", selected: func(int) bool { return false }},
{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: "random", selected: func(i int) bool { return
getTakeIndicesBenchmarkRandom(i) }},
{name: "nullable-long-runs", selected: func(i int) bool { return i%1024
< 900 }, nullable: true},
}
+func getTakeIndicesBenchmarkRandom(i int) bool {
+ x := uint32(i)*747796405 + 2891336453
+ x = ((x >> ((x >> 28) + 4)) ^ x) * 277803737
+ x = (x >> 22) ^ x
+ return x&1 == 0
+}
+
func makeFilterBenchmarkSpan(tb testing.TB, n int, pattern
filterBenchmarkPattern) *exec.ArraySpan {
tb.Helper()
bldr := array.NewBooleanBuilder(memory.DefaultAllocator)
diff --git a/arrow/compute/internal/kernels/vector_selection_test.go
b/arrow/compute/internal/kernels/vector_selection_test.go
index 46bf8143..f4c69338 100644
--- a/arrow/compute/internal/kernels/vector_selection_test.go
+++ b/arrow/compute/internal/kernels/vector_selection_test.go
@@ -201,3 +201,157 @@ func TestGetTakeIndicesBatchedRanges(t *testing.T) {
})
}
}
+
+func TestGetTakeIndicesUint32Coverage(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ const length = 65536
+ values := make([]bool, length)
+ want := make([]uint32, 0, length/2)
+ for i := range values {
+ mask := byte(i / 8)
+ values[i] = mask&(1<<uint(i%8)) != 0
+ if values[i] {
+ want = append(want, uint32(i))
+ }
+ }
+
+ tests := []struct {
+ name string
+ filter arrow.Array
+ want []uint32
+ }{
+ {
+ name: "all_masks",
+ filter: makeBooleanFilter(t, values, nil, mem),
+ want: want,
+ },
+ {
+ name: "aligned_offset",
+ filter: makeSlicedBooleanFilter(t, values, nil, 8, mem),
+ want: want,
+ },
+ {
+ name: "unaligned_offset",
+ filter: makeSlicedBooleanFilter(t, values, nil, 1, mem),
+ want: want,
+ },
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ defer tc.filter.Release()
+ var span exec.ArraySpan
+ span.SetMembers(tc.filter.Data())
+
+ result, err := GetTakeIndices(mem, &span, DropNulls)
+ require.NoError(t, err)
+ defer result.Release()
+ require.Equal(t, arrow.PrimitiveTypes.Uint32.ID(),
result.DataType().ID())
+ got :=
arrow.GetData[uint32](result.Buffers()[1].Bytes())
+ require.Len(t, got, len(tc.want))
+ for i := range got {
+ if got[i] != tc.want[i] {
+ t.Fatalf("value at index %d: got %d,
want %d", i, got[i], tc.want[i])
+ }
+ }
+ })
+ }
+
+ t.Run("nullable_filter_uses_scalar_path", func(t *testing.T) {
+ valid := make([]bool, length)
+ for i := range valid {
+ valid[i] = true
+ }
+ valid[8] = false
+ filter := makeBooleanFilter(t, values, valid, mem)
+ defer filter.Release()
+ var span exec.ArraySpan
+ span.SetMembers(filter.Data())
+
+ wantDrop := make([]uint32, 0, len(want)-1)
+ for _, value := range want {
+ if value != 8 {
+ wantDrop = append(wantDrop, value)
+ }
+ }
+ result, err := GetTakeIndices(mem, &span, DropNulls)
+ require.NoError(t, err)
+ assertTakeIndices(t, result, wantDrop, nil)
+ result.Release()
+
+ result, err = GetTakeIndices(mem, &span, EmitNulls)
+ require.NoError(t, err)
+ defer result.Release()
+ wantValid := make([]bool, len(want))
+ for i := range wantValid {
+ wantValid[i] = true
+ }
+ for i, value := range want {
+ if value == 8 {
+ wantValid[i] = false
+ }
+ }
+ assertTakeIndices(t, result, want, wantValid)
+ })
+
+ t.Run("tail_ignores_padding_bits", func(t *testing.T) {
+ const tailLength = int64(length + 3)
+ data := make([]byte, int(bitutil.BytesForBits(tailLength)))
+ wantTail := make([]uint32, 0, len(want))
+ for i := int64(0); i < tailLength; i++ {
+ mask := byte(i / 8)
+ if mask&(1<<uint(i%8)) != 0 {
+ bitutil.SetBit(data, int(i))
+ wantTail = append(wantTail, uint32(i))
+ }
+ }
+ data[len(data)-1] |= 0xe0
+
+ filterData := array.NewData(arrow.FixedWidthTypes.Boolean,
int(tailLength), []*memory.Buffer{
+ nil,
+ memory.NewBufferBytes(data),
+ }, nil, 0, 0)
+ defer filterData.Release()
+ var span exec.ArraySpan
+ span.SetMembers(filterData)
+
+ result, err := GetTakeIndices(mem, &span, DropNulls)
+ require.NoError(t, err)
+ defer result.Release()
+ got := arrow.GetData[uint32](result.Buffers()[1].Bytes())
+ require.Len(t, got, len(wantTail))
+ for i := range got {
+ if got[i] != wantTail[i] {
+ t.Fatalf("value at index %d: got %d, want %d",
i, got[i], wantTail[i])
+ }
+ }
+ })
+
+ for _, tc := range []struct {
+ name string
+ length int
+ wantID arrow.Type
+ }{
+ {name: "uint16_boundary", length: 65534, wantID: arrow.UINT16},
+ {name: "uint32_boundary", length: 65535, wantID: arrow.UINT32},
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ values := make([]bool, tc.length)
+ for i := range values {
+ values[i] = true
+ }
+ filter := makeBooleanFilter(t, values, nil, mem)
+ defer filter.Release()
+ var span exec.ArraySpan
+ span.SetMembers(filter.Data())
+
+ result, err := GetTakeIndices(mem, &span, DropNulls)
+ require.NoError(t, err)
+ defer result.Release()
+ require.Equal(t, tc.wantID, result.DataType().ID())
+ require.Equal(t, tc.length, result.Len())
+ })
+ }
+}
diff --git a/arrow/compute/selection_filter_record_batch_benchmark_test.go
b/arrow/compute/selection_filter_record_batch_benchmark_test.go
index 0bb210d6..af468303 100644
--- a/arrow/compute/selection_filter_record_batch_benchmark_test.go
+++ b/arrow/compute/selection_filter_record_batch_benchmark_test.go
@@ -59,6 +59,34 @@ func BenchmarkFilterRecordBatchSerial(b *testing.B) {
}
}
+var benchmarkLargeFilterRecordBatchRows int64
+
+func BenchmarkFilterRecordBatchGetTakeIndices(b *testing.B) {
+ for _, numRows := range []int{64 * 1024, 1024 * 1024} {
+ b.Run(fmt.Sprintf("rows=%d", numRows), func(b *testing.B) {
+ batch, filter := makeFilterRecordBatchBenchmarkInput(b,
1, numRows)
+ defer batch.Release()
+ defer filter.Release()
+
+ execCtx := compute.DefaultExecCtx()
+ execCtx.NumParallel = 1
+ ctx := compute.SetExecCtx(context.Background(), execCtx)
+
+ b.ReportAllocs()
+ b.SetBytes(int64(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)
+ }
+ benchmarkLargeFilterRecordBatchRows =
result.NumRows()
+ result.Release()
+ }
+ })
+ }
+}
+
func makeFilterRecordBatchBenchmarkInput(b *testing.B, numCols, numRows int)
(arrow.RecordBatch, arrow.Array) {
b.Helper()
mem := memory.DefaultAllocator
