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

[derekperkins]

Rationale for this change

DictByteArrayEncoder.PutByteArray inserts through the untyped MemoTable.GetOrInsert(interface{}), which boxes the parquet.ByteArray on every value written. Boxing a slice is a heap allocation (runtime.convTslice), and the memo table discards it immediately:

// parquet/internal/encoding/byte_array_encoder.go
func (enc *DictByteArrayEncoder) PutByteArray(in parquet.ByteArray) {
	memoIdx, found, err := enc.memo.GetOrInsert(in)   // convTslice per value

hashing.BinaryMemoTable already implements the allocation-free typed entry point, and GetOrInsert is a thin boxing wrapper over it:

func (b *BinaryMemoTable) GetOrInsert(val interface{}) (int, bool, error) {
	return b.InsertOrGet(b.valAsByteSlice(val))
}

The encoder can't reach it, because encoding.BinaryMemoTable doesn't list InsertOrGet among its methods.

This is the same change already made for the numeric paths in #1178 and #1251. #1178 notes that it "keep[s] the existing byte-array and fixed-length byte-array paths unchanged", so this is the remaining half of that work rather than a new direction.

Benchmark — BenchmarkEncodeDictByteArray, already in the tree (65,535 values, 100 unique, 8–32 byte strings). benchstat, n=10, Apple M3 Max, Go 1.27, against main at 7efe1c0:

	before	after	change
sec/op	3.266m ± 1%	2.495m ± 2%	−23.61% (p=0.000)
B/op	3.551Mi ± 0%	2.051Mi ± 0%	−42.26% (p=0.000)
allocs/op	131.12k ± 0%	65.58k ± 0%	−49.98% (p=0.000)

The allocation delta is exactly 65,536 — one per value, plus one.

The benchmark understates the production effect, because its 100 distinct values keep the memo table small. In a low-cardinality column the boxing dominates: every value allocates, and every value is then found to be a duplicate. We hit this writing Iceberg tables through iceberg-go, which writes via pqarrow. In a four-hour production CPU and heap profile of a single streaming writer, DictByteArrayEncoder.PutByteArray was the fifth-largest allocation site in the whole process — 25.5M objects, 7.7% of everything allocated, all of it this one site. typedDictEncoder[int64].Put was number one in the same profile at 15.1% before #1178 landed; together the two accounted for roughly a quarter of the process's allocations, which showed up as ~12% of CPU in GC mark.

What changes are included in this PR?

• Add InsertOrGet(val []byte) to the encoding.BinaryMemoTable interface.
• Call it from DictByteArrayEncoder.PutByteArray instead of GetOrInsert.
• Add InsertOrGet to binaryMemoTableImpl so it still satisfies the interface.
• Add TestBinaryInsertOrGet, covering both implementations.

Notes:

• encoding.BinaryMemoTable lives under parquet/internal/, so widening it is not a public API change.
• The production implementation (hashing.BinaryMemoTable, via NewBinaryDictionary) already satisfied the wider interface with no changes.
• The only other implementation is binaryMemoTableImpl, which the source marks deprecated and benchmark-only ("will be removed in a future release"); the method added there mirrors its existing GetOrInsert.
• DictFixedLenByteArrayEncoder has the same pattern. I left it out to keep this focused and because I have no production numbers for it — happy to follow up.

Are these changes tested?

Yes. TestBinaryInsertOrGet runs against both BinaryMemoTable implementations and checks index assignment, the found flag on re-insertion, nil treated as the empty value, agreement with GetOrInsert, and that stored values do not alias the caller's buffer. I confirmed the test fails when the implementation is deliberately broken.

• go build ./parquet/...
• go vet ./parquet/internal/encoding/
• go test ./parquet/internal/encoding/... — pass
• go test -race ./parquet/internal/encoding/... — pass
• go test ./parquet/pqarrow/... ./parquet/file/... — the only failures are pre-existing and byte-identical to unpatched main (the parquet-testing submodule data is not checked out locally); verified by stashing the patch and re-running
• gofmt -l clean, git diff --check clean

Benchmark command:

go test ./parquet/internal/encoding -run '^$' -bench '^BenchmarkEncodeDictByteArray$' -benchmem -benchtime=500ms -count=10

Are there any user-facing changes?

No. The modified interface is in an internal package, and encoded output is unchanged — only the insertion path differs.

[derekperkins]

I think those are flaky tests, they're not related to this PR. I don't have the rights to restart the failed jobs

[zeroshade]

Nice work, and thank you for the unusually complete write-up — the production profile context made this easy to evaluate. Approving.

I reproduced your benchmark independently (Linux, Intel Ultra 7 155H, benchstat n=12, interleaved against 7efe1c01):

	before	after	change
sec/op	5.346m ± 30%	3.869m ± 17%	−27.62% (p=0.001)
B/op	3.552Mi ± 0%	2.051Mi ± 0%	−42.26% (p=0.000)
allocs/op	131.12k ± 0%	65.58k ± 0%	−49.98% (p=0.000)

B/op and allocs/op match your numbers to the digit, and the delta is exactly one allocation per value as you said. My box is too noisy to confirm your −23.61% precisely, but the direction and magnitude agree.

Also verified: go build ./..., go vet, full go test ./..., and -race on parquet/internal/encoding and internal/hashing all pass. On the interface widening — the only two implementers of encoding.BinaryMemoTable are hashing.BinaryMemoTable (already satisfied it) and binaryMemoTableImpl (handled here). Every other reference in the repo is to the concrete *hashing.BinaryMemoTable, so nothing else is affected. The new test is good: covering both implementations, nil, and non-aliasing is more than I expected.

There is a second boxing allocation, and it's a one-liner

You removed one of two boxes per value. The other is still there, in hashing.BinaryMemoTable.InsertOrGet itself:

func (b *BinaryMemoTable) InsertOrGet(val []byte) (idx int, found bool, err error) {
	h := b.getHash(val)   // getHash(interface{}) — boxes the []byte again

I isolated it against the obvious alternative suspect, the comparison closure in lookup (allocs/op on the already-present path):

method	hash call	closure?	allocs/op
InsertOrGet	getHash(interface{})	yes	1
GetOrInsertBytes	Hash(val, 0)	yes	0
ExistsDirect	Hash(val, 0)	no	0
Exists	getHash(interface{})	yes	1

GetOrInsertBytes shares the same closure and costs nothing, so it is the boxing, not the closure.

GetOrInsertBytes sits two functions above InsertOrGet and is otherwise identical. Changing that one line to h := Hash(val, 0) is semantically a no-op — getHash's []byte case is Hash(v, 0), and the method only accepts []byte. On top of your branch, benchstat n=12:

	this PR	+ Hash(val, 0)	change
sec/op	5.734m ± 30%	2.570m ± 40%	−55.17% (p=0.000)
B/op	2100.2Ki ± 0%	563.2Ki ± 0%	−73.18% (p=0.000)
allocs/op	65579 ± 0%	44 ± 0%	−99.93% (p=0.000)

65,579 → 44. Full go test ./... passes with it. Given your production profile, that looks worth having. Fold it in here or do it as a follow-up, whichever you prefer — this PR stands on its own either way. Exists has the same one-line fix available.

Minor

binaryMemoTableImpl.InsertOrGet allocates once per call. I measured 1 alloc/op against 0 for its GetOrInsert, because key := string(val) copies even on the found path; the compiler only elides that for m[string(b)] used directly as an index expression. It's the deprecated impl, but it exists for benchmark comparison, so an extra allocation there quietly skews exactly those comparisons:

func (m *binaryMemoTableImpl) InsertOrGet(val []byte) (idx int, found bool, err error) {
	if idx, found = m.table[string(val)]; !found {   // no alloc on the found path
		idx = m.Size()
		key := string(val)
		m.builder.AppendString(key)
		m.table[key] = idx
	}
	return
}

Worth saying that your version is safer than the existing GetOrInsert regardless: valAsString casts with unsafe, so the map key aliases the caller's buffer and mutating it afterwards corrupts the table. Your test pins the copy behavior.

Type assertion placement (no action needed). enc.memo.(BinaryMemoTable) is inside PutByteArray, which Put and PutSpaced both call per value, whereas #1178 hoists it out of the loop. I built the hoisted variant expecting a win and measured none (p=0.713, n=12) — flagging it only so you know it was checked, and because matching #1178's shape has some readability value.

FLBA follow-up. Agreed on deferring; for whoever picks it up, it's fixed_len_byte_array_encoder.go:138 and :202, both inside loops.