apache · fallintoplace · Aug 28, 2026
@@ -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

@@ -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) {

@@ -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)