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 2a18bcee perf(parquet): encode RLE levels in batches (#1177)
2a18bcee is described below
commit 2a18bcee01b9194f0820523ae62bb90e0603a5f2
Author: Minh Vu <[email protected]>
AuthorDate: Wed Aug 26 20:11:02 2026 +0200
perf(parquet): encode RLE levels in batches (#1177)
### What changed
- add a batch path for repetition and definition level encoding
- scan repeated runs once instead of calling the scalar encoder for
every level
- keep the scalar path around literal boundaries
- preserve the existing encoded bytes across split streaming calls
- add level encoder and scalar-versus-batch benchmarks
### Benchmarks
Apple M1 Pro:
| RLE workload | Before | After | Change |
| --- | ---: | ---: | ---: |
| 65,536 all defined | 151,051 ns/op | 31,143 ns/op | -79% |
| 65,536 mostly defined | 324,228 ns/op | 252,108 ns/op | -22% |
| 65,536 alternating | 339,128 ns/op | 341,149 ns/op | +0.6% |
The alternating case is the fallback path. Allocations remain at zero.
---
parquet/internal/encoding/levels.go | 14 +--
.../encoding/levels_encode_benchmark_test.go | 138 +++++++++++++++++++++
parquet/internal/utils/bit_reader_test.go | 59 +++++++++
parquet/internal/utils/rle.go | 55 +++++++-
parquet/internal/utils/rle_internal_test.go | 76 ++++++++++++
5 files changed, 329 insertions(+), 13 deletions(-)
diff --git a/parquet/internal/encoding/levels.go
b/parquet/internal/encoding/levels.go
index e88a18ef..3da9a575 100644
--- a/parquet/internal/encoding/levels.go
+++ b/parquet/internal/encoding/levels.go
@@ -98,12 +98,7 @@ func (l *LevelEncoder) EncodeNoFlush(lvls []int16) (nencoded
int, err error) {
switch l.encoding {
case format.Encoding_RLE:
- for _, level := range lvls {
- if err = l.rle.Put(uint64(level)); err != nil {
- return
- }
- nencoded++
- }
+ nencoded, err = l.rle.PutBatchLevels(lvls)
default:
for _, level := range lvls {
if err = l.bit.WriteValue(uint64(level),
uint(l.bitWidth)); err != nil {
@@ -140,12 +135,7 @@ func (l *LevelEncoder) Encode(lvls []int16) (nencoded int,
err error) {
switch l.encoding {
case format.Encoding_RLE:
defer func() { l.rleLen = l.rle.Flush() }()
- for _, level := range lvls {
- if err = l.rle.Put(uint64(level)); err != nil {
- return
- }
- nencoded++
- }
+ nencoded, err = l.rle.PutBatchLevels(lvls)
default:
defer l.bit.Flush(false)
diff --git a/parquet/internal/encoding/levels_encode_benchmark_test.go
b/parquet/internal/encoding/levels_encode_benchmark_test.go
new file mode 100644
index 00000000..3b03b503
--- /dev/null
+++ b/parquet/internal/encoding/levels_encode_benchmark_test.go
@@ -0,0 +1,138 @@
+// 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 (
+ "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"
+ parquetutils "github.com/apache/arrow-go/v18/parquet/internal/utils"
+)
+
+func BenchmarkLevelEncoder(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)
+ output :=
encoding.NewBufferWriter(encoding.LevelEncodingMaxBufferSize(parquet.Encodings.RLE,
maxLevel, size), memory.DefaultAllocator)
+ defer output.Release()
+
+ var encoder encoding.LevelEncoder
+ encoder.Init(parquet.Encodings.RLE,
maxLevel, output)
+ b.ReportAllocs()
+ b.SetBytes(int64(len(levels) *
arrow.Int16SizeBytes))
+ b.ResetTimer()
+ for b.Loop() {
+ encoder.Reset(maxLevel)
+ encoded, err :=
encoder.Encode(levels)
+ if err != nil {
+ b.Fatal(err)
+ }
+ if encoded != size {
+ b.Fatalf("encoded %d
levels, want %d", encoded, size)
+ }
+ }
+ })
+ }
+ }
+ }
+}
+
+func BenchmarkRleLevelEncoder(b *testing.B) {
+ patterns := []struct {
+ name string
+ fill func([]int16)
+ }{
+ {"all_defined", func(levels []int16) {
+ for i := range levels {
+ levels[i] = 1
+ }
+ }},
+ {"mostly_defined", func(levels []int16) {
+ for i := range levels {
+ if i%20 != 0 {
+ levels[i] = 1
+ }
+ }
+ }},
+ {"alternating", func(levels []int16) {
+ for i := range levels {
+ levels[i] = int16(i % 2)
+ }
+ }},
+ }
+
+ for _, pattern := range patterns {
+ b.Run(pattern.name, func(b *testing.B) {
+ levels := make([]int16, 64*1024)
+ pattern.fill(levels)
+ for _, mode := range []string{"scalar", "batch"} {
+ b.Run(mode, func(b *testing.B) {
+ output := make([]byte,
parquetutils.MaxRLEBufferSize(1, len(levels)))
+ encoder :=
parquetutils.NewRleEncoder(parquetutils.NewWriterAtBuffer(output), 1)
+
+ b.ReportAllocs()
+ b.SetBytes(int64(len(levels) *
arrow.Int16SizeBytes))
+ for b.Loop() {
+ encoder.Clear()
+ if mode == "scalar" {
+ for _, level := range
levels {
+ if err :=
encoder.Put(uint64(level)); err != nil {
+
b.Fatal(err)
+ }
+ }
+ } else if _, err :=
encoder.PutBatchLevels(levels); err != nil {
+ b.Fatal(err)
+ }
+ encoder.Flush()
+ }
+ })
+ }
+ })
+ }
+}
diff --git a/parquet/internal/utils/bit_reader_test.go
b/parquet/internal/utils/bit_reader_test.go
index 0468ff13..44804c24 100644
--- a/parquet/internal/utils/bit_reader_test.go
+++ b/parquet/internal/utils/bit_reader_test.go
@@ -661,6 +661,65 @@ func TestRleRandom(t *testing.T) {
suite.Run(t, new(RLERandomSuite))
}
+func TestRleBatchLevelsMatchesScalar(t *testing.T) {
+ patterns := map[string][]int16{
+ "all defined": make([]int16, 257),
+ "alternating": make([]int16, 257),
+ "long alternating": make([]int16, 63*8+17),
+ "literal boundary": make([]int16, 63*8+17),
+ "mixed runs": {},
+ }
+ for i := range patterns["all defined"] {
+ patterns["all defined"][i] = 3
+ patterns["alternating"][i] = int16(i % 2)
+ }
+ for i := range patterns["long alternating"] {
+ patterns["long alternating"][i] = int16(i % 2)
+ }
+ for i := range patterns["literal boundary"] {
+ if i < 63*8-3 {
+ patterns["literal boundary"][i] = int16(i % 2)
+ } else {
+ patterns["literal boundary"][i] = 3
+ }
+ }
+ for runLength := 1; runLength <= 32; runLength++ {
+ for range runLength {
+ patterns["mixed runs"] = append(patterns["mixed runs"],
int16(runLength%4))
+ }
+ }
+
+ for name, levels := range patterns {
+ t.Run(name, func(t *testing.T) {
+ bufSize := utils.MaxRLEBufferSize(2, len(levels))
+ scalarBuf := make([]byte, bufSize)
+ scalar :=
utils.NewRleEncoder(utils.NewWriterAtBuffer(scalarBuf), 2)
+ for _, level := range levels {
+ if err := scalar.Put(uint64(level)); err != nil
{
+ t.Fatal(err)
+ }
+ }
+ scalarLen := scalar.Flush()
+
+ batchBuf := make([]byte, bufSize)
+ batch :=
utils.NewRleEncoder(utils.NewWriterAtBuffer(batchBuf), 2)
+ chunkSizes := []int{1, 7, 8, 9, 31}
+ for offset, chunk := 0, 0; offset < len(levels);
chunk++ {
+ end :=
min(offset+chunkSizes[chunk%len(chunkSizes)], len(levels))
+ n, err :=
batch.PutBatchLevels(levels[offset:end])
+ if err != nil {
+ t.Fatal(err)
+ }
+ assert.Equal(t, end-offset, n)
+ offset = end
+ }
+ batchLen := batch.Flush()
+
+ assert.Equal(t, scalarBuf[:scalarLen],
batchBuf[:batchLen])
+ })
+ }
+}
+
func (r *RLETestSuite) ValidateRle(vals []uint64, width int, expected []byte,
explen int) {
const buflen = 64 * 1024
buf := make([]byte, buflen)
diff --git a/parquet/internal/utils/rle.go b/parquet/internal/utils/rle.go
index 596f2018..29013c53 100644
--- a/parquet/internal/utils/rle.go
+++ b/parquet/internal/utils/rle.go
@@ -407,7 +407,7 @@ func (r *RleEncoder) flushLiteral(updateIndicator bool)
(err error) {
}
func (r *RleEncoder) flushRepeated() (ret bool) {
- indicator := r.repCount << 1
+ indicator := uint64(r.repCount) << 1
ret = r.w.WriteVlqInt(uint64(indicator))
ret = ret && r.w.WriteAligned(r.curVal,
int(bitutil.BytesForBits(int64(r.BitWidth))))
@@ -444,6 +444,59 @@ func (r *RleEncoder) Put(value uint64) error {
return nil
}
+// PutBatchLevels encodes a batch of repetition or definition levels.
+func (r *RleEncoder) PutBatchLevels(values []int16) (int, error) {
+ encoded := 0
+ for encoded < len(values) {
+ value := values[encoded]
+ if r.repCount >= 8 && r.curVal == uint64(value) {
+ runEnd := encoded + 1
+ for runEnd < len(values) && values[runEnd] == value {
+ runEnd++
+ }
+
+ runLength := min(runEnd-encoded,
int(math.MaxInt32-r.repCount))
+ r.repCount += int32(runLength)
+ encoded += runLength
+ if r.repCount == math.MaxInt32 && encoded < len(values)
&& values[encoded] == value {
+ if !r.flushRepeated() {
+ return encoded, errors.New("failed to
flush repeated value")
+ }
+ }
+ continue
+ }
+ if r.repCount == 0 && len(r.buffer) == 0 && len(values)-encoded
>= 8 && values[encoded+7] == value {
+ runEnd := encoded + 1
+ for runEnd < len(values) && values[runEnd] == value {
+ runEnd++
+ }
+ if runEnd-encoded >= 8 {
+ r.curVal = uint64(value)
+ if r.litCount != 0 {
+ r.repCount = 8
+ if err := r.flushLiteral(true); err !=
nil {
+ return encoded, err
+ }
+ encoded += 8
+ }
+ runLength := min(runEnd-encoded,
int(math.MaxInt32-r.repCount))
+ r.repCount += int32(runLength)
+ encoded += runLength
+ continue
+ }
+ }
+
+ batchEnd := min(len(values), encoded+8-len(r.buffer))
+ for encoded < batchEnd {
+ if err := r.Put(uint64(values[encoded])); err != nil {
+ return encoded, err
+ }
+ encoded++
+ }
+ }
+ return encoded, nil
+}
+
func (r *RleEncoder) Clear() {
r.curVal = 0
r.repCount = 0
diff --git a/parquet/internal/utils/rle_internal_test.go
b/parquet/internal/utils/rle_internal_test.go
new file mode 100644
index 00000000..db02f9cf
--- /dev/null
+++ b/parquet/internal/utils/rle_internal_test.go
@@ -0,0 +1,76 @@
+// 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 utils
+
+import (
+ "math"
+ "testing"
+)
+
+func TestRleEncoderMaximumRepeatedRun(t *testing.T) {
+ buf := make([]byte, 16)
+ enc := NewRleEncoder(NewWriterAtBuffer(buf), 1)
+ enc.curVal = 1
+ enc.repCount = math.MaxInt32
+
+ if n := enc.Flush(); n != 6 {
+ t.Fatalf("encoded %d bytes, want 6", n)
+ }
+ want := []byte{0xfe, 0xff, 0xff, 0xff, 0x0f, 1}
+ for i, value := range want {
+ if buf[i] != value {
+ t.Fatalf("byte %d = %#x, want %#x", i, buf[i], value)
+ }
+ }
+}
+
+func TestRleEncoderBatchSplitsMaximumRepeatedRun(t *testing.T) {
+ buf := make([]byte, 16)
+ enc := NewRleEncoder(NewWriterAtBuffer(buf), 1)
+ enc.curVal = 1
+ enc.repCount = math.MaxInt32 - 4
+
+ n, err := enc.PutBatchLevels([]int16{1, 1, 1, 1, 1, 1})
+ if err != nil {
+ t.Fatal(err)
+ }
+ if n != 6 {
+ t.Fatalf("encoded %d values, want 6", n)
+ }
+ if n := enc.Flush(); n != 8 {
+ t.Fatalf("encoded %d bytes, want 8", n)
+ }
+
+ want := []byte{0xfe, 0xff, 0xff, 0xff, 0x0f, 1, 4, 1}
+ for i, value := range want {
+ if buf[i] != value {
+ t.Fatalf("byte %d = %#x, want %#x", i, buf[i], value)
+ }
+ }
+}
+
+func TestRleEncoderBatchReportsValuesBeforeFlushError(t *testing.T) {
+ enc := NewRleEncoder(NewWriterAtBuffer(nil), 1)
+ values := []int16{0, 1, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1}
+
+ n, err := enc.PutBatchLevels(values)
+ if err == nil {
+ t.Fatal("expected a write error")
+ }
+ if n != 8 {
+ t.Fatalf("encoded %d values, want 8", n)
+ }
+}
