perf: pool hash tables in Go encode paths to reduce allocations

[huynhanx03]

Summary

Pool S2 Go-path hash tables (320KB–4.6MB/call) via sync.Pool, matching the existing amd64 asm pattern (encPools in encode_amd64.go). Eliminates per-block stack growth and GC pressure on arm64.

Changes

• s2/hashtable_pool.go (new): Pooled structs + shared size constants.
• s2/encode_better.go: encodeBlockBetterGo, encodeBlockBetterSnappyGo, encodeBlockBetterDict use pooled tables.
• s2/encode_best.go: encodeBlockBestGo, encodeBlockBestSnappy use pooled tables.

Benchmark

goos: darwin
goarch: arm64
cpu: Apple M1

name                                               old ns/op    new ns/op     delta
RandomEncodeBetterBlock16MB/better-8                 810,000      770,000     -4.9%
RandomEncodeBetterBlock16MB/best-8                34,000,000   32,000,000     -5.9%
RandomEncodeBetterBlock16MB/snappy-better-8        1,900,000    1,800,000     -5.3%

name                                                old B/op     new B/op     delta
RandomEncodeBetterBlock16MB/better-8                 524,288          380    -99.9%
RandomEncodeBetterBlock16MB/best-8                 4,718,592            0    -99.9%
RandomEncodeBetterBlock16MB/snappy-better-8          262,144            0    -99.9%

name                                              old allocs   new allocs     delta
RandomEncodeBetterBlock16MB/better-8                       1            0     -100%
RandomEncodeBetterBlock16MB/best-8                         2            0     -100%
RandomEncodeBetterBlock16MB/snappy-better-8                1            0     -100%

Checklist

• Tests, race detector, vet pass

Summary by CodeRabbit

• Refactor
  • Compression internals now use pooled, reusable hash-table storage and shared sizing parameters across variants.
  • Reduced stack usage and runtime allocations, improving memory efficiency and consistency across encoding paths.
  • Expected to yield more stable performance with lower GC pressure; no public API changes.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 1661650a-fb3c-4545-827c-957f20cda0c0

📥 Commits

Reviewing files that changed from the base of the PR and between 577674e and e967be7.

📒 Files selected for processing (3)

• s2/encode_best.go
• s2/encode_better.go
• s2/hashtable_pool.go

✅ Files skipped from review due to trivial changes (1)

• s2/hashtable_pool.go

🚧 Files skipped from review as they are similar to previous changes (1)

• s2/encode_best.go

---

📝 Walkthrough

Walkthrough

Replaces per-call stack-allocated hash-table arrays in S2 encoders with pooled, preallocated table structs; introduces pool-backed table types and shared sizing constants, and updates encoder functions to borrow and return tables via pool accessors.

Changes

Cohort / File(s)	Summary
Hash Table Pool Infrastructure s2/hashtable_pool.go	New file adding sizing constants, three table struct types (betterTables, betterSnappyTables, bestTables), three sync.Pool instances, and get*Tables() helpers that return zeroed pooled table pointers.
Pooled Table Usage — Best encoder s2/encode_best.go	Replaced local stack arrays with pooled bestTables via getBestTables(), switched to shared best*TableBits/Size constants, and return tables to bestTablePool with defer.
Pooled Table Usage — Better encoder s2/encode_better.go	Replaced local stack arrays in encodeBlockBetterGo / encodeBlockBetterSnappyGo / encodeBlockBetterDict with pooled betterTables/betterSnappyTables, adopted shared sizing constants, and deferred pool.Put for returned tables.

Sequence Diagram(s)

sequenceDiagram
    participant Encoder as Encoder (encodeBlock*)
    participant Pool as TablePool (sync.Pool)
    participant Tables as Tables (better/best struct)

    Encoder->>Pool: get*Tables()
    Pool-->>Tables: return pooled *Tables
    Encoder->>Tables: use lTable / sTable for hashing
    Encoder--)Tables: finish using tables
    Encoder->>Pool: Pool.Put(tbl)

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

• s2: Don't use stack for index tables #1014: Implements same pattern of replacing stack-allocated index/hash tables with sync.Pool-backed pooled tables/buffers.
• s2: Improve small block compression speed w/o asm #1048: Related adjustments to S2 encoder hash-table sizing and allocation used across encode_better/encode_best.

Suggested reviewers

• klauspost

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The pull request title 'perf: pool hash tables in Go encode paths to reduce allocations' clearly and concisely describes the main change: introducing pooling of hash tables in Go encode paths to reduce memory allocations.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

s2/hashtable_pool.go (1)

23-42: Consider wrapping pool checkout/reset in helpers.

Get + full reset + defer Put is now repeated in s2/encode_better.go at Lines 70-75, 332-337, and 932-937, and in s2/encode_best.go at Lines 47-52 and 481-486. A small helper per pool that resets the whole struct before handing it out would keep future table-layout changes from drifting out of sync.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@s2/hashtable_pool.go` around lines 23 - 42, Create helper getters for each
pool (e.g., getBetterTables, getBetterSnappyTables, getBestTables) that
encapsulate pool.Get(), cast to the correct pointer type (betterTables,
betterSnappyTables, bestTables), reset the entire struct to zero (e.g., assign a
zero-value struct to *ptr) and return the pointer; replace the repeated pattern
of betterTablePool.Get()+manual reset+defer Put with calls to these helpers so
callers receive an already-zeroed table object and still call Put when done. Use
the pool variables betterTablePool, betterSnappyTablePool, bestTablePool and the
struct types betterTables, betterSnappyTables, bestTables to locate where to add
the helpers.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@s2/hashtable_pool.go`:
- Around line 23-42: Create helper getters for each pool (e.g., getBetterTables,
getBetterSnappyTables, getBestTables) that encapsulate pool.Get(), cast to the
correct pointer type (betterTables, betterSnappyTables, bestTables), reset the
entire struct to zero (e.g., assign a zero-value struct to *ptr) and return the
pointer; replace the repeated pattern of betterTablePool.Get()+manual
reset+defer Put with calls to these helpers so callers receive an already-zeroed
table object and still call Put when done. Use the pool variables
betterTablePool, betterSnappyTablePool, bestTablePool and the struct types
betterTables, betterSnappyTables, bestTables to locate where to add the helpers.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: d6a14610-8ea5-4caa-99f4-8066f002aca8

📥 Commits

Reviewing files that changed from the base of the PR and between 15967de and 8f6a238.

📒 Files selected for processing (3)

• s2/encode_best.go
• s2/encode_better.go
• s2/hashtable_pool.go