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 85cf9f68 perf(parquet): decode levels directly into int16 batches
(#1176)
85cf9f68 is described below
commit 85cf9f685b7a54aff1f5b7103c8885e3b5e5f287
Author: Minh Vu <[email protected]>
AuthorDate: Fri Aug 28 22:29:57 2026 +0200
perf(parquet): decode levels directly into int16 batches (#1176)
### What changed
- decode RLE and bit-packed levels directly into `[]int16`
- count max definition levels during decoding
- remove the intermediate `uint64` level batch and its second conversion
pass
- add benchmarks for all-defined, mostly-defined, and alternating levels
- check partial value counts when a literal run is truncated
### Benchmarks
Apple M1 Pro, `-benchtime=300ms -count=5`:
| RLE workload | Before | After | Change |
| --- | ---: | ---: | ---: |
| 1,024 all defined | 2,137 ns/op | 1,495 ns/op | -30% |
| 1,024 mostly defined | 5,740 ns/op | 5,329 ns/op | -7% |
| 65,536 all defined | 61,213 ns/op | 24,601 ns/op | -60% |
| 65,536 alternating | 120,994 ns/op | 89,543 ns/op | -26% |
Allocations were unchanged.
---
parquet/internal/encoding/levels.go | 15 +--
parquet/internal/encoding/levels_benchmark_test.go | 105 +++++++++++++++++++++
parquet/internal/encoding/levels_test.go | 23 ++++-
parquet/internal/utils/bit_reader.go | 65 +++++++++++++
parquet/internal/utils/rle.go | 42 +++++++++
5 files changed, 235 insertions(+), 15 deletions(-)
diff --git a/parquet/internal/encoding/levels.go
b/parquet/internal/encoding/levels.go
index 3da9a575..c823b396 100644
--- a/parquet/internal/encoding/levels.go
+++ b/parquet/internal/encoding/levels.go
@@ -247,10 +247,10 @@ func (l *LevelDecoder) SetDataV2(nbytes int32, maxLvl
int16, nbuffered int, data
// values exist to be read, along with any decoding error.
func (l *LevelDecoder) Decode(levels []int16) (int, int64, error) {
var (
- buf [1024]uint64
totaldecoded int
decoded int
valsToRead int64
+ maxCount int64
err error
)
@@ -259,21 +259,14 @@ func (l *LevelDecoder) Decode(levels []int16) (int,
int64, error) {
batch := shared_utils.Min(1024, n)
switch l.encoding {
case format.Encoding_RLE:
- decoded, err = l.rle.GetBatch(buf[:batch])
+ decoded, maxCount, err =
l.rle.GetBatchLevels(levels[:batch], l.maxLvl)
case format.Encoding_BIT_PACKED:
- decoded, err = l.bit.GetBatch(uint(l.bitWidth),
buf[:batch])
+ decoded, maxCount, err =
l.bit.GetBatchLevels(uint(l.bitWidth), levels[:batch], l.maxLvl)
}
l.remaining -= decoded
totaldecoded += decoded
n -= int64(decoded)
-
- for idx, val := range buf[:decoded] {
- lvl := int16(val)
- levels[idx] = lvl
- if lvl == l.maxLvl {
- valsToRead++
- }
- }
+ valsToRead += maxCount
levels = levels[decoded:]
if err != nil {
return totaldecoded, valsToRead, err
diff --git a/parquet/internal/encoding/levels_benchmark_test.go
b/parquet/internal/encoding/levels_benchmark_test.go
new file mode 100644
index 00000000..f206dfcd
--- /dev/null
+++ b/parquet/internal/encoding/levels_benchmark_test.go
@@ -0,0 +1,105 @@
+// 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_test
+
+import (
+ "encoding/binary"
+ "fmt"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/internal/encoding"
+)
+
+func benchmarkLevelData(b *testing.B, levels []int16, maxLevel int16) []byte {
+ buf := encoding.NewBufferWriter(2*len(levels), memory.DefaultAllocator)
+ defer buf.Release()
+
+ buf.SetOffset(arrow.Int32SizeBytes)
+ var encoder encoding.LevelEncoder
+ encoder.Init(parquet.Encodings.RLE, maxLevel, buf)
+ encoded, err := encoder.Encode(levels)
+ if err != nil {
+ b.Fatal(err)
+ }
+ if encoded != len(levels) {
+ b.Fatalf("encoded %d levels, want %d", encoded, len(levels))
+ }
+
+ buf.SetOffset(0)
+ binary.LittleEndian.PutUint32(buf.Bytes(), uint32(encoder.Len()))
+ return append([]byte(nil), buf.Bytes()...)
+}
+
+func BenchmarkLevelDecoder(b *testing.B) {
+ patterns := []struct {
+ name string
+ fill func([]int16, int16)
+ }{
+ {"all_defined", func(levels []int16, maxLevel int16) {
+ for i := range levels {
+ levels[i] = maxLevel
+ }
+ }},
+ {"mostly_defined", func(levels []int16, maxLevel int16) {
+ for i := range levels {
+ if i%20 != 0 {
+ levels[i] = maxLevel
+ }
+ }
+ }},
+ {"alternating", func(levels []int16, maxLevel int16) {
+ for i := range levels {
+ if i%2 != 0 {
+ levels[i] = maxLevel
+ }
+ }
+ }},
+ }
+
+ for _, size := range []int{1024, 64 * 1024} {
+ for _, maxLevel := range []int16{1, 3} {
+ for _, pattern := range patterns {
+ b.Run(fmt.Sprintf("%s/max_level=%d/levels=%d",
pattern.name, maxLevel, size), func(b *testing.B) {
+ levels := make([]int16, size)
+ pattern.fill(levels, maxLevel)
+ data := benchmarkLevelData(b, levels,
maxLevel)
+ output := make([]int16, size)
+
+ b.ReportAllocs()
+ b.SetBytes(int64(len(output) *
arrow.Int16SizeBytes))
+ b.ResetTimer()
+ for b.Loop() {
+ var decoder
encoding.LevelDecoder
+ if _, err :=
decoder.SetData(parquet.Encodings.RLE, maxLevel, size, data); err != nil {
+ b.Fatal(err)
+ }
+ decoded, _, err :=
decoder.Decode(output)
+ if err != nil {
+ b.Fatal(err)
+ }
+ if decoded != size {
+ b.Fatalf("decoded %d
levels, want %d", decoded, size)
+ }
+ }
+ })
+ }
+ }
+ }
+}
diff --git a/parquet/internal/encoding/levels_test.go
b/parquet/internal/encoding/levels_test.go
index 52827d86..b6092e17 100644
--- a/parquet/internal/encoding/levels_test.go
+++ b/parquet/internal/encoding/levels_test.go
@@ -90,10 +90,12 @@ func verifyDecodingLvls(t *testing.T, enc parquet.Encoding,
maxLvl int16, input
// try multiple decoding on a single setdata call
for ct := 0; ct < decodeCount; ct++ {
offset := ct * numInnerLevels
- lvlCount, _, err = decoder.Decode(output[:numInnerLevels])
+ var valsToRead int64
+ lvlCount, valsToRead, err =
decoder.Decode(output[:numInnerLevels])
assert.NoError(t, err)
assert.Equal(t, numInnerLevels, lvlCount)
assert.Equal(t, input[offset:offset+numInnerLevels],
output[:numInnerLevels])
+ assert.EqualValues(t,
countLevels(input[offset:offset+numInnerLevels], maxLvl), valsToRead)
}
// check the remaining levels
@@ -103,10 +105,12 @@ func verifyDecodingLvls(t *testing.T, enc
parquet.Encoding, maxLvl int16, input
)
if remaining > 0 {
- lvlCount, _, err = decoder.Decode(output[:remaining])
+ var valsToRead int64
+ lvlCount, valsToRead, err = decoder.Decode(output[:remaining])
assert.NoError(t, err)
assert.Equal(t, remaining, lvlCount)
assert.Equal(t, input[levelsCompleted:], output[:remaining])
+ assert.EqualValues(t, countLevels(input[levelsCompleted:],
maxLvl), valsToRead)
}
// test decode zero values
lvlCount, _, err = decoder.Decode(output[:1])
@@ -114,6 +118,16 @@ func verifyDecodingLvls(t *testing.T, enc
parquet.Encoding, maxLvl int16, input
assert.Zero(t, lvlCount)
}
+func countLevels(levels []int16, target int16) int {
+ count := 0
+ for _, level := range levels {
+ if level == target {
+ count++
+ }
+ }
+ return count
+}
+
func verifyDecodingMultipleSetData(t *testing.T, enc parquet.Encoding, max
int16, input []int16, buf [][]byte) {
var (
decoder encoding.LevelDecoder
@@ -173,7 +187,7 @@ func TestLevelsDecodeMultipleBitWidth(t *testing.T) {
}
func TestLevelDecoderPropagatesTruncatedRleLiteralRun(t *testing.T) {
- encoded := append([]byte{17}, make([]byte, 7)...)
+ encoded := append([]byte{17}, []byte{0xff, 0xff, 0xff, 0xff, 0xff,
0xff, 0xff}...)
data := make([]byte, 4, 4+len(encoded))
binary.LittleEndian.PutUint32(data, uint32(len(encoded)))
data = append(data, encoded...)
@@ -183,9 +197,10 @@ func TestLevelDecoderPropagatesTruncatedRleLiteralRun(t
*testing.T) {
assert.NoError(t, err)
levels := make([]int16, 64)
- n, _, err := decoder.Decode(levels)
+ n, valsToRead, err := decoder.Decode(levels)
assert.ErrorIs(t, err, io.ErrUnexpectedEOF)
assert.Equal(t, 32, n)
+ assert.EqualValues(t, 32, valsToRead)
}
func TestLevelsDecodeMultipleSetData(t *testing.T) {
diff --git a/parquet/internal/utils/bit_reader.go
b/parquet/internal/utils/bit_reader.go
index 39575d0c..22f5dc49 100644
--- a/parquet/internal/utils/bit_reader.go
+++ b/parquet/internal/utils/bit_reader.go
@@ -593,6 +593,71 @@ func (b *BitReader) GetBatch(bits uint, out []uint64)
(int, error) {
return i, nil
}
+// GetBatchLevels fills out with bit-packed levels and returns the number equal
+// to maxLevel.
+func (b *BitReader) GetBatchLevels(bits uint, out []int16, maxLevel int16)
(int, int64, error) {
+ if bits > 16 {
+ return 0, 0, errors.New("must be 16 bits or less per read")
+ }
+
+ var maxCount int64
+ length := len(out)
+ i := 0
+
+ for ; i < length && b.bitoffset != 0; i++ {
+ val, err := b.next(bits)
+ if err != nil {
+ return i, maxCount, err
+ }
+ level := int16(val)
+ out[i] = level
+ if level == maxLevel {
+ maxCount++
+ }
+ }
+
+ if _, err := b.reader.Seek(b.byteoffset, io.SeekStart); err != nil {
+ return i, maxCount, err
+ }
+ for i < length {
+ unpackSize := utils.Min(buflen, length-i)
+ numUnpacked, err := unpack32(b.reader,
b.unpackBuf[:unpackSize], int(bits))
+
+ for k, val := range b.unpackBuf[:numUnpacked] {
+ level := int16(val)
+ out[i+k] = level
+ if level == maxLevel {
+ maxCount++
+ }
+ }
+ i += numUnpacked
+ b.byteoffset += int64(numUnpacked * int(bits) / 8)
+ if err != nil {
+ return i, maxCount, err
+ }
+ if numUnpacked == 0 {
+ break
+ }
+ }
+
+ if err := b.fillbuffer(); err != nil {
+ return i, maxCount, err
+ }
+ for ; i < length; i++ {
+ val, err := b.next(bits)
+ if err != nil {
+ return i, maxCount, err
+ }
+ level := int16(val)
+ out[i] = level
+ if level == maxLevel {
+ maxCount++
+ }
+ }
+
+ return i, maxCount, nil
+}
+
// GetValue returns a single value that is bit packed using width as the
number of bits
// and returns false if there weren't enough bits remaining.
func (b *BitReader) GetValue(width int) (uint64, bool) {
diff --git a/parquet/internal/utils/rle.go b/parquet/internal/utils/rle.go
index 29013c53..623cd5f7 100644
--- a/parquet/internal/utils/rle.go
+++ b/parquet/internal/utils/rle.go
@@ -239,6 +239,48 @@ func (r *RleDecoder) GetBatch(values []uint64) (int,
error) {
return read, nil
}
+// GetBatchLevels decodes levels directly into values and returns the number
+// equal to maxLevel.
+func (r *RleDecoder) GetBatchLevels(values []int16, maxLevel int16) (int,
int64, error) {
+ read := 0
+ var maxCount int64
+ out := values
+ for read < len(values) {
+ remain := len(values) - read
+
+ if r.repCount > 0 {
+ repbatch := min(remain, int(r.repCount))
+ level := int16(r.curVal)
+ for i := range out[:repbatch] {
+ out[i] = level
+ }
+ if level == maxLevel {
+ maxCount += int64(repbatch)
+ }
+
+ r.repCount -= int32(repbatch)
+ read += repbatch
+ out = out[repbatch:]
+ } else if r.litCount > 0 {
+ litbatch := min(remain, int(r.litCount))
+ n, count, err := r.r.GetBatchLevels(uint(r.bitWidth),
out[:litbatch], maxLevel)
+ r.litCount -= int32(n)
+ read += n
+ maxCount += count
+ out = out[n:]
+ if err != nil {
+ return read, maxCount, err
+ }
+ if n != litbatch {
+ return read, maxCount, nil
+ }
+ } else if !r.Next() {
+ return read, maxCount, nil
+ }
+ }
+ return read, maxCount, nil
+}
+
func (r *RleDecoder) GetBatchSpaced(vals []uint64, nullcount int, validBits
[]byte, validBitsOffset int64) (int, error) {
if nullcount == 0 {
return r.GetBatch(vals)
