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 85610f71 perf(parquet/internal/utils): pack narrow RLE literal levels
eight at a time (#1419)
85610f71 is described below
commit 85610f718a241fa3f77e77650e889f0846fe4c58
Author: Minh Vu <[email protected]>
AuthorDate: Tue Oct 27 17:19:07 2026 +0100
perf(parquet/internal/utils): pack narrow RLE literal levels eight at a
time (#1419)
## What
Pack narrow Parquet RLE literal levels eight at a time.
## Why
The encoder currently handles each literal level separately.
## Implementation
- Pack full groups directly for widths 1, 2, 4, and 8.
- Keep the existing path for repeated runs, partial groups, and wider
levels.
- Add byte parity, boundary, and fallback coverage with a paired
benchmark.
## Benchmark
**2.61x faster on random levels and 2.25x on alternating levels** in an
isolated encoder model. Other patterns vary.
---------
Co-authored-by: Matt Topol <[email protected]>
---
parquet/internal/utils/rle.go | 50 +++++++
parquet/internal/utils/rle_internal_test.go | 217 ++++++++++++++++++++++++++++
2 files changed, 267 insertions(+)
diff --git a/parquet/internal/utils/rle.go b/parquet/internal/utils/rle.go
index 9b34b350..a1655b12 100644
--- a/parquet/internal/utils/rle.go
+++ b/parquet/internal/utils/rle.go
@@ -609,6 +609,32 @@ func (r *RleEncoder) PutBatchLevels(values []int16) (int,
error) {
}
}
+ // Literal levels usually have a width of only one or two bits.
Pack a
+ // complete group at once instead of feeding eight values
through Put,
+ // then unpacking the uint64 buffer again in WriteValues.
+ if r.repCount == 0 && len(r.buffer) == 0 && len(values)-encoded
>= 8 &&
+ r.BitWidth > 0 && r.BitWidth <= 8 {
+ group := values[encoded : encoded+8]
+ var packed uint64
+ fits := true
+ maxValue := uint64(1<<r.BitWidth) - 1
+ for i, level := range group {
+ if level < 0 || uint64(level) > maxValue {
+ fits = false
+ break
+ }
+ packed |= uint64(level) << (uint(i) *
uint(r.BitWidth))
+ }
+ if fits {
+ if err := r.putPackedLevelGroup(packed); err !=
nil {
+ return encoded + 7, err
+ }
+ r.curVal = uint64(group[7])
+ encoded += 8
+ continue
+ }
+ }
+
batchEnd := min(len(values), encoded+8-len(r.buffer))
for encoded < batchEnd {
if err := r.Put(uint64(values[encoded])); err != nil {
@@ -620,6 +646,30 @@ func (r *RleEncoder) PutBatchLevels(values []int16) (int,
error) {
return encoded, nil
}
+// putPackedLevelGroup writes eight literal levels to the current RLE literal
+// run. PutBatchLevels calls this only when there is no buffered/repeated run,
+// and each level fits in BitWidth bits.
+func (r *RleEncoder) putPackedLevelGroup(packed uint64) error {
+ if r.literalIndicatorOffset == -1 {
+ var err error
+ r.literalIndicatorOffset, err = r.w.SkipBytes(1)
+ if err != nil {
+ return err
+ }
+ }
+
+ if err := r.w.WriteValue(packed, uint(r.BitWidth*8)); err != nil {
+ return err
+ }
+
+ r.litCount += 8
+ // One byte is reserved for the literal header: at most 63 groups.
+ if r.litCount/8+1 >= 1<<6 {
+ return r.flushLiteral(true)
+ }
+ return nil
+}
+
// PutBatchIndices encodes a batch of dictionary indices.
func (r *RleEncoder) PutBatchIndices(values []int32) (int, error) {
encoded := 0
diff --git a/parquet/internal/utils/rle_internal_test.go
b/parquet/internal/utils/rle_internal_test.go
index db02f9cf..a8ad7e52 100644
--- a/parquet/internal/utils/rle_internal_test.go
+++ b/parquet/internal/utils/rle_internal_test.go
@@ -16,6 +16,8 @@
package utils
import (
+ "bytes"
+ "fmt"
"math"
"testing"
)
@@ -74,3 +76,218 @@ func TestRleEncoderBatchReportsValuesBeforeFlushError(t
*testing.T) {
t.Fatalf("encoded %d values, want 8", n)
}
}
+
+func makePackedLevelValues(width, n int, pattern string) []int16 {
+ levels := make([]int16, n)
+ mask := uint32((1 << min(width, 15)) - 1)
+ seed := uint32(0x12345678)
+ for i := range levels {
+ seed ^= seed << 13
+ seed ^= seed >> 17
+ seed ^= seed << 5
+ switch pattern {
+ case "random":
+ levels[i] = int16(seed & mask)
+ case "alternating":
+ levels[i] = int16(uint32(i%2) & mask)
+ case "mostly_defined":
+ if i%20 != 0 {
+ levels[i] = int16(mask)
+ }
+ case "repeated":
+ if (i/13)%2 != 0 {
+ levels[i] = int16(mask)
+ }
+ case "boundary":
+ levels[i] = int16(uint32(i*17+i/3) & mask)
+ }
+ }
+ return levels
+}
+
+func TestRleBatchLevelsPackedMatchesScalar(t *testing.T) {
+ for _, width := range []int{0, 1, 2, 3, 4, 8, 9, 15, 16} {
+ for _, pattern := range []string{"random", "alternating",
"mostly_defined", "repeated", "boundary"} {
+ values := makePackedLevelValues(width, 63*8+17, pattern)
+ t.Run(fmt.Sprintf("width=%d/%s", width, pattern),
func(t *testing.T) {
+ size := MaxRLEBufferSize(width, len(values)) +
MinRLEBufferSize(width) + 32
+ expected := make([]byte, size)
+ scalar :=
NewRleEncoder(NewWriterAtBuffer(expected), width)
+ for _, value := range values {
+ if err := scalar.Put(uint64(value));
err != nil {
+ t.Fatal(err)
+ }
+ }
+ expected = expected[:scalar.Flush()]
+
+ for _, mode := range []string{"whole",
"chunked", "interleaved"} {
+ t.Run(mode, func(t *testing.T) {
+ got := make([]byte, size)
+ enc :=
NewRleEncoder(NewWriterAtBuffer(got), width)
+ chunkSizes := []int{1, 7, 8, 9,
31, 65}
+ for offset, chunk := 0, 0;
offset < len(values); chunk++ {
+ end := len(values)
+ if mode != "whole" {
+ end = min(end,
offset+chunkSizes[chunk%len(chunkSizes)])
+ }
+ if mode ==
"interleaved" && chunk%3 == 0 {
+ for _, value :=
range values[offset:end] {
+ if err
:= enc.Put(uint64(value)); err != nil {
+
t.Fatal(err)
+ }
+ }
+ } else {
+ n, err :=
enc.PutBatchLevels(values[offset:end])
+ if err != nil {
+
t.Fatal(err)
+ }
+ if n !=
end-offset {
+
t.Fatalf("encoded %d, want %d", n, end-offset)
+ }
+ }
+ offset = end
+ }
+ got = got[:enc.Flush()]
+ if !bytes.Equal(got, expected) {
+ t.Fatalf("encoded bytes
differ from scalar at width %d, %s (%d vs %d bytes)", width, pattern, len(got),
len(expected))
+ }
+ decoded := make([]uint64,
len(values))
+ r :=
NewRleDecoder(bytes.NewReader(got), width)
+ n, err := r.GetBatch(decoded)
+ if err != nil || n !=
len(values) {
+ t.Fatalf("decoded %d
values, err %v", n, err)
+ }
+ for i, value := range values {
+ if decoded[i] !=
uint64(value) {
+
t.Fatalf("decoded[%d] = %d, want %d", i, decoded[i], value)
+ }
+ }
+ })
+ }
+ })
+ }
+ }
+}
+
+func TestRleBatchLevelsPackedWriteErrors(t *testing.T) {
+ for _, width := range []int{1, 2, 4, 8} {
+ values := makePackedLevelValues(width, 520, "alternating")
+ for _, capacity := range []int{0, 1, 2, 8, 16, 63, 64, 128,
503} {
+ t.Run(fmt.Sprintf("width=%d/capacity=%d", width,
capacity), func(t *testing.T) {
+ expected, got := make([]byte, capacity),
make([]byte, capacity)
+ scalar :=
NewRleEncoder(NewWriterAtBuffer(expected), width)
+ scalarCount := 0
+ var scalarErr error
+ for _, value := range values {
+ if scalarErr =
scalar.Put(uint64(value)); scalarErr != nil {
+ break
+ }
+ scalarCount++
+ }
+ enc := NewRleEncoder(NewWriterAtBuffer(got),
width)
+ n, err := enc.PutBatchLevels(values)
+ if (err != nil) != (scalarErr != nil) || n !=
scalarCount {
+ t.Fatalf("count/error: got (%d, %v),
want (%d, %v)", n, err, scalarCount, scalarErr)
+ }
+ if !bytes.Equal(got, expected) {
+ t.Fatal("partial encoded output
differs")
+ }
+ })
+ }
+ }
+}
+
+// Compare partial output and errors against encoding each level separately.
+func TestRleBatchLevelsPackedBatchWriteErrors(t *testing.T) {
+ // When a repeated run follows literals, the batch API flushes the
literal
+ // header before consuming the next eight values. Put has already
accepted
+ // seven of those values when it encounters the same error. Keep
explicit
+ // batch counts for these mostly_defined cases (width -> capacity ->
count).
+ flushErrorCounts := map[int]map[int]int{
+ 1: {0: 8, 8: 48, 16: 88, 64: 328},
+ 2: {0: 8},
+ 4: {0: 8, 63: 188},
+ }
+ for _, width := range []int{1, 2, 4, 8} {
+ for _, pattern := range []string{"random", "mostly_defined",
"repeated", "alternating"} {
+ values := makePackedLevelValues(width, 520, pattern)
+ for _, capacity := range []int{0, 1, 2, 8, 16, 63, 64,
128, 503} {
+ t.Run(fmt.Sprintf("width=%d/%s/capacity=%d",
width, pattern, capacity), func(t *testing.T) {
+ expected, got := make([]byte,
capacity), make([]byte, capacity)
+ scalar :=
NewRleEncoder(NewWriterAtBuffer(expected), width)
+ scalarCount := 0
+ var scalarErr error
+ for _, value := range values {
+ if scalarErr =
scalar.Put(uint64(value)); scalarErr != nil {
+ break
+ }
+ scalarCount++
+ }
+ current :=
NewRleEncoder(NewWriterAtBuffer(got), width)
+ n, err := current.PutBatchLevels(values)
+ wantCount := scalarCount
+ if count, ok :=
flushErrorCounts[width][capacity]; pattern == "mostly_defined" && ok {
+ wantCount = count
+ if err == nil {
+ t.Fatal("expected a
literal header write error")
+ }
+ }
+ if n != wantCount || (err != nil) !=
(scalarErr != nil) || !bytes.Equal(expected, got) {
+ t.Fatalf("batch (%d, %v), want
count %d, scalar (%d, %v), bytes equal: %t", n, err, wantCount, scalarCount,
scalarErr, bytes.Equal(expected, got))
+ }
+ })
+ }
+ }
+ }
+}
+
+func TestRleBatchLevelsPackedInvalidValues(t *testing.T) {
+ // Invalid levels preserve the per-value writer's behavior, including
+ // negative input cast directly to uint64.
+ for _, width := range []int{1, 2, 4, 8} {
+ values := []int16{0, 1, 0, 1, 0, 1, 0, -1, 0, 1, 0, 1, 0, 1, 0,
1024}
+ size := MaxRLEBufferSize(width, len(values)) +
MinRLEBufferSize(width) + 32
+ expected, got := make([]byte, size), make([]byte, size)
+ scalar := NewRleEncoder(NewWriterAtBuffer(expected), width)
+ scalarCount := 0
+ var scalarErr error
+ for _, value := range values {
+ if scalarErr = scalar.Put(uint64(value)); scalarErr !=
nil {
+ break
+ }
+ scalarCount++
+ }
+ current := NewRleEncoder(NewWriterAtBuffer(got), width)
+ n, err := current.PutBatchLevels(values)
+ if scalarCount != n || (scalarErr != nil) != (err != nil) ||
+ !bytes.Equal(expected[:scalar.Flush()],
got[:current.Flush()]) {
+ t.Fatalf("width %d: invalid-level output changed",
width)
+ }
+ }
+}
+
+func BenchmarkRleBatchLevelsPacked(b *testing.B) {
+ for _, width := range []int{1, 2, 4, 8, 16} {
+ for _, pattern := range []string{"random", "alternating",
"mostly_defined", "repeated"} {
+ for _, size := range []int{1024, 65536} {
+ values := makePackedLevelValues(width, size,
pattern)
+ b.Run(fmt.Sprintf("width=%d/%s/n=%d", width,
pattern, size), func(b *testing.B) {
+ output := make([]byte,
MaxRLEBufferSize(width, len(values))+MinRLEBufferSize(width)+32)
+ enc :=
NewRleEncoder(NewWriterAtBuffer(output), width)
+ b.SetBytes(int64(len(values) * 2))
+ b.ReportAllocs()
+ for b.Loop() {
+ enc.Clear()
+ n, err :=
enc.PutBatchLevels(values)
+ if err != nil || n !=
len(values) {
+ b.Fatalf("encoded %d
levels: %v", n, err)
+ }
+ if got := enc.Flush(); got >
len(output) {
+ b.Fatalf("encoded %d
bytes into %d byte output", got, len(output))
+ }
+ }
+ })
+ }
+ }
+ }
+}
