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 8ae60185 perf(compute): vectorize mixed 16-bit filters (#1349)
8ae60185 is described below
commit 8ae601850b60514ebf5846f94f3bb75924e350a5
Author: Minh Vu <[email protected]>
AuthorDate: Tue Oct 6 18:19:28 2026 +0200
perf(compute): vectorize mixed 16-bit filters (#1349)
## What
- Add AVX2 and ARM64 NEON paths for mixed null-free 16-bit filters.
## Why
- 16-bit primitive filters still use the generic writer while 32-bit and
64-bit filters have SIMD paths.
## Implementation
- Keep nulls, short or unaligned filters, dense masks, and noasm on the
existing path.
- Benchmark alternating, random25/50/75, clustered50, all-selected, and
none-selected at 1K, 64K, and 1M values.
- On M1 Pro, compared with noasm: 1M alternating improves from 4.23 ms
to 225 us; random50 improves from 3.52 ms to 1.09 ms. Mixed cases drop
from 33 to 32 allocs/op.
---
arrow/compute/internal/kernels/Makefile | 9 +-
.../compute/internal/kernels/_lib/filter_uint16.cc | 80 +++++++++++
.../kernels/_lib/filter_uint16_avx2_amd64.s | 92 ++++++++++++
arrow/compute/internal/kernels/filter_uint16.go | 69 +++++++++
.../internal/kernels/filter_uint16_arch_test.go | 155 +++++++++++++++++++++
.../internal/kernels/filter_uint16_avx2_amd64.go | 50 +++++++
.../internal/kernels/filter_uint16_avx2_amd64.s | 93 +++++++++++++
.../kernels/filter_uint16_avx2_amd64_test.go | 53 +++++++
.../internal/kernels/filter_uint16_avx2_noasm.go | 23 +++
.../internal/kernels/filter_uint16_neon_arm64.go | 50 +++++++
.../internal/kernels/filter_uint16_neon_arm64.s | 120 ++++++++++++++++
.../kernels/filter_uint16_neon_arm64_test.go | 53 +++++++
.../internal/kernels/filter_uint16_neon_noasm.go | 23 +++
arrow/compute/internal/kernels/vector_selection.go | 8 ++
.../internal/kernels/vector_selection_test.go | 56 ++++++++
arrow/compute/vector_selection_bench_test.go | 78 ++++++++++-
arrow/compute/vector_selection_test.go | 47 +++++++
17 files changed, 1055 insertions(+), 4 deletions(-)
diff --git a/arrow/compute/internal/kernels/Makefile
b/arrow/compute/internal/kernels/Makefile
index a1a8b419..d33a8a9a 100644
--- a/arrow/compute/internal/kernels/Makefile
+++ b/arrow/compute/internal/kernels/Makefile
@@ -42,7 +42,7 @@ 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 \
- filter_uint32_avx2_amd64.s filter_uint64_avx2_amd64.s
+ filter_uint32_avx2_amd64.s filter_uint64_avx2_amd64.s
filter_uint16_avx2_amd64.s
#
# ARROW-15336: DO NOT add the assembly target for Arm64 (ARM_SOURCES) until
c2goasm added the Arm64 support.
@@ -79,6 +79,9 @@ _lib/filter_uint32_avx2_amd64.s: _lib/filter_uint32.cc
_lib/filter_uint64_avx2_amd64.s: _lib/filter_uint64.cc
$(CXX) -std=c++17 -S $(C_FLAGS) $(ASM_FLAGS_AVX2) -mno-stackrealign
-fomit-frame-pointer $^ -o $@ ; $(PERL_FIXUP_ROTATE) $@
+_lib/filter_uint16_avx2_amd64.s: _lib/filter_uint16.cc
+ $(CXX) -std=c++17 -S $(C_FLAGS) $(ASM_FLAGS_AVX2) -mno-stackrealign
-fomit-frame-pointer $^ -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) $@
@@ -122,6 +125,10 @@ filter_uint64_avx2_amd64.s: _lib/filter_uint64_avx2_amd64.s
$(C2GOASM) -a -f $^ $@
perl -i -pe 's{^//\s*\+build\b.*}{//go:build go1.18 && amd64 && !noasm
&& !appengine}' $@
+filter_uint16_avx2_amd64.s: _lib/filter_uint16_avx2_amd64.s
+ $(C2GOASM) -a -f $^ $@
+ perl -i -pe 's{^//\s*\+build\b.*}{//go:build go1.18 && amd64 && !noasm
&& !appengine}' $@
+
clean:
rm -f $(INTEL_SOURCES)
rm -f $(addprefix _lib/,$(INTEL_SOURCES))
diff --git a/arrow/compute/internal/kernels/_lib/filter_uint16.cc
b/arrow/compute/internal/kernels/_lib/filter_uint16.cc
new file mode 100644
index 00000000..749f6db1
--- /dev/null
+++ b/arrow/compute/internal/kernels/_lib/filter_uint16.cc
@@ -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.
+
+#include <arch.h>
+#include <immintrin.h>
+#include <stdint.h>
+
+extern "C" void FULL_NAME(filter_uint16)(const uint16_t* values,
+ const uint8_t* filter,
+ uint16_t* output,
+ const uint8_t* tables,
+ const int64_t length) {
+ const uint8_t* shuffle_masks = tables;
+ const uint8_t* popcount = tables + 4096;
+ 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 uint16_t* input = values + i * 8;
+
+ if (mask == 0) {
+ continue;
+ }
+
+ if (mask == 0xff) {
+ const __m128i input_values = _mm_loadu_si128(
+ reinterpret_cast<const __m128i*>(input));
+ _mm_storeu_si128(reinterpret_cast<__m128i*>(output +
output_length), input_values);
+ output_length += 8;
+ continue;
+ }
+
+ const int count = popcount[mask];
+ const __m128i input_values = _mm_loadu_si128(
+ reinterpret_cast<const __m128i*>(input));
+ const __m128i shuffle = _mm_loadu_si128(
+ reinterpret_cast<const __m128i*>(shuffle_masks + mask * 16));
+ const __m128i compacted = _mm_shuffle_epi8(input_values, shuffle);
+
+ if (count >= 4) {
+ _mm_storel_epi64(reinterpret_cast<__m128i*>(output +
output_length), compacted);
+ output_length += 4;
+ if (count >= 7) {
+ output[output_length + 2] =
static_cast<uint16_t>(_mm_extract_epi16(compacted, 6));
+ }
+ if (count >= 6) {
+ output[output_length + 1] =
static_cast<uint16_t>(_mm_extract_epi16(compacted, 5));
+ }
+ if (count >= 5) {
+ output[output_length] =
static_cast<uint16_t>(_mm_extract_epi16(compacted, 4));
+ }
+ output_length += count - 4;
+ } else {
+ if (count >= 3) {
+ output[output_length + 2] =
static_cast<uint16_t>(_mm_extract_epi16(compacted, 2));
+ }
+ if (count >= 2) {
+ output[output_length + 1] =
static_cast<uint16_t>(_mm_extract_epi16(compacted, 1));
+ }
+ if (count >= 1) {
+ output[output_length] =
static_cast<uint16_t>(_mm_extract_epi16(compacted, 0));
+ }
+ output_length += count;
+ }
+ }
+}
diff --git a/arrow/compute/internal/kernels/_lib/filter_uint16_avx2_amd64.s
b/arrow/compute/internal/kernels/_lib/filter_uint16_avx2_amd64.s
new file mode 100644
index 00000000..b73dcacc
--- /dev/null
+++ b/arrow/compute/internal/kernels/_lib/filter_uint16_avx2_amd64.s
@@ -0,0 +1,92 @@
+ .intel_syntax noprefix
+ .file "filter_uint16.cc"
+ .text
+ .globl filter_uint16_avx2 # -- Begin function
filter_uint16_avx2
+ .p2align 4
+ .type filter_uint16_avx2,@function
+filter_uint16_avx2: # @filter_uint16_avx2
+# %bb.0:
+ lea rax, [r8 + 7]
+ test r8, r8
+ cmovns rax, r8
+ cmp r8, 8
+ jl .LBB0_21
+# %bb.1:
+ sar rax, 3
+ xor r8d, r8d
+ xor r9d, r9d
+ jmp .LBB0_2
+ .p2align 4
+.LBB0_4: # in Loop: Header=BB0_2 Depth=1
+ vmovdqu xmm0, xmmword ptr [rdi]
+ vmovdqu xmmword ptr [rdx + 2*r8], xmm0
+ add r8, 8
+.LBB0_20: # in Loop: Header=BB0_2 Depth=1
+ inc r9
+ add rdi, 16
+ cmp rax, r9
+ je .LBB0_21
+.LBB0_2: # =>This Inner Loop Header: Depth=1
+ movzx r11d, byte ptr [rsi + r9]
+ test r11d, r11d
+ je .LBB0_20
+# %bb.3: # in Loop: Header=BB0_2 Depth=1
+ cmp r11d, 255
+ je .LBB0_4
+# %bb.5: # in Loop: Header=BB0_2 Depth=1
+ movzx r10d, byte ptr [rcx + r11 + 4096]
+ shl r11d, 4
+ vmovdqu xmm0, xmmword ptr [rdi]
+ vpshufb xmm0, xmm0, xmmword ptr [rcx + r11]
+ cmp r10, 4
+ jb .LBB0_13
+# %bb.6: # in Loop: Header=BB0_2 Depth=1
+ vmovq qword ptr [rdx + 2*r8], xmm0
+ cmp r10b, 7
+ jb .LBB0_8
+# %bb.7: # in Loop: Header=BB0_2 Depth=1
+ vpextrw word ptr [rdx + 2*r8 + 12], xmm0, 6
+ jmp .LBB0_10
+.LBB0_13: # in Loop: Header=BB0_2 Depth=1
+ cmp r10d, 3
+ jne .LBB0_15
+# %bb.14: # in Loop: Header=BB0_2 Depth=1
+ vpextrw word ptr [rdx + 2*r8 + 4], xmm0, 2
+ vpextrw word ptr [rdx + 2*r8 + 2], xmm0, 1
+ jmp .LBB0_18
+.LBB0_8: # in Loop: Header=BB0_2 Depth=1
+ cmp r10, 4
+ je .LBB0_12
+# %bb.9: # in Loop: Header=BB0_2 Depth=1
+ cmp r10d, 6
+ jne .LBB0_11
+.LBB0_10: # in Loop: Header=BB0_2 Depth=1
+ vpextrw word ptr [rdx + 2*r8 + 10], xmm0, 5
+.LBB0_11: # in Loop: Header=BB0_2 Depth=1
+ vpextrw word ptr [rdx + 2*r8 + 8], xmm0, 4
+.LBB0_12: # in Loop: Header=BB0_2 Depth=1
+ add r8, 4
+ add r10d, -4
+ jmp .LBB0_19
+.LBB0_15: # in Loop: Header=BB0_2 Depth=1
+ cmp r10b, 2
+ jb .LBB0_17
+# %bb.16: # in Loop: Header=BB0_2 Depth=1
+ vpextrw word ptr [rdx + 2*r8 + 2], xmm0, 1
+ jmp .LBB0_18
+.LBB0_17: # in Loop: Header=BB0_2 Depth=1
+ test r10, r10
+ je .LBB0_19
+.LBB0_18: # in Loop: Header=BB0_2 Depth=1
+ vpextrw word ptr [rdx + 2*r8], xmm0, 0
+.LBB0_19: # in Loop: Header=BB0_2 Depth=1
+ add r8, r10
+ jmp .LBB0_20
+.LBB0_21:
+ ret
+.Lfunc_end0:
+ .size filter_uint16_avx2, .Lfunc_end0-filter_uint16_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_uint16.go
b/arrow/compute/internal/kernels/filter_uint16.go
new file mode 100644
index 00000000..51d61844
--- /dev/null
+++ b/arrow/compute/internal/kernels/filter_uint16.go
@@ -0,0 +1,69 @@
+// 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 "math/bits"
+
+func makeFilterUint16Tables() (tables [4352]byte) {
+ for mask := 0; mask < 256; mask++ {
+ pos := 0
+ for lane := 0; lane < 8; lane++ {
+ if mask&(1<<uint(lane)) == 0 {
+ continue
+ }
+ tables[mask*16+pos] = byte(lane * 2)
+ tables[mask*16+pos+1] = byte(lane*2 + 1)
+ pos += 2
+ }
+ for ; pos < 16; pos++ {
+ tables[mask*16+pos] = 0x80
+ }
+ tables[4096+mask] = byte(bits.OnesCount8(uint8(mask)))
+ }
+ return tables
+}
+
+func filterUint16VectorInput(filterData []byte, filterOffset, length int64)
([]byte, bool) {
+ if length < 64 || length%8 != 0 || filterOffset%8 != 0 {
+ return nil, false
+ }
+
+ numBytes := length / 8
+ filterByteOffset := filterOffset / 8
+ if filterByteOffset < 0 || filterByteOffset+numBytes >
int64(len(filterData)) {
+ return nil, false
+ }
+ filterBytes := filterData[filterByteOffset : filterByteOffset+numBytes]
+
+ const (
+ sampleBytes = 64
+ minMixed = 4
+ )
+ mixedBytes := 0
+ for i := int64(0); i < int64(len(filterBytes)) && i < sampleBytes; i++ {
+ mask := filterBytes[i]
+ if mask != 0 && mask != 0xff {
+ mixedBytes++
+ if mixedBytes == minMixed {
+ return filterBytes, true
+ }
+ }
+ }
+ return nil, false
+}
diff --git a/arrow/compute/internal/kernels/filter_uint16_arch_test.go
b/arrow/compute/internal/kernels/filter_uint16_arch_test.go
new file mode 100644
index 00000000..1eaf70f4
--- /dev/null
+++ b/arrow/compute/internal/kernels/filter_uint16_arch_test.go
@@ -0,0 +1,155 @@
+// 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 || arm64) && !noasm && !appengine
+
+package kernels
+
+import "testing"
+
+type filterUint16Func func(values []uint16, output []uint16, filter []byte,
filterOffset, length int64) bool
+
+func testFilterUint16AllMasks(t *testing.T, filter filterUint16Func) {
+ t.Helper()
+
+ const (
+ filterBytes = 256
+ length = filterBytes * 8
+ )
+ values := make([]uint16, length)
+ filterData := make([]byte, filterBytes)
+ for i := range values {
+ values[i] = uint16(i*977 + 0x8011)
+ }
+ for i := range filterData {
+ filterData[i] = byte(i)
+ }
+
+ expected := make([]uint16, 0, length)
+ for i, mask := range filterData {
+ for lane := 0; lane < 8; lane++ {
+ if mask&(1<<uint(lane)) != 0 {
+ expected = append(expected, values[i*8+lane])
+ }
+ }
+ }
+
+ backing := make([]uint16, len(expected)+1)
+ backing[len(expected)] = 0xdead
+ if !filter(values, backing[:len(expected)], filterData, 0, length) {
+ t.Fatal("filterUint16 did not select the mixed-mask path")
+ }
+ for i, want := range expected {
+ if got := backing[i]; got != want {
+ t.Fatalf("output[%d] = %#x, want %#x", i, got, want)
+ }
+ }
+ if got := backing[len(expected)]; got != 0xdead {
+ t.Fatalf("sentinel = %#x, want %#x", got, uint16(0xdead))
+ }
+}
+
+func testFilterUint16AlignedOffset(t *testing.T, filter filterUint16Func) {
+ t.Helper()
+
+ const (
+ offset = 8
+ length = 64
+ )
+ values := make([]uint16, length)
+ for i := range values {
+ values[i] = uint16(i*31 + 0x8000)
+ }
+ filterData := make([]byte, offset/8+length/8)
+ for i := range filterData[offset/8:] {
+ filterData[offset/8+i] = 0x55
+ }
+
+ output := make([]uint16, length/2)
+ if !filter(values, output, filterData, offset, length) {
+ t.Fatal("filterUint16 did not select the aligned-offset path")
+ }
+ for i, got := range output {
+ want := values[2*i]
+ if got != want {
+ t.Fatalf("output[%d] = %#x, want %#x", i, got, want)
+ }
+ }
+}
+
+func testFilterUint16UnalignedValues(t *testing.T, filter filterUint16Func) {
+ t.Helper()
+
+ const length = 64
+ valuesBacking := make([]uint16, length+3)
+ values := valuesBacking[3:]
+ for i := range values {
+ values[i] = uint16(i*31 + 0x8000)
+ }
+ filterData := make([]byte, length/8)
+ for i := range filterData {
+ filterData[i] = 0x55
+ }
+
+ output := make([]uint16, length/2)
+ if !filter(values, output, filterData, 0, length) {
+ t.Fatal("filterUint16 did not select the unaligned-values path")
+ }
+ for i, got := range output {
+ want := values[2*i]
+ if got != want {
+ t.Fatalf("output[%d] = %#x, want %#x", i, got, want)
+ }
+ }
+}
+
+func testFilterUint16Guards(t *testing.T, filter filterUint16Func) {
+ t.Helper()
+
+ values := make([]uint16, 128)
+ mixed := make([]byte, 16)
+ for i := range mixed {
+ mixed[i] = 0x55
+ }
+ output := make([]uint16, 64)
+ if filter(values[:63], output, mixed[:8], 0, 64) {
+ t.Fatal("filterUint16 accepted values shorter than length")
+ }
+
+ tests := []struct {
+ name string
+ filter []byte
+ offset int64
+ length int64
+ want bool
+ }{
+ {name: "short", filter: mixed, length: 56},
+ {name: "non-multiple-of-eight", filter: mixed, length: 71},
+ {name: "unaligned-filter-offset", filter: mixed, offset: 3,
length: 64},
+ {name: "negative-filter-offset", filter: mixed, offset: -8,
length: 64},
+ {name: "all-false", filter: make([]byte, 8), length: 64},
+ {name: "all-true", filter: []byte{0xff, 0xff, 0xff, 0xff, 0xff,
0xff, 0xff, 0xff}, length: 64},
+ {name: "too-few-mixed-bytes", filter: []byte{0x55, 0x55, 0x55,
0, 0, 0, 0, 0}, length: 64},
+ {name: "out-of-bounds-filter", filter: mixed[:4], length: 64},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ if got := filter(values, output, tt.filter, tt.offset,
tt.length); got != tt.want {
+ t.Fatalf("filterUint16() = %t, want %t", got,
tt.want)
+ }
+ })
+ }
+}
diff --git a/arrow/compute/internal/kernels/filter_uint16_avx2_amd64.go
b/arrow/compute/internal/kernels/filter_uint16_avx2_amd64.go
new file mode 100644
index 00000000..518c2b82
--- /dev/null
+++ b/arrow/compute/internal/kernels/filter_uint16_avx2_amd64.go
@@ -0,0 +1,50 @@
+// 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 (
+ "unsafe"
+
+ "golang.org/x/sys/cpu"
+)
+
+var filterUint16Avx2Tables = makeFilterUint16Tables()
+
+//go:noescape
+func _filter_uint16_avx2(values, filter, output, tables unsafe.Pointer, length
int64)
+
+func filterUint16Avx2(values []uint16, output []uint16, filterData []byte,
filterOffset, length int64) bool {
+ if !cpu.X86.HasAVX2 || length > int64(len(values)) || len(output) == 0 {
+ return false
+ }
+
+ filterBytes, ok := filterUint16VectorInput(filterData, filterOffset,
length)
+ if !ok {
+ return false
+ }
+
+ _filter_uint16_avx2(
+ unsafe.Pointer(unsafe.SliceData(values)),
+ unsafe.Pointer(unsafe.SliceData(filterBytes)),
+ unsafe.Pointer(unsafe.SliceData(output)),
+ unsafe.Pointer(unsafe.SliceData(filterUint16Avx2Tables[:])),
+ length,
+ )
+ return true
+}
diff --git a/arrow/compute/internal/kernels/filter_uint16_avx2_amd64.s
b/arrow/compute/internal/kernels/filter_uint16_avx2_amd64.s
new file mode 100644
index 00000000..817a9712
--- /dev/null
+++ b/arrow/compute/internal/kernels/filter_uint16_avx2_amd64.s
@@ -0,0 +1,93 @@
+//go:build go1.18 && amd64 && !noasm && !appengine
+// AUTO-GENERATED BY C2GOASM -- DO NOT EDIT
+
+TEXT ·_filter_uint16_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_21
+ 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 $0x076ffac5 // vmovdqu xmm0, oword [rdi]
+ LONG $0x7f7aa1c4; WORD $0x4204 // vmovdqu oword [rdx + 2*r8], xmm0
+ LONG $0x08c08349 // add r8, 8
+
+LBB0_20:
+ WORD $0xff49; BYTE $0xc1 // inc r9
+ LONG $0x10c78348 // add rdi, 16
+ WORD $0x394c; BYTE $0xc8 // cmp rax, r9
+ JE LBB0_21
+
+LBB0_2:
+ LONG $0x1cb60f46; BYTE $0x0e // movzx r11d, byte [rsi
+ r9]
+ WORD $0x8545; BYTE $0xdb // test r11d, r11d
+ JE LBB0_20
+ LONG $0xfffb8141; WORD $0x0000; BYTE $0x00 // cmp r11d, 255
+ JE LBB0_4
+ QUAD $0x0010001994b60f46; BYTE $0x00 // movzx r10d, byte [rcx
+ r11 + 4096]
+ LONG $0x04e3c141 // shl r11d, 4
+ LONG $0x076ffac5 // vmovdqu xmm0, oword
[rdi]
+ LONG $0x0079a2c4; WORD $0x1904 // vpshufb xmm0, xmm0,
oword [rcx + r11]
+ LONG $0x04fa8349 // cmp r10, 4
+ JB LBB0_13
+ LONG $0xd679a1c4; WORD $0x4204 // vmovq qword [rdx +
2*r8], xmm0
+ LONG $0x07fa8041 // cmp r10b, 7
+ JB LBB0_8
+ QUAD $0x060c42441579a3c4 // vpextrw word [rdx +
2*r8 + 12], xmm0, 6
+ JMP LBB0_10
+
+LBB0_13:
+ LONG $0x03fa8341 // cmp r10d, 3
+ JNE LBB0_15
+ QUAD $0x020442441579a3c4 // vpextrw word [rdx + 2*r8 + 4], xmm0, 2
+ QUAD $0x010242441579a3c4 // vpextrw word [rdx + 2*r8 + 2], xmm0, 1
+ JMP LBB0_18
+
+LBB0_8:
+ LONG $0x04fa8349 // cmp r10, 4
+ JE LBB0_12
+ LONG $0x06fa8341 // cmp r10d, 6
+ JNE LBB0_11
+
+LBB0_10:
+ QUAD $0x050a42441579a3c4 // vpextrw word [rdx + 2*r8 + 10], xmm0, 5
+
+LBB0_11:
+ QUAD $0x040842441579a3c4 // vpextrw word [rdx + 2*r8 + 8], xmm0, 4
+
+LBB0_12:
+ LONG $0x04c08349 // add r8, 4
+ LONG $0xfcc28341 // add r10d, -4
+ JMP LBB0_19
+
+LBB0_15:
+ LONG $0x02fa8041 // cmp r10b, 2
+ JB LBB0_17
+ QUAD $0x010242441579a3c4 // vpextrw word [rdx + 2*r8 + 2], xmm0, 1
+ JMP LBB0_18
+
+LBB0_17:
+ WORD $0x854d; BYTE $0xd2 // test r10, r10
+ JE LBB0_19
+
+LBB0_18:
+ LONG $0x1579a3c4; WORD $0x4204; BYTE $0x00 // vpextrw word [rdx +
2*r8], xmm0, 0
+
+LBB0_19:
+ WORD $0x014d; BYTE $0xd0 // add r8, r10
+ JMP LBB0_20
+
+LBB0_21:
+ RET
diff --git a/arrow/compute/internal/kernels/filter_uint16_avx2_amd64_test.go
b/arrow/compute/internal/kernels/filter_uint16_avx2_amd64_test.go
new file mode 100644
index 00000000..97e78ce1
--- /dev/null
+++ b/arrow/compute/internal/kernels/filter_uint16_avx2_amd64_test.go
@@ -0,0 +1,53 @@
+// 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 (
+ "testing"
+
+ "golang.org/x/sys/cpu"
+)
+
+func TestFilterUint16Avx2AllMasks(t *testing.T) {
+ if !cpu.X86.HasAVX2 {
+ t.Skip("AVX2 is not available")
+ }
+ testFilterUint16AllMasks(t, filterUint16Avx2)
+}
+
+func TestFilterUint16Avx2AlignedOffset(t *testing.T) {
+ if !cpu.X86.HasAVX2 {
+ t.Skip("AVX2 is not available")
+ }
+ testFilterUint16AlignedOffset(t, filterUint16Avx2)
+}
+
+func TestFilterUint16Avx2UnalignedValues(t *testing.T) {
+ if !cpu.X86.HasAVX2 {
+ t.Skip("AVX2 is not available")
+ }
+ testFilterUint16UnalignedValues(t, filterUint16Avx2)
+}
+
+func TestFilterUint16Avx2Guards(t *testing.T) {
+ if !cpu.X86.HasAVX2 {
+ t.Skip("AVX2 is not available")
+ }
+ testFilterUint16Guards(t, filterUint16Avx2)
+}
diff --git a/arrow/compute/internal/kernels/filter_uint16_avx2_noasm.go
b/arrow/compute/internal/kernels/filter_uint16_avx2_noasm.go
new file mode 100644
index 00000000..17e52abb
--- /dev/null
+++ b/arrow/compute/internal/kernels/filter_uint16_avx2_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 filterUint16Avx2([]uint16, []uint16, []byte, int64, int64) bool {
+ return false
+}
diff --git a/arrow/compute/internal/kernels/filter_uint16_neon_arm64.go
b/arrow/compute/internal/kernels/filter_uint16_neon_arm64.go
new file mode 100644
index 00000000..f2c9c3fe
--- /dev/null
+++ b/arrow/compute/internal/kernels/filter_uint16_neon_arm64.go
@@ -0,0 +1,50 @@
+// 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"
+
+ "golang.org/x/sys/cpu"
+)
+
+var filterUint16NeonTables = makeFilterUint16Tables()
+
+//go:noescape
+func _filter_uint16_neon(values, filter, output, tables unsafe.Pointer, length
int64)
+
+func filterUint16Neon(values []uint16, output []uint16, filterData []byte,
filterOffset, length int64) bool {
+ if !cpu.ARM64.HasASIMD || length > int64(len(values)) || len(output) ==
0 {
+ return false
+ }
+
+ filterBytes, ok := filterUint16VectorInput(filterData, filterOffset,
length)
+ if !ok {
+ return false
+ }
+
+ _filter_uint16_neon(
+ unsafe.Pointer(unsafe.SliceData(values)),
+ unsafe.Pointer(unsafe.SliceData(filterBytes)),
+ unsafe.Pointer(unsafe.SliceData(output)),
+ unsafe.Pointer(unsafe.SliceData(filterUint16NeonTables[:])),
+ length,
+ )
+ return true
+}
diff --git a/arrow/compute/internal/kernels/filter_uint16_neon_arm64.s
b/arrow/compute/internal/kernels/filter_uint16_neon_arm64.s
new file mode 100644
index 00000000..4b956b16
--- /dev/null
+++ b/arrow/compute/internal/kernels/filter_uint16_neon_arm64.s
@@ -0,0 +1,120 @@
+// 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"
+
+// Each filter byte selects up to eight uint16 values. TBL compacts one
+// sixteen-byte input vector, and stores write exactly the selected bytes.
+// func _filter_uint16_neon(values, filter, output, tables unsafe.Pointer,
length int64)
+TEXT ·_filter_uint16_neon(SB), NOSPLIT|NOFRAME, $0-40
+ MOVD values+0(FP), R0
+ MOVD filter+8(FP), R1
+ MOVD output+16(FP), R2
+ MOVD tables+24(FP), R3
+ MOVD length+32(FP), R4
+
+ LSR $3, R4, R4
+ ADD $4096, R3, R8
+
+filter_uint16_loop:
+ MOVBU (R1), R5
+ ADD $1, R1
+ CBZ R5, filter_uint16_next
+ CMPW $255, R5
+ BEQ filter_uint16_all
+
+ MOVBU (R8)(R5<<0), R6
+ CBZ R6, filter_uint16_next
+
+ LSL $4, R5, R7
+ ADD R3, R7, R7
+ VLD1 (R0), [V0.B16]
+ VLD1 (R7), [V1.B16]
+ VTBL V1.B16, [V0.B16], V2.B16
+
+ CMPW $4, R6
+ BLT filter_uint16_store_small
+ VST1 [V2.H4], (R2)
+ ADD $8, R2
+ SUB $4, R6
+ CBZ R6, filter_uint16_next
+ CMPW $1, R6
+ BEQ filter_uint16_store_high1
+ CMPW $2, R6
+ BEQ filter_uint16_store_high2
+ VMOV V2.H[4], R9
+ MOVH R9, (R2)
+ VMOV V2.H[5], R9
+ MOVH R9, 2(R2)
+ VMOV V2.H[6], R9
+ MOVH R9, 4(R2)
+ ADD $6, R2
+ B filter_uint16_next
+
+filter_uint16_store_high1:
+ VMOV V2.H[4], R9
+ MOVH R9, (R2)
+ ADD $2, R2
+ B filter_uint16_next
+
+filter_uint16_store_high2:
+ VMOV V2.H[4], R9
+ MOVH R9, (R2)
+ VMOV V2.H[5], R9
+ MOVH R9, 2(R2)
+ ADD $4, R2
+ B filter_uint16_next
+
+filter_uint16_store_small:
+ CMPW $1, R6
+ BEQ filter_uint16_store1
+ CMPW $2, R6
+ BEQ filter_uint16_store2
+ VMOV V2.H[0], R9
+ MOVH R9, (R2)
+ VMOV V2.H[1], R9
+ MOVH R9, 2(R2)
+ VMOV V2.H[2], R9
+ MOVH R9, 4(R2)
+ ADD $6, R2
+ B filter_uint16_next
+
+filter_uint16_store1:
+ VMOV V2.H[0], R9
+ MOVH R9, (R2)
+ ADD $2, R2
+ B filter_uint16_next
+
+filter_uint16_store2:
+ VMOV V2.H[0], R9
+ MOVH R9, (R2)
+ VMOV V2.H[1], R9
+ MOVH R9, 2(R2)
+ ADD $4, R2
+ B filter_uint16_next
+
+filter_uint16_all:
+ VLD1 (R0), [V0.B16]
+ VST1 [V0.B16], (R2)
+ ADD $16, R2
+
+filter_uint16_next:
+ ADD $16, R0
+ SUBS $1, R4, R4
+ BNE filter_uint16_loop
+ RET
diff --git a/arrow/compute/internal/kernels/filter_uint16_neon_arm64_test.go
b/arrow/compute/internal/kernels/filter_uint16_neon_arm64_test.go
new file mode 100644
index 00000000..e46a0f10
--- /dev/null
+++ b/arrow/compute/internal/kernels/filter_uint16_neon_arm64_test.go
@@ -0,0 +1,53 @@
+// 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 (
+ "testing"
+
+ "golang.org/x/sys/cpu"
+)
+
+func TestFilterUint16NeonAllMasks(t *testing.T) {
+ if !cpu.ARM64.HasASIMD {
+ t.Skip("ARM64 SIMD is not available")
+ }
+ testFilterUint16AllMasks(t, filterUint16Neon)
+}
+
+func TestFilterUint16NeonAlignedOffset(t *testing.T) {
+ if !cpu.ARM64.HasASIMD {
+ t.Skip("ARM64 SIMD is not available")
+ }
+ testFilterUint16AlignedOffset(t, filterUint16Neon)
+}
+
+func TestFilterUint16NeonUnalignedValues(t *testing.T) {
+ if !cpu.ARM64.HasASIMD {
+ t.Skip("ARM64 SIMD is not available")
+ }
+ testFilterUint16UnalignedValues(t, filterUint16Neon)
+}
+
+func TestFilterUint16NeonGuards(t *testing.T) {
+ if !cpu.ARM64.HasASIMD {
+ t.Skip("ARM64 SIMD is not available")
+ }
+ testFilterUint16Guards(t, filterUint16Neon)
+}
diff --git a/arrow/compute/internal/kernels/filter_uint16_neon_noasm.go
b/arrow/compute/internal/kernels/filter_uint16_neon_noasm.go
new file mode 100644
index 00000000..265cf13d
--- /dev/null
+++ b/arrow/compute/internal/kernels/filter_uint16_neon_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 || !arm64 || appengine)
+
+package kernels
+
+func filterUint16Neon([]uint16, []uint16, []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 8ea29377..174148f2 100644
--- a/arrow/compute/internal/kernels/vector_selection.go
+++ b/arrow/compute/internal/kernels/vector_selection.go
@@ -495,6 +495,14 @@ func PrimitiveFilter(ctx *exec.KernelCtx, batch
*exec.ExecSpan, out *exec.ExecRe
return nil
}
}
+ if bitWidth == 16 && values.Nulls == 0 && filter.Nulls == 0 {
+ valuesData := exec.GetSpanValues[uint16](values, 1)
+ outData := exec.GetSpanValues[uint16](out, 1)
+ if filterUint16Avx2(valuesData, outData, filter.Buffers[1].Buf,
filter.Offset, values.Len) ||
+ filterUint16Neon(valuesData, outData,
filter.Buffers[1].Buf, filter.Offset, values.Len) {
+ return nil
+ }
+ }
if bitWidth == 64 && values.Nulls == 0 && filter.Nulls == 0 {
valuesData := exec.GetSpanValues[uint64](values, 1)
outData := exec.GetSpanValues[uint64](out, 1)
diff --git a/arrow/compute/internal/kernels/vector_selection_test.go
b/arrow/compute/internal/kernels/vector_selection_test.go
index f4c69338..0721f6b7 100644
--- a/arrow/compute/internal/kernels/vector_selection_test.go
+++ b/arrow/compute/internal/kernels/vector_selection_test.go
@@ -20,6 +20,7 @@
package kernels
import (
+ "context"
"testing"
"github.com/apache/arrow-go/v18/arrow"
@@ -112,6 +113,61 @@ func assertTakeIndices[T arrow.IntType | arrow.UintType](t
*testing.T, data arro
require.Equal(t, nulls, data.NullN())
}
+func TestPrimitiveFilterInt16MixedMaskOffsets(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ const (
+ offset = 8
+ length = 128
+ )
+ valuesBuilder := array.NewInt16Builder(mem)
+ valuesBuilder.Reserve(offset + length)
+ for i := 0; i < offset+length; i++ {
+ valuesBuilder.Append(int16(uint16(i*997 + 0x8123)))
+ }
+ valuesBase := valuesBuilder.NewInt16Array()
+ valuesBuilder.Release()
+ values := array.NewSlice(valuesBase, offset, offset+length)
+ valuesBase.Release()
+ defer values.Release()
+
+ filterBuilder := array.NewBooleanBuilder(mem)
+ filterBuilder.Reserve(offset + length)
+ for i := 0; i < offset+length; i++ {
+ filterBuilder.Append(i%5 == 0 || i%5 == 2)
+ }
+ filterBase := filterBuilder.NewBooleanArray()
+ filterBuilder.Release()
+ filter := array.NewSlice(filterBase, offset, offset+length)
+ filterBase.Release()
+ defer filter.Release()
+
+ var batch exec.ExecSpan
+ batch.Len = length
+ batch.Values = make([]exec.ExecValue, 2)
+ batch.Values[0].Array.SetMembers(values.Data())
+ batch.Values[1].Array.SetMembers(filter.Data())
+ ctx := &exec.KernelCtx{
+ Ctx: exec.WithAllocator(context.Background(), mem),
+ State: FilterState{NullSelection: DropNulls},
+ }
+ var out exec.ExecResult
+ require.NoError(t, PrimitiveFilter(ctx, &batch, &out))
+ defer out.Release()
+
+ got := exec.GetSpanValues[int16](&out, 1)
+ valuesData := exec.GetSpanValues[int16](&batch.Values[0].Array, 1)
+ position := 0
+ for i := offset; i < offset+length; i++ {
+ if i%5 == 0 || i%5 == 2 {
+ require.Equal(t, valuesData[i-offset], got[position])
+ position++
+ }
+ }
+ require.Equal(t, position, int(out.Len))
+}
+
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())
diff --git a/arrow/compute/vector_selection_bench_test.go
b/arrow/compute/vector_selection_bench_test.go
index 8bef0201..b0a5c930 100644
--- a/arrow/compute/vector_selection_bench_test.go
+++ b/arrow/compute/vector_selection_bench_test.go
@@ -14,7 +14,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
-//go:build go1.18
+//go:build go1.24
package compute_test
@@ -35,6 +35,10 @@ func BenchmarkFilterInt32MixedMasks(b *testing.B) {
benchmarkFilter32MixedMasks(b, makeFilterInt32BenchmarkInput)
}
+func BenchmarkFilterInt16MixedMasks(b *testing.B) {
+ benchmarkFilterMixedMasks(b, makeFilterInt16BenchmarkInput, 2)
+}
+
func BenchmarkFilterFloat32MixedMasks(b *testing.B) {
benchmarkFilter32MixedMasks(b, makeFilterFloat32BenchmarkInput)
}
@@ -66,6 +70,12 @@ func benchmarkFilterMixedMasks(b *testing.B,
name string
selected func(int) bool
}{
+ {name: "mixed-prefix-long-runs", selected: func(i int) bool {
+ if i < 32 {
+ return i%2 == 0
+ }
+ return (i/4096)%2 == 0
+ }},
{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) }},
@@ -112,7 +122,6 @@ func BenchmarkFilterInt32FallbackControls(b *testing.B) {
}
for _, control := range controls {
- control := control
b.Run(control.name, func(b *testing.B) {
values, filter :=
makeFilterInt32BenchmarkInputWithOffset(b, control.size, control.offset,
func(i int) bool { return i%2 == 0 })
@@ -122,7 +131,40 @@ func BenchmarkFilterInt32FallbackControls(b *testing.B) {
b.ReportAllocs()
b.SetBytes(int64(control.size * 4))
b.ResetTimer()
- for i := 0; i < b.N; i++ {
+ for b.Loop() {
+ result, err :=
compute.FilterArray(context.Background(), values, filter,
*compute.DefaultFilterOptions())
+ if err != nil {
+ b.Fatal(err)
+ }
+ benchmarkFilterOutputLength = result.Len()
+ result.Release()
+ }
+ })
+ }
+}
+
+func BenchmarkFilterInt16FallbackControls(b *testing.B) {
+ controls := []struct {
+ name string
+ size int
+ offset int64
+ }{
+ {name: "unaligned-filter-offset", size: 1 << 16, offset: 3},
+ {name: "short", size: 56},
+ {name: "non-multiple-of-eight", size: 1<<16 + 1},
+ }
+
+ for _, control := range controls {
+ b.Run(control.name, func(b *testing.B) {
+ values, filter :=
makeFilterInt16BenchmarkInputWithOffset(b, control.size, control.offset,
+ func(i int) bool { return i%2 == 0 })
+ defer values.Release()
+ defer filter.Release()
+
+ b.ReportAllocs()
+ b.SetBytes(int64(control.size * 2))
+ b.ResetTimer()
+ for b.Loop() {
result, err :=
compute.FilterArray(context.Background(), values, filter,
*compute.DefaultFilterOptions())
if err != nil {
b.Fatal(err)
@@ -134,6 +176,36 @@ func BenchmarkFilterInt32FallbackControls(b *testing.B) {
}
}
+func makeFilterInt16BenchmarkInput(b *testing.B, size int, selected func(int)
bool) (arrow.Array, arrow.Array) {
+ return makeFilterInt16BenchmarkInputWithOffset(b, size, 0, selected)
+}
+
+func makeFilterInt16BenchmarkInputWithOffset(b *testing.B, size int, offset
int64, selected func(int) bool) (arrow.Array, arrow.Array) {
+ b.Helper()
+ mem := memory.DefaultAllocator
+
+ valuesBuilder := array.NewInt16Builder(mem)
+ valuesBuilder.Reserve(size + int(offset))
+ for i := 0; i < size+int(offset); i++ {
+ valuesBuilder.Append(int16(i))
+ }
+ valuesBase := valuesBuilder.NewInt16Array()
+ valuesBuilder.Release()
+ values := array.NewSlice(valuesBase, offset, offset+int64(size))
+ valuesBase.Release()
+
+ filterBuilder := array.NewBooleanBuilder(mem)
+ filterBuilder.Reserve(size + int(offset))
+ for i := 0; i < size+int(offset); i++ {
+ filterBuilder.Append(selected(i))
+ }
+ filterBase := filterBuilder.NewBooleanArray()
+ filterBuilder.Release()
+ filter := array.NewSlice(filterBase, offset, offset+int64(size))
+ filterBase.Release()
+ return values, filter
+}
+
func makeFilterInt32BenchmarkInput(b *testing.B, size int, selected func(int)
bool) (arrow.Array, arrow.Array) {
return makeFilterInt32BenchmarkInputWithOffset(b, size, 0, selected)
}
diff --git a/arrow/compute/vector_selection_test.go
b/arrow/compute/vector_selection_test.go
index 0d9d9f04..1e9233e6 100644
--- a/arrow/compute/vector_selection_test.go
+++ b/arrow/compute/vector_selection_test.go
@@ -2079,6 +2079,53 @@ func TestFilterInt32MixedMaskOffsets(t *testing.T) {
}
}
+func TestFilterInt16MixedMaskOffsets(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.NewInt16Builder(mem)
+ valuesBuilder.Reserve(int(offset) + length)
+ for i := int64(0); i < offset+length; i++ {
+ valuesBuilder.Append(int16(uint16(i*977 +
0x8011)))
+ }
+ valuesBase := valuesBuilder.NewInt16Array()
+ 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%5 == 0 || i%5 == 2)
+ }
+ filterBase := filterBuilder.NewBooleanArray()
+ filterBuilder.Release()
+ filter := array.NewSlice(filterBase, offset,
offset+length)
+ filterBase.Release()
+ defer filter.Release()
+
+ expectedBuilder := array.NewInt16Builder(mem)
+ for i := offset; i < offset+length; i++ {
+ if i%5 == 0 || i%5 == 2 {
+
expectedBuilder.Append(int16(uint16(i*977 + 0x8011)))
+ }
+ }
+ expected := expectedBuilder.NewInt16Array()
+ 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).
