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 55d009b9 perf(parquet): encode dictionary indices in RLE batches
(#1328)
55d009b9 is described below
commit 55d009b9e23b2966788f7078d1f567c165445dd9
Author: Minh Vu <[email protected]>
AuthorDate: Mon Sep 21 19:48:03 2026 +0200
perf(parquet): encode dictionary indices in RLE batches (#1328)
### What
- Encode Parquet dictionary indices in RLE batches.
### Why
- `WriteIndices` currently sends every index through `Put`, including
long repeated runs and complete literal groups.
### Implementation
- Add `PutBatchIndices` for `int32` indices.
- Encode repeated runs and complete literal groups directly.
- Pack literal groups through one batched bit-writer call.
- Keep scalar handling for partial literal groups.
- Preserve encoded bytes and dictionary state on write errors.
### Benchmark
Command, run six times per version with before/after order alternated:
`go test -run '^$' -bench '^BenchmarkDictEncoderWriteIndices$' -benchmem
-benchtime=500ms -count=1 ./parquet/internal/encoding`
Apple M1 Pro, Go 1.26.3, `GOMAXPROCS=1`. `WriteIndices` on 65,536
indices.
| Input | Before | After | Time change |
| --- | ---: | ---: | ---: |
| Constant | 187.7 µs | 26.0 µs | -86.1% |
| Runs of 8 | 388.9 µs | 200.9 µs | -48.3% |
| Runs of 32 | 219.4 µs | 67.0 µs | -69.5% |
| Runs of 256 | 197.6 µs | 41.9 µs | -78.8% |
| Random, 16 entries | 580.7 µs | 306.3 µs | -47.3% |
| Random, 256 entries | 452.8 µs | 302.8 µs | -33.1% |
| Alternating | 400.7 µs | 270.6 µs | -32.5% |
Allocations stay at 232 B/op and 4 allocs/op. The benchmark also covers
64-index batches.
### Compatibility
- Encoded bytes match the existing scalar path.
- No Parquet format or public API changes.
### Tests
- `go test ./parquet/...`
- `go test -race ./parquet/internal/utils ./parquet/internal/encoding`
- `go test -tags noasm ./parquet/internal/utils
./parquet/internal/encoding`
- `go vet -composites=false ./parquet/internal/utils
./parquet/internal/encoding`
---
.../encoding/dict_indices_benchmark_test.go | 63 +++++++
parquet/internal/encoding/dict_indices_test.go | 93 +++++++++
parquet/internal/encoding/encoder.go | 6 +-
parquet/internal/utils/bit_writer.go | 28 +++
parquet/internal/utils/bit_writer_batch_test.go | 57 ++++++
parquet/internal/utils/rle.go | 85 ++++++++-
parquet/internal/utils/rle_indices_test.go | 207 +++++++++++++++++++++
7 files changed, 531 insertions(+), 8 deletions(-)
diff --git a/parquet/internal/encoding/dict_indices_benchmark_test.go
b/parquet/internal/encoding/dict_indices_benchmark_test.go
new file mode 100644
index 00000000..6914876c
--- /dev/null
+++ b/parquet/internal/encoding/dict_indices_benchmark_test.go
@@ -0,0 +1,63 @@
+// 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 (
+ "fmt"
+ "math/rand/v2"
+ "testing"
+)
+
+func BenchmarkDictEncoderWriteIndices(b *testing.B) {
+ for _, length := range []int{64, 65536} {
+ for _, pattern := range []struct {
+ name string
+ cardinality, runLength int
+ }{
+ {"constant", 1, 65536},
+ {"runs_8", 256, 8},
+ {"runs_32", 256, 32},
+ {"runs_256", 256, 256},
+ {"random_16", 16, 0},
+ {"random_256", 256, 0},
+ {"alternating", 2, 1},
+ } {
+ b.Run(fmt.Sprintf("%s/indices=%d", pattern.name,
length), func(b *testing.B) {
+ indices := make([]int32, length)
+ rng := rand.New(rand.NewPCG(0, 0))
+ for i := range indices {
+ if pattern.runLength == 0 {
+ indices[i] =
int32(rng.IntN(pattern.cardinality))
+ } else {
+ indices[i] = int32((i /
pattern.runLength) % pattern.cardinality)
+ }
+ }
+ enc := makeDictIndicesEncoder(b,
pattern.cardinality)
+ enc.idxValues = indices
+ output := make([]byte,
enc.EstimatedDataEncodedSize())
+ b.ReportAllocs()
+ b.SetBytes(int64(length * 4))
+ for b.Loop() {
+ enc.idxValues = indices
+ if _, err := enc.WriteIndices(output);
err != nil {
+ b.Fatal(err)
+ }
+ }
+ })
+ }
+ }
+}
diff --git a/parquet/internal/encoding/dict_indices_test.go
b/parquet/internal/encoding/dict_indices_test.go
new file mode 100644
index 00000000..3bede9f3
--- /dev/null
+++ b/parquet/internal/encoding/dict_indices_test.go
@@ -0,0 +1,93 @@
+// 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 (
+ "fmt"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/parquet/internal/utils"
+ "github.com/stretchr/testify/require"
+)
+
+func makeDictIndicesEncoder(t testing.TB, cardinality int) dictEncoder {
+ t.Helper()
+ memo := NewInt32Dictionary()
+ for i := range cardinality {
+ _, _, err := memo.GetOrInsert(int32(i))
+ require.NoError(t, err)
+ }
+ t.Cleanup(memo.Reset)
+ return dictEncoder{memo: memo}
+}
+
+func TestDictEncoderWriteIndicesMatchesScalar(t *testing.T) {
+ for _, tc := range []struct{ cardinality, width int }{
+ {0, 0}, {1, 1}, {2, 1}, {256, 8}, {257, 9}, {65537, 17},
+ } {
+ t.Run(fmt.Sprint(tc.cardinality), func(t *testing.T) {
+ enc := makeDictIndicesEncoder(t, tc.cardinality)
+ var indices []int32
+ if tc.cardinality > 0 {
+ indices = make([]int32, 1031)
+ for i := range indices {
+ indices[i] = int32(tc.cardinality - 1)
+ if i%31 < 9 {
+ indices[i] = int32(i %
tc.cardinality)
+ }
+ }
+ }
+ for range 2 {
+ enc.idxValues = indices
+ enc.rawDataSize = int64(len(indices) * 4)
+ expected := make([]byte,
enc.EstimatedDataEncodedSize())
+ expected[0] = byte(tc.width)
+ scalar :=
utils.NewRleEncoder(utils.NewWriterAtBuffer(expected[1:]), tc.width)
+ for _, index := range indices {
+ require.NoError(t,
scalar.Put(uint64(index)))
+ }
+ expectedSize := scalar.Flush() + 1
+ output := make([]byte, len(expected))
+ n, err := enc.WriteIndices(output)
+ require.NoError(t, err)
+ require.Equal(t, expected[:expectedSize],
output[:n])
+ require.Empty(t, enc.idxValues)
+ require.Zero(t, enc.rawDataSize)
+ require.Equal(t, tc.cardinality,
enc.NumEntries())
+ }
+ })
+ }
+}
+
+func TestDictEncoderWriteIndicesPreservesStateOnError(t *testing.T) {
+ enc := makeDictIndicesEncoder(t, 2)
+ indices := []int32{0, 0, 0, 0, 0, 0, 0, 0, 0, 1}
+ enc.idxValues = indices
+ enc.rawDataSize = int64(len(indices) * 4)
+ n, err := enc.WriteIndices(make([]byte, 1))
+ require.Error(t, err)
+ require.Equal(t, -1, n)
+ require.Equal(t, indices, enc.idxValues)
+ require.EqualValues(t, len(indices)*4, enc.rawDataSize)
+
+ output := make([]byte, enc.EstimatedDataEncodedSize())
+ n, err = enc.WriteIndices(output)
+ require.NoError(t, err)
+ require.Equal(t, []byte{1, 18, 0, 2, 1}, output[:n])
+ require.Empty(t, enc.idxValues)
+ require.Zero(t, enc.rawDataSize)
+}
diff --git a/parquet/internal/encoding/encoder.go
b/parquet/internal/encoding/encoder.go
index 58851b30..09009246 100644
--- a/parquet/internal/encoding/encoder.go
+++ b/parquet/internal/encoding/encoder.go
@@ -391,10 +391,8 @@ func (d *dictEncoder) WriteIndices(out []byte) (int,
error) {
out[0] = byte(d.BitWidth())
enc := utils.NewRleEncoder(utils.NewWriterAtBuffer(out[1:]),
d.BitWidth())
- for _, idx := range d.idxValues {
- if err := enc.Put(uint64(idx)); err != nil {
- return -1, err
- }
+ if _, err := enc.PutBatchIndices(d.idxValues); err != nil {
+ return -1, err
}
nbytes := enc.Flush()
diff --git a/parquet/internal/utils/bit_writer.go
b/parquet/internal/utils/bit_writer.go
index dec87598..e0d963a5 100644
--- a/parquet/internal/utils/bit_writer.go
+++ b/parquet/internal/utils/bit_writer.go
@@ -129,6 +129,34 @@ func (b *BitWriter) WriteValue(v uint64, nbits uint) error
{
return nil
}
+// WriteValues writes values using nbits to pack them into the stream.
+func (b *BitWriter) WriteValues(values []uint64, nbits uint) error {
+ buffer := b.buffer
+ bitoffset := b.bitoffset
+ byteoffset := b.byteoffset
+ for _, v := range values {
+ buffer |= v << bitoffset
+ bitoffset += nbits
+
+ if bitoffset >= 64 {
+ binary.LittleEndian.PutUint64(b.raw[:], buffer)
+ if _, err := b.wr.WriteAt(b.raw[:], int64(byteoffset));
err != nil {
+ b.buffer = buffer
+ b.bitoffset = bitoffset
+ b.byteoffset = byteoffset
+ return err
+ }
+ bitoffset -= 64
+ buffer = v >> (nbits - bitoffset)
+ byteoffset += 8
+ }
+ }
+ b.buffer = buffer
+ b.bitoffset = bitoffset
+ b.byteoffset = byteoffset
+ return nil
+}
+
// Flush will flush any buffered data to the underlying writer, pass true if
// the next write should be byte-aligned after this flush.
func (b *BitWriter) Flush(align bool) {
diff --git a/parquet/internal/utils/bit_writer_batch_test.go
b/parquet/internal/utils/bit_writer_batch_test.go
new file mode 100644
index 00000000..29363466
--- /dev/null
+++ b/parquet/internal/utils/bit_writer_batch_test.go
@@ -0,0 +1,57 @@
+// 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_test
+
+import (
+ "fmt"
+ "math"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow/bitutil"
+ "github.com/apache/arrow-go/v18/parquet/internal/utils"
+ "github.com/stretchr/testify/require"
+)
+
+func TestBitWriterWriteValuesMatchesScalar(t *testing.T) {
+ const nvalues = 17
+ for width := uint(0); width <= 64; width++ {
+ t.Run(fmt.Sprintf("width=%d", width), func(t *testing.T) {
+ mask := uint64(math.MaxUint64)
+ if width < 64 {
+ mask = (uint64(1) << width) - 1
+ }
+ values := make([]uint64, nvalues)
+ for i := range values {
+ values[i] = (uint64(i)*0x9e3779b97f4a7c15 +
uint64(i/3)) & mask
+ }
+
+ outputSize :=
int(bitutil.BytesForBits(int64(3+width*nvalues))) + 8
+ scalarOutput := make([]byte, outputSize)
+ batchOutput := make([]byte, outputSize)
+ scalar :=
utils.NewBitWriter(utils.NewWriterAtBuffer(scalarOutput))
+ batch :=
utils.NewBitWriter(utils.NewWriterAtBuffer(batchOutput))
+ require.NoError(t, scalar.WriteValue(5, 3))
+ require.NoError(t, batch.WriteValue(5, 3))
+ for _, value := range values {
+ require.NoError(t, scalar.WriteValue(value,
width))
+ }
+ require.NoError(t, batch.WriteValues(values, width))
+ scalar.Flush(false)
+ batch.Flush(false)
+ require.Equal(t, scalarOutput[:scalar.Written()],
batchOutput[:batch.Written()])
+ })
+ }
+}
diff --git a/parquet/internal/utils/rle.go b/parquet/internal/utils/rle.go
index 2d25ad95..57fe6a1a 100644
--- a/parquet/internal/utils/rle.go
+++ b/parquet/internal/utils/rle.go
@@ -464,10 +464,8 @@ func (r *RleEncoder) flushLiteral(updateIndicator bool)
(err error) {
}
}
- for _, val := range r.buffer {
- if err = r.w.WriteValue(val, uint(r.BitWidth)); err != nil {
- return
- }
+ if err = r.w.WriteValues(r.buffer, uint(r.BitWidth)); err != nil {
+ return
}
r.buffer = r.buffer[:0]
@@ -575,6 +573,85 @@ func (r *RleEncoder) PutBatchLevels(values []int16) (int,
error) {
return encoded, nil
}
+// PutBatchIndices encodes a batch of dictionary indices.
+func (r *RleEncoder) PutBatchIndices(values []int32) (int, error) {
+ encoded := 0
+ for encoded < len(values) {
+ value := values[encoded]
+ if r.repCount >= 8 {
+ if r.curVal != uint64(value) {
+ if !r.flushRepeated() {
+ return encoded, errors.New("failed to
flush repeated value")
+ }
+ } else {
+ 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
+ }
+ }
+
+ if r.repCount == 0 && len(r.buffer) == 0 && len(values)-encoded
>= 8 {
+ if err := r.putBatchIndicesLiteral(values[encoded:]);
err != nil {
+ return encoded + 7, err
+ }
+ encoded += 8
+ 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
+}
+
+// putBatchIndicesLiteral writes the first complete literal group directly.
+// The caller ensures that values contains at least eight entries and that the
+// encoder has no active repeated run or buffered values.
+func (r *RleEncoder) putBatchIndicesLiteral(values []int32) error {
+ r.buffer = r.buffer[:8]
+ for i, index := range values[:8] {
+ r.buffer[i] = uint64(index)
+ }
+ r.curVal = r.buffer[7]
+ return r.flushBuffered(false)
+}
+
func (r *RleEncoder) Clear() {
r.curVal = 0
r.repCount = 0
diff --git a/parquet/internal/utils/rle_indices_test.go
b/parquet/internal/utils/rle_indices_test.go
new file mode 100644
index 00000000..5aed2ec1
--- /dev/null
+++ b/parquet/internal/utils/rle_indices_test.go
@@ -0,0 +1,207 @@
+// 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 (
+ "bytes"
+ "fmt"
+ "math"
+ "math/rand/v2"
+ "testing"
+
+ "github.com/stretchr/testify/require"
+)
+
+func TestRleBatchIndicesMatchesScalar(t *testing.T) {
+ for _, width := range []int{0, 1, 4, 8, 16, 17, 31} {
+ maxValue := int32((uint32(1) << width) - 1)
+ patterns := map[string][]int32{
+ "empty": nil,
+ "single": {maxValue},
+ "constant": make([]int32, 1024),
+ "alternating": make([]int32, 1024),
+ "random": make([]int32, 1024),
+ "literal boundary": make([]int32, 63*8+17),
+ }
+ rng := rand.New(rand.NewPCG(0, 0))
+ for i := range patterns["constant"] {
+ patterns["constant"][i] = maxValue
+ patterns["alternating"][i] = int32(i%2) & maxValue
+ patterns["random"][i] = int32(rng.Uint32() &
uint32(maxValue))
+ }
+ for i := range patterns["literal boundary"] {
+ patterns["literal boundary"][i] = maxValue
+ if i < 63*8-3 {
+ patterns["literal boundary"][i] = int32(i%2) &
maxValue
+ }
+ }
+ for _, runLength := range []int{8, 9, 16, 32} {
+ values := make([]int32, 1024)
+ for i := range values {
+ values[i] = int32((i/runLength)%2) & maxValue
+ }
+ patterns[fmt.Sprintf("consecutive runs=%d", runLength)]
= values
+ }
+ for offset := range 8 {
+ for _, runLength := range []int{7, 8, 9, 32} {
+ values := make([]int32, offset+runLength+9)
+ for i := range values {
+ values[i] = int32(i%2) & maxValue
+ }
+ for i := offset; i < offset+runLength; i++ {
+ values[i] = maxValue
+ }
+ patterns[fmt.Sprintf("offset=%d/run=%d",
offset, runLength)] = values
+ }
+ }
+ for name, values := range patterns {
+ t.Run(fmt.Sprintf("width=%d/%s", width, name), func(t
*testing.T) {
+ outputSize := MaxRLEBufferSize(width,
len(values)) + MinRLEBufferSize(width)
+ scalarOutput := make([]byte, outputSize)
+ scalar :=
NewRleEncoder(NewWriterAtBuffer(scalarOutput), width)
+ for _, value := range values {
+ require.NoError(t,
scalar.Put(uint64(value)))
+ }
+ scalarSize := scalar.Flush()
+ for _, mode := range []string{"whole",
"chunked", "interleaved"} {
+ t.Run(mode, func(t *testing.T) {
+ output := make([]byte,
outputSize)
+ batch :=
NewRleEncoder(NewWriterAtBuffer(output), width)
+ for range 2 {
+ batch.Clear()
+ n, err :=
batch.PutBatchIndices(nil)
+ require.NoError(t, err)
+ require.Zero(t, n)
+ chunkSizes := []int{1,
7, 8, 9, 31}
+ for offset, chunk := 0,
0; offset < len(values); chunk++ {
+ end :=
len(values)
+ if mode !=
"whole" {
+ end =
min(offset+chunkSizes[chunk%len(chunkSizes)], end)
+ }
+ if mode ==
"interleaved" && chunk%3 == 0 {
+ for _,
value := range values[offset:end] {
+
require.NoError(t, batch.Put(uint64(value)))
+ }
+ } else {
+ n, err
:= batch.PutBatchIndices(values[offset:end])
+
require.NoError(t, err)
+
require.Equal(t, end-offset, n)
+ }
+ offset = end
+ }
+ size := batch.Flush()
+ require.Equal(t,
scalarOutput[:scalarSize], output[:size])
+ decoded :=
make([]uint64, len(values))
+ decoder :=
NewRleDecoder(bytes.NewReader(output[:size]), width)
+ n, err =
decoder.GetBatch(decoded)
+ require.NoError(t, err)
+ require.Equal(t,
len(values), n)
+ for i, value := range
values {
+
require.Equal(t, uint64(value), decoded[i])
+ }
+ }
+ })
+ }
+ })
+ }
+ }
+}
+
+func TestRleBatchIndicesSplitsMaximumRepeatedRun(t *testing.T) {
+ for _, initialCount := range []int32{math.MaxInt32 - 4, math.MaxInt32} {
+ t.Run(fmt.Sprint(initialCount), func(t *testing.T) {
+ output := make([]byte, 32)
+ enc := NewRleEncoder(NewWriterAtBuffer(output), 1)
+ enc.curVal = 1
+ enc.repCount = initialCount
+ n, err := enc.PutBatchIndices([]int32{1, 1, 1, 1, 1, 1})
+ require.NoError(t, err)
+ require.Equal(t, 6, n)
+ want := []byte{0xfe, 0xff, 0xff, 0xff, 0x0f, 1,
byte((int64(initialCount) + 6 - math.MaxInt32) * 2), 1}
+ require.Equal(t, want, output[:enc.Flush()])
+ })
+ }
+}
+
+func TestRleBatchIndicesMaximumRunAtBatchBoundary(t *testing.T) {
+ output := make([]byte, 16)
+ enc := NewRleEncoder(NewWriterAtBuffer(output), 1)
+ enc.curVal = 1
+ enc.repCount = math.MaxInt32 - 4
+ n, err := enc.PutBatchIndices([]int32{1, 1, 1, 1})
+ require.NoError(t, err)
+ require.Equal(t, 4, n)
+ n, err = enc.PutBatchIndices([]int32{0})
+ require.NoError(t, err)
+ require.Equal(t, 1, n)
+ require.Equal(t, []byte{0xfe, 0xff, 0xff, 0xff, 0x0f, 1, 2, 0},
output[:enc.Flush()])
+}
+
+func TestRleBatchIndicesWriteErrors(t *testing.T) {
+ for _, tc := range []struct {
+ name string
+ repeated int32
+ values []int32
+ encoded int
+ }{
+ {"literal header", 0, []int32{0, 1, 0, 1, 0, 1, 0, 1, 1, 1, 1,
1, 1, 1, 1, 1}, 8},
+ {"repeated run", 8, []int32{0}, 0},
+ {"maximum repeated run", math.MaxInt32 - 1, []int32{1, 1}, 1},
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ enc := NewRleEncoder(NewWriterAtBuffer(nil), 1)
+ enc.curVal = 1
+ enc.repCount = tc.repeated
+ n, err := enc.PutBatchIndices(tc.values)
+ require.Error(t, err)
+ require.Equal(t, tc.encoded, n)
+ })
+ }
+}
+
+func TestRleBatchIndicesLiteralWriteErrorsMatchScalar(t *testing.T) {
+ for _, width := range []int{1, 8, 17, 31} {
+ values := make([]int32, 520)
+ mask := int32((uint32(1) << width) - 1)
+ for i := range values {
+ values[i] = int32(i*1234567) & mask
+ }
+ 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) {
+ scalarOutput, batchOutput := make([]byte,
capacity), make([]byte, capacity)
+ scalar :=
NewRleEncoder(NewWriterAtBuffer(scalarOutput), width)
+ scalarCount := 0
+ var scalarErr error
+ for _, value := range values {
+ if scalarErr =
scalar.Put(uint64(value)); scalarErr != nil {
+ break
+ }
+ scalarCount++
+ }
+ batch :=
NewRleEncoder(NewWriterAtBuffer(batchOutput), width)
+ n, err := batch.PutBatchIndices(values)
+ require.Equal(t, scalarErr, err)
+ require.Equal(t, scalarCount, n)
+ require.Equal(t, scalarOutput, batchOutput)
+ require.Equal(t, scalar.buffer, batch.buffer)
+ require.Equal(t, scalar.curVal, batch.curVal)
+ require.Equal(t, scalar.repCount,
batch.repCount)
+ require.Equal(t, scalar.litCount,
batch.litCount)
+ require.Equal(t, scalar.literalIndicatorOffset,
batch.literalIndicatorOffset)
+ })
+ }
+ }
+}
