perf(parquet): use typed memo insertion for byte-array dictionaries

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

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 {

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)
+		})
+	}
+}