From d0e337f8d882fe69b14c76904bdc2c35827a653f Mon Sep 17 00:00:00 2001 From: Derek Perkins Date: Mon, 31 Aug 2026 15:26:05 -0600 Subject: [PATCH] perf(parquet): use typed memo insertion for byte-array dictionaries DictByteArrayEncoder.PutByteArray inserted through the untyped MemoTable.GetOrInsert, boxing the parquet.ByteArray on every value written. Boxing a slice is a heap allocation that the memo table discards immediately. hashing.BinaryMemoTable already implements the allocation-free InsertOrGet, and GetOrInsert is a boxing wrapper over it. Expose InsertOrGet on the encoding.BinaryMemoTable interface so the encoder can call it directly. This applies the treatment from #1178 and #1251 to the byte-array path that #1178 explicitly left unchanged. BenchmarkEncodeDictByteArray, benchstat n=10: -23.61% sec/op, -42.26% B/op, -49.98% allocs/op. The allocation delta is exactly one per value. Co-Authored-By: Claude Opus 5 --- .../internal/encoding/byte_array_encoder.go | 2 +- parquet/internal/encoding/memo_table.go | 14 ++++ parquet/internal/encoding/memo_table_test.go | 69 +++++++++++++++++++ 3 files changed, 84 insertions(+), 1 deletion(-) diff --git a/parquet/internal/encoding/byte_array_encoder.go b/parquet/internal/encoding/byte_array_encoder.go index cb80fbaba..95ea3e11d 100644 --- a/parquet/internal/encoding/byte_array_encoder.go +++ b/parquet/internal/encoding/byte_array_encoder.go @@ -106,7 +106,7 @@ func (enc *DictByteArrayEncoder) WriteDict(out []byte) { // PutByteArray adds a single byte array to buffer, updating the dictionary // and encoded size if it's a new value func (enc *DictByteArrayEncoder) PutByteArray(in parquet.ByteArray) { - memoIdx, found, err := enc.memo.GetOrInsert(in) + memoIdx, found, err := enc.memo.(BinaryMemoTable).InsertOrGet(in) if err != nil { panic(err) } diff --git a/parquet/internal/encoding/memo_table.go b/parquet/internal/encoding/memo_table.go index 062104b6b..a0cab8511 100644 --- a/parquet/internal/encoding/memo_table.go +++ b/parquet/internal/encoding/memo_table.go @@ -99,6 +99,9 @@ type TypedMemoTable[T hashing.MemoTypes] interface { // for handling byte arrays/strings/fixed length byte arrays. type BinaryMemoTable interface { MemoTable + // InsertOrGet is the typed equivalent of MemoTable.GetOrInsert, avoiding the + // interface boxing of the value on every call. + InsertOrGet(val []byte) (idx int, found bool, err error) // ValuesSize returns the total number of bytes needed to copy all of the values // from this table. ValuesSize() int @@ -240,6 +243,17 @@ func (m *binaryMemoTableImpl) GetOrInsert(val interface{}) (idx int, found bool, return } +func (m *binaryMemoTableImpl) InsertOrGet(val []byte) (idx int, found bool, err error) { + key := string(val) + idx, found = m.table[key] + if !found { + idx = m.Size() + m.builder.AppendString(key) + m.table[key] = idx + } + return +} + func (m *binaryMemoTableImpl) GetOrInsertNull() (idx int, found bool) { idx, found = m.GetNull() if !found { diff --git a/parquet/internal/encoding/memo_table_test.go b/parquet/internal/encoding/memo_table_test.go index 348cda76c..fb2b27d4c 100644 --- a/parquet/internal/encoding/memo_table_test.go +++ b/parquet/internal/encoding/memo_table_test.go @@ -291,3 +291,72 @@ func (m *MemoTableTestSuite) TestBinaryEmpty() { table.CopyOffsetsSubset(0, offsets) m.Equal(int32(0), offsets[0]) } + +// InsertOrGet is the typed entry point used by the byte-array dictionary +// encoders. It must agree with GetOrInsert on index assignment and on whether +// the value already existed, for every BinaryMemoTable implementation. +func (m *MemoTableTestSuite) TestBinaryInsertOrGet() { + const ( + A = "" + B = "a" + C = "foo" + D = "\000" + E = "\000trailing" + ) + + for _, tt := range []struct { + name string + table func() encoding.BinaryMemoTable + }{ + {"hashing", func() encoding.BinaryMemoTable { + return encoding.NewBinaryDictionary(memory.DefaultAllocator) + }}, + {"legacy", func() encoding.BinaryMemoTable { + return encoding.NewBinaryMemoTable(memory.DefaultAllocator) + }}, + } { + m.Run(tt.name, func() { + table := tt.table() + defer table.Release() + + for idx, val := range []string{A, B, C, D, E} { + got, found, err := table.InsertOrGet([]byte(val)) + m.Require().NoError(err) + m.False(found, "value %q should be inserted, not found", val) + m.Equal(idx, got) + } + m.Equal(5, table.Size()) + + // Re-inserting must return the original index and report found. + for idx, val := range []string{A, B, C, D, E} { + got, found, err := table.InsertOrGet([]byte(val)) + m.Require().NoError(err) + m.True(found, "value %q should already exist", val) + m.Equal(idx, got) + } + m.Equal(5, table.Size()) + + // A nil slice is the empty value, which was inserted first. + got, found, err := table.InsertOrGet(nil) + m.Require().NoError(err) + m.True(found) + m.Equal(0, got) + + // InsertOrGet and GetOrInsert must agree. + got, found, err = table.GetOrInsert([]byte(C)) + m.Require().NoError(err) + m.True(found) + m.Equal(2, got) + + // The inserted value must not alias the caller's buffer. + buf := []byte("mutable") + inserted, _, err := table.InsertOrGet(buf) + m.Require().NoError(err) + buf[0] = 'X' + again, found, err := table.InsertOrGet([]byte("mutable")) + m.Require().NoError(err) + m.True(found, "stored value must be a copy, not a view of the caller's slice") + m.Equal(inserted, again) + }) + } +}