From 2ec735cf2e3e4b4e9207254be80978a106a495f4 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Fri, 28 Aug 2026 10:47:12 +0200 Subject: [PATCH] perf(parquet): reuse dictionary RLE decoder --- parquet/internal/encoding/decoder.go | 15 +++- .../encoding/encoding_benchmarks_test.go | 71 +++++++++++++++++++ parquet/internal/encoding/encoding_test.go | 50 +++++++++++++ 3 files changed, 134 insertions(+), 2 deletions(-) diff --git a/parquet/internal/encoding/decoder.go b/parquet/internal/encoding/decoder.go index 9c124da29..456bb8d2e 100644 --- a/parquet/internal/encoding/decoder.go +++ b/parquet/internal/encoding/decoder.go @@ -107,6 +107,7 @@ type dictDecoder[T parquet.ColumnTypes] struct { mem memory.Allocator dictValueDecoder utils.DictionaryConverter[T] idxDecoder *utils.TypedRleDecoder[T] + idxReader bytes.Reader idxScratchSpace []uint64 } @@ -126,7 +127,7 @@ func (d *dictDecoder[T]) SetData(nvals int, data []byte) error { d.nvals = nvals if len(data) == 0 { // no data, bitwidth can safely be 0 - d.idxDecoder = utils.NewTypedRleDecoder[T](bytes.NewReader(data), 0 /* bitwidth */) + d.resetIndexDecoder(data, 0 /* bitwidth */) return nil } @@ -137,10 +138,20 @@ func (d *dictDecoder[T]) SetData(nvals int, data []byte) error { } // pass the rest of the data, minus that first byte, to the decoder - d.idxDecoder = utils.NewTypedRleDecoder[T](bytes.NewReader(data[1:]), int(width)) + d.resetIndexDecoder(data[1:], int(width)) return nil } +func (d *dictDecoder[T]) resetIndexDecoder(data []byte, width int) { + d.idxReader.Reset(data) + if d.idxDecoder == nil { + d.idxDecoder = utils.NewTypedRleDecoder[T](&d.idxReader, width) + return + } + + d.idxDecoder.Reset(&d.idxReader, width) +} + func (d *dictDecoder[T]) discard(n int) (int, error) { n = d.idxDecoder.Discard(n) d.nvals -= n diff --git a/parquet/internal/encoding/encoding_benchmarks_test.go b/parquet/internal/encoding/encoding_benchmarks_test.go index 781967a96..eb4e7e3a4 100644 --- a/parquet/internal/encoding/encoding_benchmarks_test.go +++ b/parquet/internal/encoding/encoding_benchmarks_test.go @@ -544,6 +544,77 @@ func BenchmarkDecodeDictByteArray(b *testing.B) { } } +func BenchmarkDecodeDictByteArrayReuse(b *testing.B) { + const nunique = 100 + + for _, nvalues := range []int{64, 1024, 8192} { + b.Run(fmt.Sprintf("values=%d", nvalues), func(b *testing.B) { + rag := testutils.NewRandomArrayGenerator(0) + dict := rag.ByteArray(nunique, 32, 32, 0).(*array.String) + indices := rag.Int32(int64(nvalues), 0, nunique-1, 0) + + values := make([]parquet.ByteArray, nvalues) + for idx := range values { + values[idx] = []byte(dict.Value(int(indices.Value(idx)))) + } + + col := schema.NewColumn(schema.NewByteArrayNode("bytearray", parquet.Repetitions.Required, -1), 0, 0) + enc := encoding.NewEncoder(parquet.Types.ByteArray, parquet.Encodings.PlainDict, true, col, memory.DefaultAllocator).(*encoding.DictByteArrayEncoder) + enc.Put(values) + + dictBuf := make([]byte, enc.DictEncodedSize()) + enc.WriteDict(dictBuf) + idxBuf := make([]byte, enc.EstimatedDataEncodedSize()) + n, err := enc.WriteIndices(idxBuf) + if err != nil { + b.Fatal(err) + } + idxBuf = idxBuf[:n] + + dec := encoding.NewDecoder(parquet.Types.ByteArray, parquet.Encodings.Plain, col, memory.DefaultAllocator) + if err := dec.SetData(nunique, dictBuf); err != nil { + b.Fatal(err) + } + dictDec := encoding.NewDictDecoder(parquet.Types.ByteArray, col, memory.DefaultAllocator).(*encoding.DictByteArrayDecoder) + dictDec.SetDict(dec) + out := make([]parquet.ByteArray, nvalues) + + b.SetBytes(int64(nvalues)) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if err := dictDec.SetData(nvalues, idxBuf); err != nil { + b.Fatal(err) + } + decoded, err := dictDec.Decode(out) + if err != nil { + b.Fatal(err) + } + if decoded != nvalues { + b.Fatalf("decoded %d values, want %d", decoded, nvalues) + } + } + }) + } +} + +func BenchmarkDictDecoderSetDataReuse(b *testing.B) { + column := schema.NewColumn(schema.NewInt32Node("int32", parquet.Repetitions.Required, -1), 0, 0) + decoder := encoding.NewDictDecoder(parquet.Types.Int32, column, memory.DefaultAllocator) + data := []byte{1} + if err := decoder.SetData(1024, data); err != nil { + b.Fatal(err) + } + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if err := decoder.SetData(1024, data); err != nil { + b.Fatal(err) + } + } +} + func BenchmarkByteStreamSplitEncodingInt32(b *testing.B) { for sz := MINSIZE; sz < MAXSIZE+1; sz *= 2 { b.Run(fmt.Sprintf("len %d", sz), func(b *testing.B) { diff --git a/parquet/internal/encoding/encoding_test.go b/parquet/internal/encoding/encoding_test.go index 4b5273322..1589c1332 100644 --- a/parquet/internal/encoding/encoding_test.go +++ b/parquet/internal/encoding/encoding_test.go @@ -628,6 +628,56 @@ func TestDictionaryDecoderRejectsTruncatedPackedLiteralRun(t *testing.T) { assert.ErrorIs(t, err, io.ErrUnexpectedEOF) } +func encodeDictionaryIndexPage(t *testing.T, values []uint64, width int) []byte { + t.Helper() + + data := make([]byte, 1+utils.MaxRLEBufferSize(width, len(values))+utils.MinRLEBufferSize(width)) + data[0] = byte(width) + encoder := utils.NewRleEncoder(utils.NewWriterAtBuffer(data[1:]), width) + for _, value := range values { + require.NoError(t, encoder.Put(value)) + } + return data[:encoder.Flush()+1] +} + +func TestDictionaryDecoderResetsStateBetweenPages(t *testing.T) { + column := schema.NewColumn(schema.NewInt32Node("int32", parquet.Repetitions.Required, -1), 0, 0) + dictionaryData := make([]byte, 8) + binary.LittleEndian.PutUint32(dictionaryData, 10) + binary.LittleEndian.PutUint32(dictionaryData[4:], 20) + + dictionary := encoding.NewDecoder(parquet.Types.Int32, parquet.Encodings.Plain, column, memory.DefaultAllocator) + require.NoError(t, dictionary.SetData(2, dictionaryData)) + + decoder := encoding.NewDictDecoder(parquet.Types.Int32, column, memory.DefaultAllocator) + decoder.SetDict(dictionary) + + firstPageValues := make([]uint64, 16) + firstPage := encodeDictionaryIndexPage(t, firstPageValues, 1) + require.NoError(t, decoder.SetData(len(firstPageValues), firstPage)) + + firstOutput := make([]int32, 3) + decoded, err := decoder.(encoding.Int32Decoder).Decode(firstOutput) + require.NoError(t, err) + assert.Equal(t, len(firstOutput), decoded) + assert.Equal(t, []int32{10, 10, 10}, firstOutput) + assert.Equal(t, len(firstPageValues)-len(firstOutput), decoder.ValuesLeft()) + + require.NoError(t, decoder.SetData(1, nil)) + _, err = decoder.(encoding.Int32Decoder).Decode(make([]int32, 1)) + assert.Error(t, err) + + secondPageValues := []uint64{1, 0, 1, 1, 0, 1, 0, 1} + secondPage := encodeDictionaryIndexPage(t, secondPageValues, 1) + require.NoError(t, decoder.SetData(len(secondPageValues), secondPage)) + + secondOutput := make([]int32, len(secondPageValues)) + decoded, err = decoder.(encoding.Int32Decoder).Decode(secondOutput) + require.NoError(t, err) + assert.Equal(t, len(secondOutput), decoded) + assert.Equal(t, []int32{20, 10, 20, 20, 10, 20, 10, 20}, secondOutput) +} + func TestWriteDeltaBitPackedInt32(t *testing.T) { column := schema.NewColumn(schema.NewInt32Node("int32", parquet.Repetitions.Required, -1), 0, 0)