From eafc06d9e4f3a0d235f05801e8247e2afa6890be Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Sat, 29 Aug 2026 14:27:44 +0200 Subject: [PATCH 1/2] perf(parquet): SIMD BYTE_STREAM_SPLIT encoding --- .../internal/encoding/byte_stream_split.go | 9 +- .../encoding/byte_stream_split_amd64.go | 24 ++ .../encoding/byte_stream_split_arm64.go | 24 ++ .../byte_stream_split_encode_amd64_test.go | 39 ++++ .../byte_stream_split_encode_arm64_test.go | 39 ++++ .../byte_stream_split_encode_avx2_amd64.s | 205 ++++++++++++++++++ .../byte_stream_split_encode_neon_arm64.s | 161 ++++++++++++++ .../encoding/byte_stream_split_encode_test.go | 103 +++++++++ .../encoding/fixed_len_byte_array_encoder.go | 4 +- 9 files changed, 604 insertions(+), 4 deletions(-) create mode 100644 parquet/internal/encoding/byte_stream_split_encode_amd64_test.go create mode 100644 parquet/internal/encoding/byte_stream_split_encode_arm64_test.go create mode 100644 parquet/internal/encoding/byte_stream_split_encode_avx2_amd64.s create mode 100644 parquet/internal/encoding/byte_stream_split_encode_neon_arm64.s create mode 100644 parquet/internal/encoding/byte_stream_split_encode_test.go diff --git a/parquet/internal/encoding/byte_stream_split.go b/parquet/internal/encoding/byte_stream_split.go index 1854c6c4e..970e0f88c 100644 --- a/parquet/internal/encoding/byte_stream_split.go +++ b/parquet/internal/encoding/byte_stream_split.go @@ -88,6 +88,11 @@ func encodeByteStreamSplitWidth8(data []byte, in []byte) { } } +var ( + encodeByteStreamSplitWidth4Impl = encodeByteStreamSplitWidth4 + encodeByteStreamSplitWidth8Impl = encodeByteStreamSplitWidth8 +) + func releaseBufferToPool(pooled *PooledBufferWriter) { buf := pooled.buf memory.Set(buf.Buf(), 0) @@ -126,9 +131,9 @@ func (enc *byteStreamSplitEncoder[T]) FlushValues() (Buffer, error) { var z T switch any(z).(type) { case int32, float32: - encodeByteStreamSplitWidth4(enc.flushBuffer.Bytes(), in.Bytes()) + encodeByteStreamSplitWidth4Impl(enc.flushBuffer.Bytes(), in.Bytes()) case int64, float64: - encodeByteStreamSplitWidth8(enc.flushBuffer.Bytes(), in.Bytes()) + encodeByteStreamSplitWidth8Impl(enc.flushBuffer.Bytes(), in.Bytes()) } return enc.flushBuffer.Finish(), nil diff --git a/parquet/internal/encoding/byte_stream_split_amd64.go b/parquet/internal/encoding/byte_stream_split_amd64.go index b6aad9d7a..8c8342b49 100644 --- a/parquet/internal/encoding/byte_stream_split_amd64.go +++ b/parquet/internal/encoding/byte_stream_split_amd64.go @@ -29,9 +29,33 @@ func init() { if cpu.X86.HasAVX2 { decodeByteStreamSplitBatchWidth4InByteOrder = decodeByteStreamSplitBatchWidth4AVX2 decodeByteStreamSplitBatchWidth8InByteOrder = decodeByteStreamSplitBatchWidth8AVX2 + encodeByteStreamSplitWidth4Impl = encodeByteStreamSplitWidth4AVX2 + encodeByteStreamSplitWidth8Impl = encodeByteStreamSplitWidth8AVX2 } } +//go:noescape +func _encodeByteStreamSplitWidth4AVX2(in, out unsafe.Pointer, nValues int) + +//go:noescape +func _encodeByteStreamSplitWidth8AVX2(in, out unsafe.Pointer, nValues int) + +func encodeByteStreamSplitWidth4AVX2(data, in []byte) { + if len(in) == 0 { + return + } + debug.Assert(len(data) >= len(in), "not enough space in destination buffer for encoding") + _encodeByteStreamSplitWidth4AVX2(unsafe.Pointer(&in[0]), unsafe.Pointer(&data[0]), len(in)/4) +} + +func encodeByteStreamSplitWidth8AVX2(data, in []byte) { + if len(in) == 0 { + return + } + debug.Assert(len(data) >= len(in), "not enough space in destination buffer for encoding") + _encodeByteStreamSplitWidth8AVX2(unsafe.Pointer(&in[0]), unsafe.Pointer(&data[0]), len(in)/8) +} + //go:noescape func _decodeByteStreamSplitWidth4AVX2(data, out unsafe.Pointer, nValues, stride int) diff --git a/parquet/internal/encoding/byte_stream_split_arm64.go b/parquet/internal/encoding/byte_stream_split_arm64.go index 6f638812d..40e139195 100644 --- a/parquet/internal/encoding/byte_stream_split_arm64.go +++ b/parquet/internal/encoding/byte_stream_split_arm64.go @@ -29,9 +29,33 @@ func init() { if cpu.ARM64.HasASIMD { decodeByteStreamSplitBatchWidth4InByteOrder = decodeByteStreamSplitBatchWidth4NEON decodeByteStreamSplitBatchWidth8InByteOrder = decodeByteStreamSplitBatchWidth8NEON + encodeByteStreamSplitWidth4Impl = encodeByteStreamSplitWidth4NEON + encodeByteStreamSplitWidth8Impl = encodeByteStreamSplitWidth8NEON } } +//go:noescape +func _encodeByteStreamSplitWidth4NEON(in, out unsafe.Pointer, nValues int) + +//go:noescape +func _encodeByteStreamSplitWidth8NEON(in, out unsafe.Pointer, nValues int) + +func encodeByteStreamSplitWidth4NEON(data, in []byte) { + if len(in) == 0 { + return + } + debug.Assert(len(data) >= len(in), "not enough space in destination buffer for encoding") + _encodeByteStreamSplitWidth4NEON(unsafe.Pointer(&in[0]), unsafe.Pointer(&data[0]), len(in)/4) +} + +func encodeByteStreamSplitWidth8NEON(data, in []byte) { + if len(in) == 0 { + return + } + debug.Assert(len(data) >= len(in), "not enough space in destination buffer for encoding") + _encodeByteStreamSplitWidth8NEON(unsafe.Pointer(&in[0]), unsafe.Pointer(&data[0]), len(in)/8) +} + //go:noescape func _decodeByteStreamSplitWidth4NEON(data, out unsafe.Pointer, nValues, stride int) diff --git a/parquet/internal/encoding/byte_stream_split_encode_amd64_test.go b/parquet/internal/encoding/byte_stream_split_encode_amd64_test.go new file mode 100644 index 000000000..0a38f627a --- /dev/null +++ b/parquet/internal/encoding/byte_stream_split_encode_amd64_test.go @@ -0,0 +1,39 @@ +// 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 !noasm && !appengine + +package encoding + +import ( + "testing" + + "golang.org/x/sys/cpu" +) + +func TestEncodeByteStreamSplitWidth4AVX2(t *testing.T) { + if !cpu.X86.HasAVX2 { + t.Skip("AVX2 is not available") + } + testEncodeByteStreamSplit(t, 4, encodeByteStreamSplitWidth4AVX2) +} + +func TestEncodeByteStreamSplitWidth8AVX2(t *testing.T) { + if !cpu.X86.HasAVX2 { + t.Skip("AVX2 is not available") + } + testEncodeByteStreamSplit(t, 8, encodeByteStreamSplitWidth8AVX2) +} diff --git a/parquet/internal/encoding/byte_stream_split_encode_arm64_test.go b/parquet/internal/encoding/byte_stream_split_encode_arm64_test.go new file mode 100644 index 000000000..b53aa25ca --- /dev/null +++ b/parquet/internal/encoding/byte_stream_split_encode_arm64_test.go @@ -0,0 +1,39 @@ +// 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 !noasm && !appengine + +package encoding + +import ( + "testing" + + "golang.org/x/sys/cpu" +) + +func TestEncodeByteStreamSplitWidth4NEON(t *testing.T) { + if !cpu.ARM64.HasASIMD { + t.Skip("ASIMD is not available") + } + testEncodeByteStreamSplit(t, 4, encodeByteStreamSplitWidth4NEON) +} + +func TestEncodeByteStreamSplitWidth8NEON(t *testing.T) { + if !cpu.ARM64.HasASIMD { + t.Skip("ASIMD is not available") + } + testEncodeByteStreamSplit(t, 8, encodeByteStreamSplitWidth8NEON) +} diff --git a/parquet/internal/encoding/byte_stream_split_encode_avx2_amd64.s b/parquet/internal/encoding/byte_stream_split_encode_avx2_amd64.s new file mode 100644 index 000000000..7d654e925 --- /dev/null +++ b/parquet/internal/encoding/byte_stream_split_encode_avx2_amd64.s @@ -0,0 +1,205 @@ +//+build !noasm !appengine + +// AVX2 implementation of BYTE_STREAM_SPLIT encoding. + +#include "textflag.h" + +// func _encodeByteStreamSplitWidth4AVX2(in, out unsafe.Pointer, nValues int) +// +// Transposes eight 4-byte values at a time. The unpack stages produce two +// 8-byte streams per output register. +TEXT ·_encodeByteStreamSplitWidth4AVX2(SB), NOSPLIT, $0-24 + MOVQ in+0(FP), SI + MOVQ out+8(FP), DI + MOVQ nValues+16(FP), CX + TESTQ CX, CX + JZ encode_w4_done + + MOVQ CX, AX + SHRQ $3, AX + LEAQ (DI)(CX*1), R8 + LEAQ (DI)(CX*2), R9 + LEAQ (R8)(CX*2), R10 + MOVQ AX, R11 + SHLQ $3, R11 + SUBQ R11, CX + + TESTQ AX, AX + JZ encode_w4_tail + +encode_w4_vector: + VMOVD 0(SI), X0 + VMOVD 4(SI), X1 + VMOVD 8(SI), X2 + VMOVD 12(SI), X3 + VMOVD 16(SI), X4 + VMOVD 20(SI), X5 + VMOVD 24(SI), X6 + VMOVD 28(SI), X7 + + VPUNPCKLBW X1, X0, X8 + VPUNPCKLBW X3, X2, X9 + VPUNPCKLBW X5, X4, X10 + VPUNPCKLBW X7, X6, X11 + + VPUNPCKLWD X9, X8, X12 + VPUNPCKLWD X11, X10, X13 + + VPUNPCKLDQ X13, X12, X0 + VPUNPCKHDQ X13, X12, X1 + + VMOVQ X0, (DI) + VPSRLDQ $8, X0, X0 + VMOVQ X0, (R8) + VMOVQ X1, (R9) + VPSRLDQ $8, X1, X1 + VMOVQ X1, (R10) + + ADDQ $32, SI + ADDQ $8, DI + ADDQ $8, R8 + ADDQ $8, R9 + ADDQ $8, R10 + DECQ AX + JNZ encode_w4_vector + +encode_w4_tail: + TESTQ CX, CX + JZ encode_w4_done + +encode_w4_tail_loop: + MOVBQZX 0(SI), BX + MOVB BX, (DI) + MOVBQZX 1(SI), BX + MOVB BX, (R8) + MOVBQZX 2(SI), BX + MOVB BX, (R9) + MOVBQZX 3(SI), BX + MOVB BX, (R10) + + ADDQ $4, SI + INCQ DI + INCQ R8 + INCQ R9 + INCQ R10 + DECQ CX + JNZ encode_w4_tail_loop + +encode_w4_done: + VZEROUPPER + RET + +// func _encodeByteStreamSplitWidth8AVX2(in, out unsafe.Pointer, nValues int) +// +// Transposes eight 8-byte values at a time. The unpack stages produce two +// 8-byte streams per output register. +TEXT ·_encodeByteStreamSplitWidth8AVX2(SB), NOSPLIT, $0-24 + MOVQ in+0(FP), SI + MOVQ out+8(FP), DI + MOVQ nValues+16(FP), CX + TESTQ CX, CX + JZ encode_w8_done + + MOVQ CX, AX + SHRQ $3, AX + LEAQ (DI)(CX*1), R8 + LEAQ (DI)(CX*2), R9 + LEAQ (R8)(CX*2), R10 + LEAQ (DI)(CX*4), R11 + LEAQ (R8)(CX*4), R12 + LEAQ (R9)(CX*4), R13 + LEAQ (R10)(CX*4), R14 + MOVQ AX, R15 + SHLQ $3, R15 + SUBQ R15, CX + + TESTQ AX, AX + JZ encode_w8_tail + +encode_w8_vector: + VMOVQ 0(SI), X0 + VMOVQ 8(SI), X1 + VMOVQ 16(SI), X2 + VMOVQ 24(SI), X3 + VMOVQ 32(SI), X4 + VMOVQ 40(SI), X5 + VMOVQ 48(SI), X6 + VMOVQ 56(SI), X7 + + VPUNPCKLBW X1, X0, X8 + VPUNPCKLBW X3, X2, X9 + VPUNPCKLBW X5, X4, X10 + VPUNPCKLBW X7, X6, X11 + + VPUNPCKLWD X9, X8, X12 + VPUNPCKHWD X9, X8, X13 + VPUNPCKLWD X11, X10, X14 + VPUNPCKHWD X11, X10, X15 + + VPUNPCKLDQ X14, X12, X0 + VPUNPCKHDQ X14, X12, X1 + VPUNPCKLDQ X15, X13, X2 + VPUNPCKHDQ X15, X13, X3 + + VMOVQ X0, (DI) + VPSRLDQ $8, X0, X0 + VMOVQ X0, (R8) + VMOVQ X1, (R9) + VPSRLDQ $8, X1, X1 + VMOVQ X1, (R10) + VMOVQ X2, (R11) + VPSRLDQ $8, X2, X2 + VMOVQ X2, (R12) + VMOVQ X3, (R13) + VPSRLDQ $8, X3, X3 + VMOVQ X3, (R14) + + ADDQ $64, SI + ADDQ $8, DI + ADDQ $8, R8 + ADDQ $8, R9 + ADDQ $8, R10 + ADDQ $8, R11 + ADDQ $8, R12 + ADDQ $8, R13 + ADDQ $8, R14 + DECQ AX + JNZ encode_w8_vector + +encode_w8_tail: + TESTQ CX, CX + JZ encode_w8_done + +encode_w8_tail_loop: + MOVBQZX 0(SI), BX + MOVB BX, (DI) + MOVBQZX 1(SI), BX + MOVB BX, (R8) + MOVBQZX 2(SI), BX + MOVB BX, (R9) + MOVBQZX 3(SI), BX + MOVB BX, (R10) + MOVBQZX 4(SI), BX + MOVB BX, (R11) + MOVBQZX 5(SI), BX + MOVB BX, (R12) + MOVBQZX 6(SI), BX + MOVB BX, (R13) + MOVBQZX 7(SI), BX + MOVB BX, (R14) + + ADDQ $8, SI + INCQ DI + INCQ R8 + INCQ R9 + INCQ R10 + INCQ R11 + INCQ R12 + INCQ R13 + INCQ R14 + DECQ CX + JNZ encode_w8_tail_loop + +encode_w8_done: + VZEROUPPER + RET diff --git a/parquet/internal/encoding/byte_stream_split_encode_neon_arm64.s b/parquet/internal/encoding/byte_stream_split_encode_neon_arm64.s new file mode 100644 index 000000000..00420c378 --- /dev/null +++ b/parquet/internal/encoding/byte_stream_split_encode_neon_arm64.s @@ -0,0 +1,161 @@ +//+build !noasm !appengine + +// NEON implementation of BYTE_STREAM_SPLIT encoding. + +#include "textflag.h" + +// func _encodeByteStreamSplitWidth4NEON(in, out unsafe.Pointer, nValues int) +// +// VLD4 deinterleaves sixteen 4-byte values into the four output streams. +TEXT ·_encodeByteStreamSplitWidth4NEON(SB), NOSPLIT, $0-24 + MOVD in+0(FP), R0 + MOVD out+8(FP), R1 + MOVD nValues+16(FP), R2 + CBZ R2, encode_w4_neon_done + + MOVD R2, R3 + LSR $4, R3, R3 + LSL $4, R3, R12 + SUB R12, R2, R13 + + MOVD R1, R4 + ADD R2, R1, R5 + ADD R2, R5, R6 + ADD R2, R6, R7 + + CBZ R3, encode_w4_neon_tail + +encode_w4_neon_vector: + VLD4 (R0), [V0.B16, V1.B16, V2.B16, V3.B16] + VST1 [V0.B16], (R4) + VST1 [V1.B16], (R5) + VST1 [V2.B16], (R6) + VST1 [V3.B16], (R7) + + ADD $64, R0, R0 + ADD $16, R4, R4 + ADD $16, R5, R5 + ADD $16, R6, R6 + ADD $16, R7, R7 + SUB $1, R3, R3 + CBNZ R3, encode_w4_neon_vector + +encode_w4_neon_tail: + CBZ R13, encode_w4_neon_done + +encode_w4_neon_tail_loop: + MOVBU 0(R0), R14 + MOVB R14, (R4) + MOVBU 1(R0), R14 + MOVB R14, (R5) + MOVBU 2(R0), R14 + MOVB R14, (R6) + MOVBU 3(R0), R14 + MOVB R14, (R7) + + ADD $4, R0, R0 + ADD $1, R4, R4 + ADD $1, R5, R5 + ADD $1, R6, R6 + ADD $1, R7, R7 + SUB $1, R13, R13 + CBNZ R13, encode_w4_neon_tail_loop + +encode_w4_neon_done: + RET + +// func _encodeByteStreamSplitWidth8NEON(in, out unsafe.Pointer, nValues int) +// +// VLD4 first separates each 8-byte value into its low and high four-byte +// groups. VUZP then separates those groups into the eight byte streams. +TEXT ·_encodeByteStreamSplitWidth8NEON(SB), NOSPLIT, $0-24 + MOVD in+0(FP), R0 + MOVD out+8(FP), R1 + MOVD nValues+16(FP), R2 + CBZ R2, encode_w8_neon_done + + MOVD R2, R3 + LSR $3, R3, R3 + LSL $3, R3, R12 + SUB R12, R2, R13 + + MOVD R1, R4 + ADD R2, R1, R5 + ADD R2, R5, R6 + ADD R2, R6, R7 + LSL $2, R2, R12 + ADD R12, R1, R8 + ADD R12, R5, R9 + ADD R12, R6, R10 + ADD R12, R7, R11 + + CBZ R3, encode_w8_neon_tail + +encode_w8_neon_vector: + VLD4 (R0), [V0.B16, V1.B16, V2.B16, V3.B16] + + VUZP1 V0.B16, V0.B16, V4.B16 + VUZP2 V0.B16, V0.B16, V8.B16 + VUZP1 V1.B16, V1.B16, V5.B16 + VUZP2 V1.B16, V1.B16, V9.B16 + VUZP1 V2.B16, V2.B16, V6.B16 + VUZP2 V2.B16, V2.B16, V10.B16 + VUZP1 V3.B16, V3.B16, V7.B16 + VUZP2 V3.B16, V3.B16, V11.B16 + + VST1 [V4.B8], (R4) + VST1 [V5.B8], (R5) + VST1 [V6.B8], (R6) + VST1 [V7.B8], (R7) + VST1 [V8.B8], (R8) + VST1 [V9.B8], (R9) + VST1 [V10.B8], (R10) + VST1 [V11.B8], (R11) + + ADD $64, R0, R0 + ADD $8, R4, R4 + ADD $8, R5, R5 + ADD $8, R6, R6 + ADD $8, R7, R7 + ADD $8, R8, R8 + ADD $8, R9, R9 + ADD $8, R10, R10 + ADD $8, R11, R11 + SUB $1, R3, R3 + CBNZ R3, encode_w8_neon_vector + +encode_w8_neon_tail: + CBZ R13, encode_w8_neon_done + +encode_w8_neon_tail_loop: + MOVBU 0(R0), R14 + MOVB R14, (R4) + MOVBU 1(R0), R14 + MOVB R14, (R5) + MOVBU 2(R0), R14 + MOVB R14, (R6) + MOVBU 3(R0), R14 + MOVB R14, (R7) + MOVBU 4(R0), R14 + MOVB R14, (R8) + MOVBU 5(R0), R14 + MOVB R14, (R9) + MOVBU 6(R0), R14 + MOVB R14, (R10) + MOVBU 7(R0), R14 + MOVB R14, (R11) + + ADD $8, R0, R0 + ADD $1, R4, R4 + ADD $1, R5, R5 + ADD $1, R6, R6 + ADD $1, R7, R7 + ADD $1, R8, R8 + ADD $1, R9, R9 + ADD $1, R10, R10 + ADD $1, R11, R11 + SUB $1, R13, R13 + CBNZ R13, encode_w8_neon_tail_loop + +encode_w8_neon_done: + RET diff --git a/parquet/internal/encoding/byte_stream_split_encode_test.go b/parquet/internal/encoding/byte_stream_split_encode_test.go new file mode 100644 index 000000000..7e065fe73 --- /dev/null +++ b/parquet/internal/encoding/byte_stream_split_encode_test.go @@ -0,0 +1,103 @@ +// 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. + +package encoding + +import ( + "bytes" + "fmt" + "testing" +) + +func TestEncodeByteStreamSplitWidth4(t *testing.T) { + testEncodeByteStreamSplit(t, 4, encodeByteStreamSplitWidth4Impl) +} + +func TestEncodeByteStreamSplitWidth8(t *testing.T) { + testEncodeByteStreamSplit(t, 8, encodeByteStreamSplitWidth8Impl) +} + +func testEncodeByteStreamSplit(t *testing.T, width int, implementation func([]byte, []byte)) { + for _, nValues := range []int{0, 1, 2, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129} { + t.Run(fmt.Sprintf("nValues=%d", nValues), func(t *testing.T) { + in := make([]byte, nValues*width) + for i := range in { + in[i] = byte((i*37 + width) ^ (i >> 3)) + } + + want := make([]byte, len(in)) + got := make([]byte, len(in)) + switch width { + case 4: + encodeByteStreamSplitWidth4(want, in) + case 8: + encodeByteStreamSplitWidth8(want, in) + } + implementation(got, in) + if !bytes.Equal(got, want) { + t.Fatalf("encoded output mismatch: got %x, want %x", got, want) + } + }) + } +} + +func BenchmarkEncodeByteStreamSplitWidth4(b *testing.B) { + benchmarkEncodeByteStreamSplit(b, 4) +} + +func BenchmarkEncodeByteStreamSplitWidth8(b *testing.B) { + benchmarkEncodeByteStreamSplit(b, 8) +} + +func benchmarkEncodeByteStreamSplit(b *testing.B, width int) { + for _, nValues := range []int{8, 1024, 65536} { + b.Run(fmt.Sprintf("nValues=%d", nValues), func(b *testing.B) { + in := make([]byte, nValues*width) + for i := range in { + in[i] = byte((i*37 + width) ^ (i >> 3)) + } + out := make([]byte, len(in)) + + b.Run("scalar", func(b *testing.B) { + b.SetBytes(int64(len(in))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + switch width { + case 4: + encodeByteStreamSplitWidth4(out, in) + case 8: + encodeByteStreamSplitWidth8(out, in) + } + } + }) + + b.Run("dispatch", func(b *testing.B) { + b.SetBytes(int64(len(in))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + switch width { + case 4: + encodeByteStreamSplitWidth4Impl(out, in) + case 8: + encodeByteStreamSplitWidth8Impl(out, in) + } + } + }) + }) + } +} diff --git a/parquet/internal/encoding/fixed_len_byte_array_encoder.go b/parquet/internal/encoding/fixed_len_byte_array_encoder.go index 190802eec..e4c3c6504 100644 --- a/parquet/internal/encoding/fixed_len_byte_array_encoder.go +++ b/parquet/internal/encoding/fixed_len_byte_array_encoder.go @@ -110,9 +110,9 @@ func (enc *ByteStreamSplitFixedLenByteArrayEncoder) FlushValues() (Buffer, error case 2: encodeByteStreamSplitWidth2(enc.flushBuffer.Bytes(), in.Bytes()) case 4: - encodeByteStreamSplitWidth4(enc.flushBuffer.Bytes(), in.Bytes()) + encodeByteStreamSplitWidth4Impl(enc.flushBuffer.Bytes(), in.Bytes()) case 8: - encodeByteStreamSplitWidth8(enc.flushBuffer.Bytes(), in.Bytes()) + encodeByteStreamSplitWidth8Impl(enc.flushBuffer.Bytes(), in.Bytes()) default: encodeByteStreamSplit(enc.flushBuffer.Bytes(), in.Bytes(), enc.typeLen) } From 8ecb995814bba7568e3b74f2d2313b6c93f1fab7 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Sun, 30 Aug 2026 22:42:12 +0200 Subject: [PATCH 2/2] fix(parquet): honor noasm for byte-stream split encoding --- .../byte_stream_split_encode_avx2_amd64.s | 19 ++++++++++++++++++- .../byte_stream_split_encode_neon_arm64.s | 19 ++++++++++++++++++- .../encoding/byte_stream_split_encode_test.go | 15 ++++++++++++--- 3 files changed, 48 insertions(+), 5 deletions(-) diff --git a/parquet/internal/encoding/byte_stream_split_encode_avx2_amd64.s b/parquet/internal/encoding/byte_stream_split_encode_avx2_amd64.s index 7d654e925..7c3e14438 100644 --- a/parquet/internal/encoding/byte_stream_split_encode_avx2_amd64.s +++ b/parquet/internal/encoding/byte_stream_split_encode_avx2_amd64.s @@ -1,4 +1,21 @@ -//+build !noasm !appengine +// 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 !noasm && !appengine +// +build !noasm,!appengine // AVX2 implementation of BYTE_STREAM_SPLIT encoding. diff --git a/parquet/internal/encoding/byte_stream_split_encode_neon_arm64.s b/parquet/internal/encoding/byte_stream_split_encode_neon_arm64.s index 00420c378..a1e6986da 100644 --- a/parquet/internal/encoding/byte_stream_split_encode_neon_arm64.s +++ b/parquet/internal/encoding/byte_stream_split_encode_neon_arm64.s @@ -1,4 +1,21 @@ -//+build !noasm !appengine +// 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 !noasm && !appengine +// +build !noasm,!appengine // NEON implementation of BYTE_STREAM_SPLIT encoding. diff --git a/parquet/internal/encoding/byte_stream_split_encode_test.go b/parquet/internal/encoding/byte_stream_split_encode_test.go index 7e065fe73..23b545b7a 100644 --- a/parquet/internal/encoding/byte_stream_split_encode_test.go +++ b/parquet/internal/encoding/byte_stream_split_encode_test.go @@ -31,15 +31,18 @@ func TestEncodeByteStreamSplitWidth8(t *testing.T) { } func testEncodeByteStreamSplit(t *testing.T, width int, implementation func([]byte, []byte)) { - for _, nValues := range []int{0, 1, 2, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129} { + for _, nValues := range []int{0, 1, 2, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129, 1023, 1024, 1025} { t.Run(fmt.Sprintf("nValues=%d", nValues), func(t *testing.T) { - in := make([]byte, nValues*width) + input := bytes.Repeat([]byte{0xa5}, nValues*width+2) + in := input[1 : len(input)-1] for i := range in { in[i] = byte((i*37 + width) ^ (i >> 3)) } + original := bytes.Clone(input) want := make([]byte, len(in)) - got := make([]byte, len(in)) + output := bytes.Repeat([]byte{0x5a}, len(in)+2) + got := output[1 : len(output)-1] switch width { case 4: encodeByteStreamSplitWidth4(want, in) @@ -50,6 +53,12 @@ func testEncodeByteStreamSplit(t *testing.T, width int, implementation func([]by if !bytes.Equal(got, want) { t.Fatalf("encoded output mismatch: got %x, want %x", got, want) } + if output[0] != 0x5a || output[len(output)-1] != 0x5a { + t.Fatal("encoding modified bytes outside the output slice") + } + if !bytes.Equal(input, original) { + t.Fatal("encoding modified the input") + } }) } }