fix(parquet/metadata): fix BloomFilter memory recycling, doc contract, and benchmark#864
fix(parquet/metadata): fix BloomFilter memory recycling, doc contract, and benchmark#864Patzifist wants to merge 8 commits into
Conversation
… GetColumnBloomFilter
There was a problem hiding this comment.
I ran some testing using your provided test harness configuration and it looks like your originally measured 13GB memory usage is an artifact of your test and not actually an issue in the code.
Your original test contained a make([]byte, 2048) inside the loop. Running that test on main produced ~2 MB/op, with a cumulative alloc of 13,016. The 13GB you measured is 65 * 100 * 2 MB, produced by the original test and not by this library.
If I run the test from this PR against main, I get the same reduction you did because the test in this PR hoists the slice allocation outside of the loop.
Another key point is that if BloomFilterOffset is 0, we trigger an early return and don't even go into the bloom-filter read/pool path at all (because we just return the raw bloom filter bytes as they existed). Your test uses BloomFilterOffset = 0 which means you're never actually exercising the pool path anyways in this test, proving that the original 13GB allocation was just the make which existed in the test and the reduction was solely due to hoisting it out of the loop in the test. Essentially, the allocation measurement test you're adding doesn't actually exercise any of the new code this PR adds.
If I create a proper test that uses a non-zero BloomFilterOffset against main, using a shared reader and sync.Pool recycling then I get about ~11 KB/op with a cumulative allocation of 72 MB. Even in the worst case scenario, removing the explicit returning to the pool and relying solely on the GC to recycle things, I still max out at only ~3GB, far from the 13GB you measured.
Can you try reproducing the memory allocation issue using:
- A non-zero
BloomFilterOffsetso that we actually use the pool instead of just returning the raw bytes deserialized by thrift - Use the same test against both main and against your changes.
Given the analysis I've done and my own tests, I wasn't able to reproduce the original issue you reported in #863
That's not to say there isn't any issue here: because we rely on the GC finalizers to return the buffers, they don't get promptly returned under a burst and we end up over-allocating a bit with some churn (between 1.5 - 4 GB of churn vs ~75 MB if we promptly return the buffers to the pool). So there's definitely something we can fix here, but it's not a memory leak and capacity destruction as you originally thought
| @@ -0,0 +1,162 @@ | |||
| package metadata | |||
There was a problem hiding this comment.
this is missing the apache license header that is required.
| Hasher() Hasher | ||
| CheckHash(hash uint64) bool | ||
| Size() int64 | ||
| Close() |
There was a problem hiding this comment.
this is a breaking change that we shouldn't introduce and shouldn't be needed.
| var ( | ||
| headerPool = make(chan []byte, 512) | ||
| filterPool = make(chan *memory.Buffer, 512) | ||
| ) |
There was a problem hiding this comment.
you're replacing the per-reader, allocator-aware sync.Pool with package level globals. You're bypassing the user-configured allocator (which means that if users are using the mallocator i.e. CGO/non-go memory, unreleased buffers will leak native memory rather than Go memory that can be reclaimed).
| func (b *blockSplitBloomFilter) Close() { | ||
| if b.cancelCleanup != nil { | ||
| b.cancelCleanup() | ||
| b.cancelCleanup = nil | ||
| } | ||
| } |
There was a problem hiding this comment.
this isn't actually called anywhere in the code except for in your new test, so this PR doesn't actually fix anything
|
Ultimately, there is a measurable thing to fix here regarding the latency and churn of memory and how promptly we return the buffer to the pool. But we should improve that without breaking the public For reference, here's the test file I used: const (
reproGoroutines = 65
reproIters = 100
reproOps = reproGoroutines * reproIters
reproBitsetSize = 1 * 1024 * 1024
reproScratch = 2 * 1024 * 1024
reproReadSize = int32(4096)
)
func reproHeaderBytes(t testing.TB) []byte {
t.Helper()
hb, err := thrift.NewThriftSerializer().Write(context.Background(), &format.BloomFilterHeader{
NumBytes: int32(reproBitsetSize),
Algorithm: &defaultAlgorithm,
Hash: &defaultHashStrategy,
Compression: &defaultCompression,
})
if err != nil {
t.Fatalf("serialize bloom filter header: %v", err)
}
return hb
}
// reproFileData lays out [8 byte prefix][thrift header][bitset] so a non-zero
// BloomFilterOffset points at the header.
func reproFileData(t testing.TB) []byte {
t.Helper()
fd := append(make([]byte, 8), reproHeaderBytes(t)...)
return append(fd, make([]byte, reproBitsetSize)...)
}
func reproMeta(t testing.TB, bloomFilterOffset int64) *RowGroupMetaData {
t.Helper()
length := reproReadSize
cmd := format.ColumnMetaData{
Type: format.Type_BYTE_ARRAY,
Encodings: []format.Encoding{format.Encoding_PLAIN},
PathInSchema: []string{"test_col"},
Codec: format.CompressionCodec_UNCOMPRESSED,
BloomFilterOffset: &bloomFilterOffset,
BloomFilterLength: &length,
}
rg := format.RowGroup{Columns: []*format.ColumnChunk{{MetaData: &cmd}}, TotalByteSize: 100, NumRows: 100}
node, err := schema.NewPrimitiveNode("test_col", parquet.Repetition(format.FieldRepetitionType_REQUIRED),
parquet.Type(format.Type_BYTE_ARRAY), -1, -1)
if err != nil {
t.Fatalf("create primitive node: %v", err)
}
root, err := schema.NewGroupNode("schema", parquet.Repetition(format.FieldRepetitionType_REPEATED),
schema.FieldList{node}, -1)
if err != nil {
t.Fatalf("create group node: %v", err)
}
return NewRowGroupMetaData(&rg, schema.NewSchema(root), nil, nil)
}
func reproReader(fileData []byte, meta *RowGroupMetaData, pool *sync.Pool) *RowGroupBloomFilterReader {
return &RowGroupBloomFilterReader{
input: bytes.NewReader(fileData),
sourceFileSize: int64(len(fileData)),
rgMeta: meta,
bufferPool: pool,
}
}
func reproPool() *sync.Pool {
return &sync.Pool{New: func() any {
buf := memory.NewResizableBuffer(memory.NewGoAllocator())
runtime.SetFinalizer(buf, func(o *memory.Buffer) { o.Release() })
return buf
}}
}
// reproRecycle returns the bloom filter's buffer to the pool immediately, doing
// exactly what the reader's runtime.AddCleanup callback does (ResizeNoShrink(0)
// then Put), but synchronously. The registered cleanup is cancelled first so the
// same buffer is not returned twice.
func reproRecycle(bf BloomFilter, pool *sync.Pool) {
b, ok := bf.(*blockSplitBloomFilter)
if !ok || b.data == nil {
return
}
if b.cancelCleanup != nil {
b.cancelCleanup()
}
b.data.ResizeNoShrink(0)
pool.Put(b.data)
}
func reproMeasure(t *testing.T, name string, run func() (read, nilRet int64)) (int64, float64) {
t.Helper()
runtime.GC()
var before runtime.MemStats
runtime.ReadMemStats(&before)
start := time.Now()
read, nilRet := run()
elapsed := time.Since(start)
var after runtime.MemStats
runtime.ReadMemStats(&after)
churnMB := float64(after.TotalAlloc-before.TotalAlloc) / 1024 / 1024
t.Logf("%s\n ops=%d filters-read=%d nil-returns=%d cumulative-alloc(churn)=%.1f MB mallocs=%d elapsed=%v",
name, reproOps, read, nilRet, churnMB, after.Mallocs-before.Mallocs, elapsed)
return read, churnMB
}
func reproRunParallel(work func() BloomFilter) (read, nilRet int64) {
var wg sync.WaitGroup
gate := make(chan struct{})
for range reproGoroutines {
wg.Add(1)
go func() {
defer wg.Done()
<-gate
for range reproIters {
if work() != nil {
atomic.AddInt64(&read, 1)
} else {
atomic.AddInt64(&nilRet, 1)
}
}
}()
}
close(gate)
wg.Wait()
return read, nilRet
}
// Scenario A: the original reproducer. BloomFilterOffset=0 and a fresh 2 MiB file
// buffer allocated per iteration. Allocates ~13 GB while never reading a filter.
func TestBloomFilterPoolRepro_A_OriginalHarness(t *testing.T) {
header := reproHeaderBytes(t)
meta := reproMeta(t, 0)
read, _ := reproMeasure(t, "A original harness: offset=0, fresh 2MiB file per iteration", func() (int64, int64) {
var read, nilRet int64
var wg sync.WaitGroup
gate := make(chan struct{})
for range reproGoroutines {
wg.Add(1)
go func() {
defer wg.Done()
pool := reproPool()
<-gate
for range reproIters {
scratch := make([]byte, reproScratch)
copy(scratch, header)
bf, err := reproReader(scratch, meta, pool).GetColumnBloomFilter(0)
if err != nil || bf == nil {
atomic.AddInt64(&nilRet, 1)
continue
}
atomic.AddInt64(&read, 1)
reproRecycle(bf, pool)
}
}()
}
close(gate)
wg.Wait()
return read, nilRet
})
if read != 0 {
t.Fatalf("expected 0 filters read with offset=0, got %d", read)
}
}
// Scenario B: corrected harness. Non-zero offset (read path runs), shared reader,
// buffer returned to the pool promptly. The 1 MiB bitset buffer is reused.
func TestBloomFilterPoolRepro_B_SharedReaderPromptReturn(t *testing.T) {
fileData := reproFileData(t)
meta := reproMeta(t, 8)
pool := reproPool()
rdr := reproReader(fileData, meta, pool)
bf, err := rdr.GetColumnBloomFilter(0)
if err != nil || bf == nil {
t.Fatalf("sanity read failed: bf=%v err=%v", bf, err)
}
if bf.Size() != int64(reproBitsetSize) {
t.Fatalf("unexpected filter size: got %d want %d", bf.Size(), reproBitsetSize)
}
reproRecycle(bf, pool)
read, churnMB := reproMeasure(t, "B shared reader + prompt return: offset>0", func() (int64, int64) {
return reproRunParallel(func() BloomFilter {
bf, err := rdr.GetColumnBloomFilter(0)
if err != nil || bf == nil {
return nil
}
reproRecycle(bf, pool)
return bf
})
})
if read != reproOps {
t.Fatalf("expected %d filters read, got %d", reproOps, read)
}
if churnMB > 500 {
t.Fatalf("pool not reusing buffers: churn %.1f MB (1 MiB x %d ops would be ~%d MB without reuse)",
churnMB, reproOps, reproOps)
}
}
// Scenario C: corrected harness using the reader's real return path
// (runtime.AddCleanup). GC-timing dependent; reported only.
func TestBloomFilterPoolRepro_C_SharedReaderCleanupOnly(t *testing.T) {
fileData := reproFileData(t)
meta := reproMeta(t, 8)
pool := reproPool()
rdr := reproReader(fileData, meta, pool)
reproMeasure(t, "C shared reader + cleanup-only return: offset>0 (reader's real path)", func() (int64, int64) {
return reproRunParallel(func() BloomFilter {
bf, err := rdr.GetColumnBloomFilter(0)
if err != nil {
return nil
}
return bf
})
})
}
// Scenario D1: a new reader is built every iteration (as in the original harness)
// but the file bytes are shared and the offset is non-zero. A shared pool still
// recycles, so recreating the reader is not what defeated reuse.
func TestBloomFilterPoolRepro_D1_FreshReaderSharedPool(t *testing.T) {
fileData := reproFileData(t)
meta := reproMeta(t, 8)
pool := reproPool()
read, churnMB := reproMeasure(t, "D1 fresh reader + shared pool: offset>0, file shared, prompt return", func() (int64, int64) {
return reproRunParallel(func() BloomFilter {
bf, err := reproReader(fileData, meta, pool).GetColumnBloomFilter(0)
if err != nil || bf == nil {
return nil
}
reproRecycle(bf, pool)
return bf
})
})
if read != reproOps {
t.Fatalf("expected %d filters read, got %d", reproOps, read)
}
if churnMB > 500 {
t.Fatalf("fresh reader defeated shared pool reuse: churn %.1f MB", churnMB)
}
}
// Scenario D2: a new pool is built per read (i.e. a new file.Reader per single
// read). This is the only configuration where per-reader pooling cannot help.
func TestBloomFilterPoolRepro_D2_FreshReaderFreshPool(t *testing.T) {
fileData := reproFileData(t)
meta := reproMeta(t, 8)
reproMeasure(t, "D2 fresh reader + fresh pool per read: offset>0, file shared", func() (int64, int64) {
return reproRunParallel(func() BloomFilter {
bf, err := reproReader(fileData, meta, reproPool()).GetColumnBloomFilter(0)
if err != nil {
return nil
}
return bf
})
})
}
// BenchmarkBloomFilterPoolReuse reports allocs/op for the read path with the
// buffer returned to the pool. A reused 1 MiB bitset does not show up per op.
func BenchmarkBloomFilterPoolReuse(b *testing.B) {
fileData := reproFileData(b)
meta := reproMeta(b, 8)
pool := reproPool()
rdr := reproReader(fileData, meta, pool)
b.ReportAllocs()
b.ResetTimer()
b.RunParallel(func(pb *testing.PB) {
for pb.Next() {
bf, err := rdr.GetColumnBloomFilter(0)
if err != nil || bf == nil {
continue
}
reproRecycle(bf, pool)
}
})
}Running it locally myself with go test -bench I get the following results:
Let's construct a good solution for getting more prompt returns to the pool without breaking the API or defeating the pool. Thoughts? |
|
Hi Zeroshade, Benchmark Results (64 Workers, 100 Iterations, Random NDV Sizes): BenchmarkGetColumnBloomFilter_OriginalSyncPool-12 1 1176984718 ns/op 6280746824 B/op 233435 allocs/op |
|
Test for original check go test -bench=BenchmarkGetColumnBloomFilter_OriginalSyncPool -run=^$ -memprofile=mem.out go tool pprof -http=:8080 mem.out github.com/apache/arrow-go/v18/parquet/metadata.(*RowGroupBloomFilterReader).GetColumnBloomFilter Total: 0 3.19GB (flat, cum) 99.86% |
|
Thanks for the extra benchmark @Patzifist. Before we iterate on the implementation, I want to make sure we're measuring the right thing — I don't think these numbers show what they appear to. What the benchmark is actually measuring Both rows ran at I confirmed this with the
Two related points:
What I do agree is real There is something worth fixing, and it's what we converged on earlier in the thread: today the bitset buffer is only returned to the pool via What would help: a proper reproducer To justify a change of this size, could you put together a benchmark that isolates the claim? Specifically:
If you can share a repro that shows per-call allocations and/or retained-memory growth on current |
|
Following up on my last comment with a concrete alternative direction. I measured the original path against a prompt-return path on current
So the per-reader The only change from your bf, err := threadRdr.GetColumnBloomFilter(0)
if err != nil {
return
}
// return the bitset buffer to the pool immediately instead of waiting for GC
if b, ok := bf.(*blockSplitBloomFilter); ok && b.data != nil {
if b.cancelCleanup != nil {
b.cancelCleanup() // cancel the AddCleanup so the buffer isn't returned twice
}
b.data.ResizeNoShrink(0)
originalArrowPool.Put(b.data)
}Same harness, same pool, same allocator — only the return timing changes, and that alone accounts for the 5.1 GB → 0.40 GB / 3.0 GB → 83 MB difference. (There's also a real pool-pollution issue: ProposalKeep Option A — scoped accessor (preferred): func (r *RowGroupBloomFilterReader) VisitColumnBloomFilter(i int, fn func(BloomFilter) error) error {
bf, err := r.GetColumnBloomFilter(i)
if err != nil || bf == nil {
return err
}
defer r.recycle(bf) // cancelCleanup() + data.ResizeNoShrink(0) + pool.Put(data)
return fn(bf)
}The buffer returns to the pool the instant the callback ends — deterministic, lexically-scoped lifetime, hard to misuse. Option B — explicit recycle: bf, err := rg.GetColumnBloomFilter(i)
// ... use bf ...
rg.RecycleBloomFilter(bf) // cancels the cleanup, ResizeNoShrink(0), Put back to poolMore flexible for callers who can't wrap usage in a closure, but easier to forget (forgetting just falls back to today's GC behavior, so it's safe, only slower). Both call the same helper that cancels the registered cleanup (so the buffer isn't double-returned) and puts it back. Either can pair with giving the bloom reader its own per-reader pool, separate from the page-buffer pool, to fix the pollution. Why this over the current PRThe PR reaches a similar steady-state, but the way it gets there trips each of the constraints I called out earlier:
Net: same bounded-memory result, no interface break, no allocator bypass, no global pools, and the existing safety net stays. Happy to put up a PR with Option A + the dedicated pool if that direction works for you. |
Introduce BenchmarkVisitColumnBloomFilter_SyncPool to measure the performance and memory allocations of visiting column Bloom filters when utilizing a sync.Pool. Benchmark results on Intel i7-1355U: - Operations: 9330 - Speed: 114338 ns/op - Memory: 13946 B/op - Allocations: 17 allocs/op
|
Based on your feedback, I have updated the PR to implement Option A (scoped accessor):
Please take a look when you have a moment! |
Introduce BenchmarkVisitColumnBloomFilter_SyncPool to measure the performance and memory allocations of visiting column Bloom filters when utilizing a sync.Pool. Benchmark results on Intel i7-1355U: - Operations: 9330 - Speed: 114338 ns/op - Memory: 13946 B/op - Allocations: 17 allocs/op
zeroshade
left a comment
There was a problem hiding this comment.
Thanks for reworking this to the scoped-accessor approach — this is a big improvement. It resolves the earlier concerns: GetColumnBloomFilter and its runtime.AddCleanup safety net are intact, the public BloomFilter interface is untouched (no Close()), the process-global pools and the make([]byte) allocator bypass are gone, and the benchmark now measures a single filter read. I built the branch and ran BenchmarkVisitColumnBloomFilter_SyncPool: ~17 allocs/op, ~3.8 KB/op — the ~1 MB bitset is reused across calls exactly as intended.
Two things need to be fixed before this can merge (both marked inline), plus a few minor items:
- [blocker] Missing ASF license header on the new
bloom_filter_benchmark_test.go. Apache RAT (dev/release/run_rat.sh) and the pre-commit license hook will fail CI. - [blocker]
recyclereturns the encrypted path's buffer to the pool, which pushes decrypted data into the shared pool and panics on the next reuse (nil allocator). Reproducer is inline.
Still open, but fine as a follow-up (not blocking this PR): the bloom reader still shares file.Reader.bufferPool with the column/record readers (parquet/file/file_reader.go), so the small-header vs MB-bitset pool pollution remains.
Details inline.
- Fixed `RowGroupBloomFilterReader.recycle` to only put buffers back into the sync.Pool if they were allocated by it (checking `cancelCleanup`). This prevents pool corruption and panics on the encrypted path. - Fixed `prepareTestHarness` in benchmarks to dynamically advance the payload offset and assign unique pointers, ensuring accurate multi-size testing instead of overwriting the same memory location. - Corrected `BloomFilterLength` in benchmarks to prevent slice out of bounds panic.
zeroshade
left a comment
There was a problem hiding this comment.
Re-review of the latest commit (074bfed). Thanks for pushing this forward — the two prior blockers are resolved:
- ✅ ASF license header added to
bloom_filter_benchmark_test.go. - ✅ Encrypted path no longer pushes a nil-allocator buffer into the pool —
recyclegates oncancelCleanupandRelease()s otherwise (andRelease()on aNewBufferBytesbuffer is a safe no-op, sincemem/parentare both nil).
Minor items also addressed: ✅ lifetime-contract doc comment on VisitColumnBloomFilter, ✅ dead _Channels benchmark removed, ✅ benchmark now gives each meta its own offset/size.
One new blocker in that same recycle gating (inline): it nils cancelCleanup instead of calling it, so the GC cleanup registered by addCleanup stays armed and returns the buffer to the pool a second time — a double-Put (same buffer handed to two goroutines) on Go 1.24+, and a finalizer nil-deref on Go ≤1.23. CI doesn't run the benchmark under -race, and the second Put only fires after a GC, so the suite stays green while the corruption stays latent.
Non-blocking: this PR adds VisitColumnBloomFilter but nothing in the library calls it, so the win only lands for callers who opt into the new API; GetColumnBloomFilter users still rely on GC recycling. Fine if that's the intent.
Details inline.
Co-authored-by: Matt Topol <zotthewizard@gmail.com>
…BloomFilter` upper defer block back to the original `headerBuf.ResizeNoShrink(0)`. This guarantees that the resource recycling and array length scrubbing styles remain strictly uniform and consistent across the entire metadata codebase.
|
Thank you for the guidance and the direct patch pointers, |
Rationale for this change
Previously, this package relied entirely on the Go Garbage Collector (
runtime.AddCleanup) to return buffers back to thesync.Pool. Under heavy or concurrent load, this caused uncontrolled memory growth because new buffers were being allocated much faster than the GC could clean them up and return them to the pool.To fix this memory leak, an explicit synchronous recycling model was introduced by calling
defer r.recycle(bf)insideVisitColumnBloomFilter. This ensures immediate buffer reuse.What changes are included in this PR?
parquet/metadata/bloom_filter.go:VisitColumnBloomFilterto stop uncontrolled memory growth and ensure immediate pool reuse.Are these changes tested?
Yes. Verified that the memory leak is resolved, added the regression test suite, and successfully ran the parallel benchmark:
go test -bench=^BenchmarkVisitColumnBloomFilter_SyncPool$ -run=^$ -benchmem ./parquet/metadata/...Output:
Are there any user-facing changes?
No. This is strictly an internal performance, stability, and memory footprint optimization.