This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new 0c0cd4e5 perf(parquet): use typed memo insertion for byte-array
dictionaries (#1272)
0c0cd4e5 is described below
commit 0c0cd4e5ef6417f33d2fac2cd79838afd62df31a
Author: Derek Perkins <[email protected]>
AuthorDate: Wed Sep 2 10:35:50 2026 -0600
perf(parquet): use typed memo insertion for byte-array dictionaries (#1272)
### 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:
```go
// 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:
```go
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.
Co-authored-by: Claude Opus 5 <[email protected]>
---
parquet/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 cb80fbab..95ea3e11 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 062104b6..a0cab851 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 348cda76..fb2b27d4 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)
+ })
+ }
+}
