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 e9cc75ce perf(compute): vectorize mixed int32 filters with AVX2 (#1284)
e9cc75ce is described below
commit e9cc75ce43a80350885c1a488dc1e160d082df13
Author: Minh Vu <[email protected]>
AuthorDate: Thu Sep 3 17:55:13 2026 +0200
perf(compute): vectorize mixed int32 filters with AVX2 (#1284)
## Summary
* **Add an AVX2 path** for filtering 32-bit fixed-width values with
mixed selection masks.
* **Compress two groups of four values** with `vpshufb` and use masked
stores for the compact output.
* **Keep the current path** for nulls, short or unaligned inputs, dense
masks, unsupported CPUs, and `noasm` builds.
* **Add tests and a benchmark** for mixed masks, offsets, and 1K, 64K,
and 1M rows.
## Benchmark
The benchmark is in `arrow/compute/vector_selection_bench_test.go`.
It covers alternating, random 25/50/75%, clustered 50%, all-selected,
and none-selected masks at 1K, 64K, and 1M rows.
A native AVX2 run needs an x86_64 machine with AVX2. The benchmark also
runs on the local Apple M1, but that only exercises the existing arm64
fallback, so no local AVX2 timing is included.
## Tests
* `go test ./arrow/compute/internal/kernels ./arrow/compute`
* `go test -tags noasm ./arrow/compute/internal/kernels ./arrow/compute`
* `go test -race ./arrow/compute/internal/kernels ./arrow/compute`
* amd64 cross-build for both packages
---
arrow/compute/internal/kernels/Makefile | 8 +-
.../compute/internal/kernels/_lib/filter_uint32.cc | 75 ++++++++++++++++
.../kernels/_lib/filter_uint32_avx2_amd64.s | 84 +++++++++++++++++
.../internal/kernels/filter_uint32_avx2_amd64.go | 99 ++++++++++++++++++++
.../internal/kernels/filter_uint32_avx2_amd64.s | 81 +++++++++++++++++
.../internal/kernels/filter_uint32_noasm.go | 23 +++++
arrow/compute/internal/kernels/vector_selection.go | 7 ++
arrow/compute/vector_selection_bench_test.go | 100 +++++++++++++++++++++
arrow/compute/vector_selection_test.go | 47 ++++++++++
9 files changed, 523 insertions(+), 1 deletion(-)
diff --git a/arrow/compute/internal/kernels/Makefile
b/arrow/compute/internal/kernels/Makefile
index 4e8ddd85..1edf1270 100644
--- a/arrow/compute/internal/kernels/Makefile
+++ b/arrow/compute/internal/kernels/Makefile
@@ -39,7 +39,7 @@ ALL_SOURCES := $(shell find . -path ./_lib -prune -o -name
'*.go' -name '*.s' -n
INTEL_SOURCES := \
cast_numeric_avx2_amd64.s cast_numeric_sse4_amd64.s
constant_factor_avx2_amd64.s \
constant_factor_sse4_amd64.s base_arithmetic_avx2_amd64.s
base_arithmetic_sse4_amd64.s \
- scalar_comparison_avx2_amd64.s scalar_comparison_sse4_amd64.s
+ scalar_comparison_avx2_amd64.s scalar_comparison_sse4_amd64.s
filter_uint32_avx2_amd64.s
#
# ARROW-15336: DO NOT add the assembly target for Arm64 (ARM_SOURCES) until
c2goasm added the Arm64 support.
@@ -70,6 +70,9 @@ _lib/scalar_comparison_avx2_amd64.s: _lib/scalar_comparison.cc
_lib/scalar_comparison_sse4_amd64.s: _lib/scalar_comparison.cc
$(CXX) -std=c++17 -S $(C_FLAGS) $(ASM_FLAGS_SSE4) $^ -o $@ ;
$(PERL_FIXUP_ROTATE) $@
+_lib/filter_uint32_avx2_amd64.s: _lib/filter_uint32.cc
+ $(CXX) -std=c++17 -S $(C_FLAGS) $(ASM_FLAGS_AVX2) $^ -o $@ ;
$(PERL_FIXUP_ROTATE) $@
+
_lib/base_arithmetic_neon.s: _lib/base_arithmetic.cc
$(CXX) -std=c++17 -S $(C_FLAGS_NEON) $^ -o $@ ; $(PERL_FIXUP_ROTATE) $@
@@ -106,6 +109,9 @@ scalar_comparison_avx2_amd64.s:
_lib/scalar_comparison_avx2_amd64.s
scalar_comparison_sse4_amd64.s: _lib/scalar_comparison_sse4_amd64.s
$(C2GOASM) -a -f $^ $@
+filter_uint32_avx2_amd64.s: _lib/filter_uint32_avx2_amd64.s
+ $(C2GOASM) -a -f $^ $@
+
clean:
rm -f $(INTEL_SOURCES)
rm -f $(addprefix _lib/,$(INTEL_SOURCES))
diff --git a/arrow/compute/internal/kernels/_lib/filter_uint32.cc
b/arrow/compute/internal/kernels/_lib/filter_uint32.cc
new file mode 100644
index 00000000..2e72ec3d
--- /dev/null
+++ b/arrow/compute/internal/kernels/_lib/filter_uint32.cc
@@ -0,0 +1,75 @@
+// 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.
+
+#include <arch.h>
+#include <immintrin.h>
+#include <stdint.h>
+
+extern "C" void FULL_NAME(filter_uint32)(const uint32_t* values,
+ const uint8_t* filter,
+ uint32_t* output,
+ const uint8_t* tables,
+ const int64_t length) {
+ const uint8_t* shuffle_masks = tables;
+ const int32_t* store_masks = reinterpret_cast<const int32_t*>(tables +
256);
+ const uint8_t* popcount = tables + 336;
+ int64_t output_length = 0;
+ const int64_t num_bytes = length / 8;
+
+ for (int64_t i = 0; i < num_bytes; ++i) {
+ const uint8_t mask = filter[i];
+ const uint32_t* input = values + i * 8;
+
+ if (mask == 0) {
+ continue;
+ }
+
+ if (mask == 0xff) {
+ const __m256i low = _mm256_loadu_si256(reinterpret_cast<const
__m256i*>(input));
+ _mm256_storeu_si256(reinterpret_cast<__m256i*>(output +
output_length), low);
+ output_length += 8;
+ continue;
+ }
+
+ const uint8_t low_mask = mask & 0xf;
+ const uint8_t high_mask = mask >> 4;
+ const int low_count = popcount[low_mask];
+ const int high_count = popcount[high_mask];
+ const __m128i low_values = _mm_loadu_si128(reinterpret_cast<const
__m128i*>(input));
+ const __m128i high_values = _mm_loadu_si128(reinterpret_cast<const
__m128i*>(input + 4));
+
+ if (low_count != 0) {
+ const __m128i shuffle = _mm_loadu_si128(
+ reinterpret_cast<const __m128i*>(shuffle_masks + low_mask *
16));
+ const __m128i compacted = _mm_shuffle_epi8(low_values, shuffle);
+ const __m128i store_mask = _mm_loadu_si128(
+ reinterpret_cast<const __m128i*>(store_masks + low_count * 4));
+ _mm_maskstore_epi32(reinterpret_cast<int*>(output +
output_length), store_mask, compacted);
+ output_length += low_count;
+ }
+
+ if (high_count != 0) {
+ const __m128i shuffle = _mm_loadu_si128(
+ reinterpret_cast<const __m128i*>(shuffle_masks + high_mask *
16));
+ const __m128i compacted = _mm_shuffle_epi8(high_values, shuffle);
+ const __m128i store_mask = _mm_loadu_si128(
+ reinterpret_cast<const __m128i*>(store_masks + high_count *
4));
+ _mm_maskstore_epi32(reinterpret_cast<int*>(output +
output_length), store_mask, compacted);
+ output_length += high_count;
+ }
+ }
+
+}
diff --git a/arrow/compute/internal/kernels/_lib/filter_uint32_avx2_amd64.s
b/arrow/compute/internal/kernels/_lib/filter_uint32_avx2_amd64.s
new file mode 100644
index 00000000..506e00cd
--- /dev/null
+++ b/arrow/compute/internal/kernels/_lib/filter_uint32_avx2_amd64.s
@@ -0,0 +1,84 @@
+ .intel_syntax noprefix
+ .file "filter_uint32.cc"
+ .text
+ .globl filter_uint32_avx2 # -- Begin function
filter_uint32_avx2
+ .p2align 4
+ .type filter_uint32_avx2,@function
+filter_uint32_avx2: # @filter_uint32_avx2
+# %bb.0:
+ lea rax, [r8 + 7]
+ test r8, r8
+ cmovns rax, r8
+ cmp r8, 8
+ jl .LBB0_11
+# %bb.1:
+ push rbp
+ push r14
+ push rbx
+ sar rax, 3
+ xor r8d, r8d
+ xor r9d, r9d
+ jmp .LBB0_2
+ .p2align 4
+.LBB0_4: # in Loop: Header=BB0_2 Depth=1
+ vmovdqu ymm0, ymmword ptr [rdi]
+ vmovdqu ymmword ptr [rdx + 4*r8], ymm0
+ add r8, 8
+.LBB0_9: # in Loop: Header=BB0_2 Depth=1
+ inc r9
+ add rdi, 32
+ cmp rax, r9
+ je .LBB0_10
+.LBB0_2: # =>This Inner Loop Header: Depth=1
+ movzx r10d, byte ptr [rsi + r9]
+ test r10d, r10d
+ je .LBB0_9
+# %bb.3: # in Loop: Header=BB0_2 Depth=1
+ cmp r10d, 255
+ je .LBB0_4
+# %bb.5: # in Loop: Header=BB0_2 Depth=1
+ mov r11d, r10d
+ and r11d, 15
+ mov r14d, r10d
+ shr r14d, 4
+ movzx ebx, byte ptr [rcx + r11 + 336]
+ movzx r11d, byte ptr [rcx + r14 + 336]
+ vmovdqu xmm0, xmmword ptr [rdi + 16]
+ test rbx, rbx
+ je .LBB0_7
+# %bb.6: # in Loop: Header=BB0_2 Depth=1
+ vmovdqu xmm1, xmmword ptr [rdi]
+ mov ebp, r10d
+ shl bpl, 4
+ movzx r14d, bpl
+ vpshufb xmm1, xmm1, xmmword ptr [rcx + r14]
+ mov r14d, ebx
+ shl r14d, 4
+ vmovdqu xmm2, xmmword ptr [rcx + r14 + 256]
+ vpmaskmovd xmmword ptr [rdx + 4*r8], xmm2, xmm1
+ add r8, rbx
+.LBB0_7: # in Loop: Header=BB0_2 Depth=1
+ test r11, r11
+ je .LBB0_9
+# %bb.8: # in Loop: Header=BB0_2 Depth=1
+ and r10d, -16
+ vpshufb xmm0, xmm0, xmmword ptr [rcx + r10]
+ mov r10d, r11d
+ shl r10d, 4
+ vmovdqu xmm1, xmmword ptr [rcx + r10 + 256]
+ vpmaskmovd xmmword ptr [rdx + 4*r8], xmm1, xmm0
+ add r8, r11
+ jmp .LBB0_9
+.LBB0_10:
+ pop rbx
+ pop r14
+ pop rbp
+.LBB0_11:
+ vzeroupper
+ ret
+.Lfunc_end0:
+ .size filter_uint32_avx2, .Lfunc_end0-filter_uint32_avx2
+ # -- End function
+ .ident "Apple clang version 21.0.0 (clang-2100.1.1.101)"
+ .section ".note.GNU-stack","",@progbits
+ .addrsig
diff --git a/arrow/compute/internal/kernels/filter_uint32_avx2_amd64.go
b/arrow/compute/internal/kernels/filter_uint32_avx2_amd64.go
new file mode 100644
index 00000000..4d7a283b
--- /dev/null
+++ b/arrow/compute/internal/kernels/filter_uint32_avx2_amd64.go
@@ -0,0 +1,99 @@
+// 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 && amd64 && !noasm && !appengine
+
+package kernels
+
+import (
+ "encoding/binary"
+ "math/bits"
+ "unsafe"
+
+ "golang.org/x/sys/cpu"
+)
+
+var filterUint32Tables = makeFilterUint32Tables()
+
+func makeFilterUint32Tables() (tables [352]byte) {
+ for mask := 0; mask < 16; mask++ {
+ pos := 0
+ for lane := 0; lane < 4; lane++ {
+ if mask&(1<<uint(lane)) == 0 {
+ continue
+ }
+ for byteInLane := 0; byteInLane < 4; byteInLane++ {
+ tables[mask*16+pos] = byte(lane*4 + byteInLane)
+ pos++
+ }
+ }
+ for ; pos < 16; pos++ {
+ tables[mask*16+pos] = 0x80
+ }
+ }
+
+ for count := 1; count <= 4; count++ {
+ for lane := 0; lane < count; lane++ {
+
binary.LittleEndian.PutUint32(tables[256+count*16+lane*4:], ^uint32(0))
+ }
+ }
+ for mask := 0; mask < 16; mask++ {
+ tables[336+mask] = byte(bits.OnesCount8(uint8(mask)))
+ }
+ return tables
+}
+
+//go:noescape
+func _filter_uint32_avx2(values, filter, output, tables unsafe.Pointer, length
int64)
+
+func filterUint32Avx2(values []uint32, output []uint32, filterData []byte,
filterOffset, length int64) bool {
+ if !cpu.X86.HasAVX2 || length < 64 || length%8 != 0 || filterOffset%8
!= 0 {
+ return false
+ }
+
+ numBytes := length / 8
+ filterByteOffset := filterOffset / 8
+ if filterByteOffset < 0 || filterByteOffset+numBytes >
int64(len(filterData)) {
+ return false
+ }
+
+ mixedBytes := 0
+ const sampleBytes = 64
+ for i := int64(0); i < numBytes && i < sampleBytes; i++ {
+ mask := filterData[filterByteOffset+i]
+ if mask != 0 && mask != 0xff {
+ mixedBytes++
+ if mixedBytes == 4 {
+ break
+ }
+ }
+ }
+ if mixedBytes < 4 {
+ return false
+ }
+
+ if len(output) == 0 {
+ return false
+ }
+ _filter_uint32_avx2(
+ unsafe.Pointer(&values[0]),
+ unsafe.Pointer(&filterData[filterByteOffset]),
+ unsafe.Pointer(&output[0]),
+ unsafe.Pointer(&filterUint32Tables[0]),
+ length,
+ )
+ return true
+}
diff --git a/arrow/compute/internal/kernels/filter_uint32_avx2_amd64.s
b/arrow/compute/internal/kernels/filter_uint32_avx2_amd64.s
new file mode 100644
index 00000000..6d6c0fca
--- /dev/null
+++ b/arrow/compute/internal/kernels/filter_uint32_avx2_amd64.s
@@ -0,0 +1,81 @@
+//go:build go1.18 && amd64 && !noasm && !appengine
+// AUTO-GENERATED BY C2GOASM -- DO NOT EDIT
+
+TEXT ยท_filter_uint32_avx2(SB), $0-40
+
+ MOVQ values+0(FP), DI
+ MOVQ filter+8(FP), SI
+ MOVQ output+16(FP), DX
+ MOVQ tables+24(FP), CX
+ MOVQ length+32(FP), R8
+
+ LONG $0x07408d49 // lea rax, [r8 + 7]
+ WORD $0x854d; BYTE $0xc0 // test r8, r8
+ LONG $0xc0490f49 // cmovns rax, r8
+ LONG $0x08f88349 // cmp r8, 8
+ JL LBB0_11
+ BYTE $0x55 // push rbp
+ WORD $0x5641 // push r14
+ BYTE $0x53 // push rbx
+ LONG $0x03f8c148 // sar rax, 3
+ WORD $0x3145; BYTE $0xc0 // xor r8d, r8d
+ WORD $0x3145; BYTE $0xc9 // xor r9d, r9d
+ JMP LBB0_2
+
+LBB0_4:
+ LONG $0x076ffec5 // vmovdqu ymm0, yword [rdi]
+ LONG $0x7f7ea1c4; WORD $0x8204 // vmovdqu yword [rdx + 4*r8], ymm0
+ LONG $0x08c08349 // add r8, 8
+
+LBB0_9:
+ WORD $0xff49; BYTE $0xc1 // inc r9
+ LONG $0x20c78348 // add rdi, 32
+ WORD $0x394c; BYTE $0xc8 // cmp rax, r9
+ JE LBB0_10
+
+LBB0_2:
+ LONG $0x14b60f46; BYTE $0x0e // movzx r10d, byte [rsi
+ r9]
+ WORD $0x8545; BYTE $0xd2 // test r10d, r10d
+ JE LBB0_9
+ LONG $0xfffa8141; WORD $0x0000; BYTE $0x00 // cmp r10d, 255
+ JE LBB0_4
+ WORD $0x8945; BYTE $0xd3 // mov r11d, r10d
+ LONG $0x0fe38341 // and r11d, 15
+ WORD $0x8945; BYTE $0xd6 // mov r14d, r10d
+ LONG $0x04eec141 // shr r14d, 4
+ QUAD $0x000150199cb60f42; BYTE $0x00 // movzx ebx, byte [rcx +
r11 + 336]
+ QUAD $0x000150319cb60f46; BYTE $0x00 // movzx r11d, byte [rcx
+ r14 + 336]
+ LONG $0x476ffac5; BYTE $0x10 // vmovdqu xmm0, oword
[rdi + 16]
+ WORD $0x8548; BYTE $0xdb // test rbx, rbx
+ JE LBB0_7
+ LONG $0x0f6ffac5 // vmovdqu xmm1, oword
[rdi]
+ WORD $0x8944; BYTE $0xd5 // mov ebp, r10d
+ LONG $0x04e5c040 // shl bpl, 4
+ LONG $0xf5b60f44 // movzx r14d, bpl
+ LONG $0x0071a2c4; WORD $0x310c // vpshufb xmm1, xmm1,
oword [rcx + r14]
+ WORD $0x8941; BYTE $0xde // mov r14d, ebx
+ LONG $0x04e6c141 // shl r14d, 4
+ QUAD $0x010031946f7aa1c4; WORD $0x0000 // vmovdqu xmm2, oword
[rcx + r14 + 256]
+ LONG $0x8e69a2c4; WORD $0x820c // vpmaskmovd oword [rdx
+ 4*r8], xmm2, xmm1
+ WORD $0x0149; BYTE $0xd8 // add r8, rbx
+
+LBB0_7:
+ WORD $0x854d; BYTE $0xdb // test r11, r11
+ JE LBB0_9
+ LONG $0xf0e28341 // and r10d, -16
+ LONG $0x0079a2c4; WORD $0x1104 // vpshufb xmm0, xmm0, oword
[rcx + r10]
+ WORD $0x8945; BYTE $0xda // mov r10d, r11d
+ LONG $0x04e2c141 // shl r10d, 4
+ QUAD $0x0100118c6f7aa1c4; WORD $0x0000 // vmovdqu xmm1, oword [rcx +
r10 + 256]
+ LONG $0x8e71a2c4; WORD $0x8204 // vpmaskmovd oword [rdx +
4*r8], xmm1, xmm0
+ WORD $0x014d; BYTE $0xd8 // add r8, r11
+ JMP LBB0_9
+
+LBB0_10:
+ BYTE $0x5b // pop rbx
+ WORD $0x5e41 // pop r14
+ BYTE $0x5d // pop rbp
+
+LBB0_11:
+ VZEROUPPER
+ RET
diff --git a/arrow/compute/internal/kernels/filter_uint32_noasm.go
b/arrow/compute/internal/kernels/filter_uint32_noasm.go
new file mode 100644
index 00000000..10f9b616
--- /dev/null
+++ b/arrow/compute/internal/kernels/filter_uint32_noasm.go
@@ -0,0 +1,23 @@
+// 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 || !amd64 || appengine)
+
+package kernels
+
+func filterUint32Avx2([]uint32, []uint32, []byte, int64, int64) bool {
+ return false
+}
diff --git a/arrow/compute/internal/kernels/vector_selection.go
b/arrow/compute/internal/kernels/vector_selection.go
index cbe0c220..68198628 100644
--- a/arrow/compute/internal/kernels/vector_selection.go
+++ b/arrow/compute/internal/kernels/vector_selection.go
@@ -482,6 +482,13 @@ func PrimitiveFilter(ctx *exec.KernelCtx, batch
*exec.ExecSpan, out *exec.ExecRe
allocateValidity := values.Nulls != 0 || filter.Nulls != 0
bitWidth := values.Type.(arrow.FixedWidthDataType).BitWidth()
preallocateData(ctx, outputLength, bitWidth, allocateValidity, out)
+ if bitWidth == 32 && values.Nulls == 0 && filter.Nulls == 0 {
+ valuesData := exec.GetSpanValues[uint32](values, 1)
+ outData := exec.GetSpanValues[uint32](out, 1)
+ if filterUint32Avx2(valuesData, outData, filter.Buffers[1].Buf,
filter.Offset, values.Len) {
+ return nil
+ }
+ }
var wr writeFiltered
switch bitWidth {
diff --git a/arrow/compute/vector_selection_bench_test.go
b/arrow/compute/vector_selection_bench_test.go
new file mode 100644
index 00000000..b5e8993e
--- /dev/null
+++ b/arrow/compute/vector_selection_bench_test.go
@@ -0,0 +1,100 @@
+// 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 benchmarkFilterOutputLength int
+
+func BenchmarkFilterInt32MixedMasks(b *testing.B) {
+ patterns := []struct {
+ name string
+ selected func(int) bool
+ }{
+ {name: "alternating", selected: func(i int) bool { return i%2
== 0 }},
+ {name: "random25", selected: func(i int) bool { return
filterBenchmarkRandom(i, 25) }},
+ {name: "random50", selected: func(i int) bool { return
filterBenchmarkRandom(i, 50) }},
+ {name: "random75", selected: func(i int) bool { return
filterBenchmarkRandom(i, 75) }},
+ {name: "clustered50", selected: func(i int) bool { return
(i/32)%2 == 0 }},
+ {name: "all-selected", selected: func(int) bool { return true
}},
+ {name: "none-selected", selected: func(int) bool { return false
}},
+ }
+
+ for _, size := range []int{1 << 10, 1 << 16, 1 << 20} {
+ size := size
+ for _, pattern := range patterns {
+ pattern := pattern
+ b.Run(fmt.Sprintf("size=%d/%s", size, pattern.name),
func(b *testing.B) {
+ values, filter :=
makeFilterInt32BenchmarkInput(b, size, pattern.selected)
+ defer values.Release()
+ defer filter.Release()
+
+ b.ReportAllocs()
+ b.SetBytes(int64(size * 4))
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ result, err :=
compute.FilterArray(context.Background(), values, filter,
*compute.DefaultFilterOptions())
+ if err != nil {
+ b.Fatal(err)
+ }
+ benchmarkFilterOutputLength =
result.Len()
+ result.Release()
+ }
+ })
+ }
+ }
+}
+
+func makeFilterInt32BenchmarkInput(b *testing.B, size int, selected func(int)
bool) (arrow.Array, arrow.Array) {
+ b.Helper()
+ mem := memory.DefaultAllocator
+
+ valuesBuilder := array.NewInt32Builder(mem)
+ valuesBuilder.Reserve(size)
+ for i := 0; i < size; i++ {
+ valuesBuilder.Append(int32(i))
+ }
+ values := valuesBuilder.NewInt32Array()
+ valuesBuilder.Release()
+
+ filterBuilder := array.NewBooleanBuilder(mem)
+ filterBuilder.Reserve(size)
+ for i := 0; i < size; i++ {
+ filterBuilder.Append(selected(i))
+ }
+ filter := filterBuilder.NewBooleanArray()
+ filterBuilder.Release()
+ return values, filter
+}
+
+func filterBenchmarkRandom(i, selectedPercent int) bool {
+ x := uint32(i)*747796405 + 2891336453
+ x = ((x >> ((x >> 28) + 4)) ^ x) * 277803737
+ x = (x >> 22) ^ x
+ return int(x%100) < selectedPercent
+}
diff --git a/arrow/compute/vector_selection_test.go
b/arrow/compute/vector_selection_test.go
index 08d41bae..0d9d9f04 100644
--- a/arrow/compute/vector_selection_test.go
+++ b/arrow/compute/vector_selection_test.go
@@ -2032,6 +2032,53 @@ func TestFilterKernels(t *testing.T) {
suite.Run(t, new(FilterKernelWithTable))
}
+func TestFilterInt32MixedMaskOffsets(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ const length = 128
+ for _, offset := range []int64{0, 3, 8} {
+ t.Run(fmt.Sprintf("offset=%d", offset), func(t *testing.T) {
+ valuesBuilder := array.NewInt32Builder(mem)
+ valuesBuilder.Reserve(int(offset) + length)
+ for i := int64(0); i < offset+length; i++ {
+ valuesBuilder.Append(int32(i))
+ }
+ valuesBase := valuesBuilder.NewInt32Array()
+ valuesBuilder.Release()
+ values := array.NewSlice(valuesBase, offset,
offset+length)
+ valuesBase.Release()
+ defer values.Release()
+
+ filterBuilder := array.NewBooleanBuilder(mem)
+ filterBuilder.Reserve(int(offset) + length)
+ for i := int64(0); i < offset+length; i++ {
+ filterBuilder.Append(i%2 == 0)
+ }
+ filterBase := filterBuilder.NewBooleanArray()
+ filterBuilder.Release()
+ filter := array.NewSlice(filterBase, offset,
offset+length)
+ filterBase.Release()
+ defer filter.Release()
+
+ expectedBuilder := array.NewInt32Builder(mem)
+ for i := offset; i < offset+length; i++ {
+ if i%2 == 0 {
+ expectedBuilder.Append(int32(i))
+ }
+ }
+ expected := expectedBuilder.NewInt32Array()
+ expectedBuilder.Release()
+ defer expected.Release()
+
+ actual, err :=
compute.FilterArray(context.Background(), values, filter,
*compute.DefaultFilterOptions())
+ require.NoError(t, err)
+ defer actual.Release()
+ assertArraysEqual(t, expected, actual)
+ })
+ }
+}
+
// Benchmark tests for Take operation with variable-length data
// These benchmarks test the performance improvements from buffer
pre-allocation
// in VarBinaryImpl for string/binary data reorganization (e.g., partitioning).
