-
Notifications
You must be signed in to change notification settings - Fork 417
Add space-efficient arena-based maps #9776
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add space-efficient arena-based maps #9776
Conversation
Current kv.mem iteration sorts a slice at every call to Next(). Instead use a sorted map. Will use this in `lakefs gc simulate`.
Most of the issues flagged by copilot are irrelevant. In the case of Set operations it might make some sense to copy slices - they live longer, and the caller might use them somewhere (it doesn't happen in practice). Other kv operations are Gets, and that just doesn't happen.
- Decode JSON objects
It read JSON inputs (`kv scan` output and retention rules JSON), loads them into an in-memory KV (for _speed_ 🐎) and then runs the actual GC code on it.
Controlled by flags.
Code for testing arena.Map was 🎉 > 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
Change Put method signature from `Put(k K, v V)` to `Put(k K, v V) *V` to return a pointer to the value stored in the arena. Update test to verify the returned pointer is valid and points to the correct value. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
The default growth is 2x. Manually growing size 1.15x each time _increases_ the number of copied bytes 3.5x but _decreases_ worst-case memory overuse from 2x to 1.15x. This is worth doing as GC memory pressure is currently worse than CPU time,
Extract testArenaMap helper function that tests Map interface behavior. This function is now called by both TestArenaMap (using NewMap) and TestBoundedArenaMap (using NewBoundedKeyMap). Both implementations pass the same test suite, verifying that boundedArenaMap correctly implements the Map interface. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
Fix bugs in boundedArenaMap: - Fix compareKey to properly compare zero-padded keys by using bytes.Compare instead of converting to string with trailing nulls - Remove unused cmp import - Fix Optimize to handle duplicate keys by using SortStableFunc and deduplicating entries, keeping the last occurrence Add comprehensive TestBoundedArenaMapOptimize that verifies: - Put 100 elements in shuffled order and retrieve them - Optimize and verify elements still accessible - Put 100 more elements and overwrite 50 old elements in shuffled order - Verify all 200 elements work correctly - Optimize again and verify all 200 elements still accessible Test uses shuffled insertion order to verify sorting and binary search work correctly regardless of insertion order. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
Add TestMapSizes that compares actual memory usage between: - Regular Go map[string]string - arena.Map (arenaMap) - arena.BoundedKeyMap (boundedArenaMap) before and after Optimize Use github.com/DmitriyVTitov/size to measure runtime memory consumption. Results with 1000 entries show: - Regular map: 56,578 bytes - arena.Map: 65,002 bytes (+15%) - BoundedKeyMap: 65,034 bytes (+15%) - BoundedKeyMap (optimized): 40,722 bytes (-28% vs regular map) Demonstrate that boundedArenaMap.Optimize() provides significant memory savings by packing keys into fixed-size arrays and values into a sorted slice. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
nopcoder
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
As a general use of a container I think that a slice backed by a map for fast lookup or set of function (like Go's heap) that will produce a packed data structure will do a good job.
I think that IndexedList https://share.google/aimode/SvvBgwdSKbgMJI86J or similar construct can help building the final structure we like to work with. In case we don't want to have the backed map - we can just use slice and Optimize function.
pkg/arena/arena.go
Outdated
| m.bigMap[writeIdx] = m.bigMap[readIdx] | ||
| } else { | ||
| // Same key - keep the later one (at readIdx) | ||
| m.bigMap[writeIdx] = m.bigMap[readIdx] |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
do we need the else block? or we like to skip the duplicates
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We actually need this!
Your way of doing it also works, and has less indentation 👉🏽 better.
| for readIdx := 1; readIdx < len(m.bigMap); readIdx++ { | ||
| if m.compareEntries(m.bigMap[writeIdx], m.bigMap[readIdx]) != 0 { | ||
| writeIdx++ | ||
| m.bigMap[writeIdx] = m.bigMap[readIdx] | ||
| } else { | ||
| // Same key - keep the later one (at readIdx) | ||
| m.bigMap[writeIdx] = m.bigMap[readIdx] | ||
| } | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
we like to skip duplicates - so we like to skip the assignment or we like to move the assignment outside the if blocks.
did you mean:
if m.compareEntries(m.bigMap[writeIdx], m.bigMap[readIdx]) != 0 {
writeIdx++
}
m.bigMap[writeIdx] = m.bigMap[readIdx]There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That doesn't work, this loop needs the invariant writeIdx < readIdx.
Actually I had another look, and this is almost just a poor implementation of slices.CompactFunc -- except we want to keep the last equivalent element where that one keeps the first. We can't use it, but reading its code reminds me to clear the last elements at the end - which is good for GC if V holds any references.
Anyway, simplified the implementation slightly, but added a call to clear; hope it's good now.
pkg/arena/arena.go
Outdated
| } | ||
| } | ||
|
|
||
| const KEY_SIZE_BOUND = 16 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
| const KEY_SIZE_BOUND = 16 | |
| const KeySizeBound = 16 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
True, sorry: writing contiguous-memory data structures put me into C++/Rust mode.
pkg/arena/arena.go
Outdated
| // NewBoundedKeyMap returns a Map that uses string-like keys of bounded length. Keys are | ||
| // zero-padded, so must not end in zero bytes. This Map is not thread-safe. | ||
| // | ||
| // It to keep keys in an Arena. The map *panics* if it encounters a longer key. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
| // It to keep keys in an Arena. The map *panics* if it encounters a longer key. | |
| // It allows keeping keys in an Arena. The map *panics* if it encounters a longer key. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks!
pkg/arena/arena.go
Outdated
|
|
||
| func newArenaMap[K comparable, V any]() *arenaMap[K, V] { | ||
| return &arenaMap[K, V]{ | ||
| indices: make(map[K]Index, 0), |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit; 0 is the default.
| indices: make(map[K]Index, 0), | |
| indices: make(map[K]Index), |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done, thanks!
| if index, ok := m.indices[k]; ok { | ||
| ptr := m.arena.Get(index) | ||
| *ptr = v | ||
| return ptr | ||
| } else { | ||
| index = m.arena.Add(v) | ||
| m.indices[k] = index | ||
| return m.arena.Get(index) | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
if we use return, can we avoid the else. https://go.dev/doc/effective_go#if
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done.
| if index, ok := m.indices[k]; ok { | ||
| return m.arena.Get(index) | ||
| } else { | ||
| return nil | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
| if index, ok := m.indices[k]; ok { | |
| return m.arena.Get(index) | |
| } else { | |
| return nil | |
| } | |
| if index, ok := m.indices[k]; ok { | |
| return m.arena.Get(index) | |
| } | |
| return nil |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Early return adopted throughout, thanks!
| type Index struct { | ||
| o int | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
do we need a struct? can we use type Index int
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
But then Index is an int. And then callers can write things like
var i Index = myArena.Add(..)
i++
myArena.Get(i) // blows upor even just
i := Index(17)
myArena.Get(i)I miss Rust and C++ 😭 .
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't see an issue - it is an integer.
If you call Get with invalid index you get nil by design.
You can always use Index as pointer to the data which the user can't increment.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I do not see the harm in type safety here.
- Incrementing an index does not yield a meaningful new index. I would like indexes to be safe.
- Using a pointer, however, would be unsafe! As soon you add elements and the arena resizes, the pointer goes bad. (This is exactly what happens in the current use case - we keep an Index in the map and then keep on growing the arena.)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
not blocking.
about the pointer - it will keep a reference to the data, but it is not what you want.
| // Arena holds Ts packed together for fast random access. It is an efficient way to store many | ||
| // Ts. If T is a pointer (or smaller), this does not save anything. | ||
| type Arena[T any] interface { | ||
| // Add copies t into the arena and returns its index. It invalidates results from all | ||
| // previous Gets. The returned Index is valid for the lifetime of the arena. | ||
| Add(t T) Index | ||
| // Get returns the T at index or nil. The returned pointer is valid until the next time | ||
| // the Arena is mutated (New or Add). | ||
| Get(index Index) *T | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
- Can we have the container itself public, as we like to use a specific type when possible without going through interface. I don't mind having an interface too.
- The Arena implementation (not just the selected line) is a slice with alternative way to control the growth? wanted to understand better why we need a type vs use slice.
- Consider: Add => Append, adding Len method
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
- Done: renamed
arena→SliceArena. FTR, that will give a zero growth factor, which is usually bad. - Why Arena and not just a slice? I think the abstraction is worthwhile!
- It exactly controls growth. Really important, otherwise we sometimes take 2x as much space as needed this is a bad memory/time tradeoff for GC which takes up the majority of process memory.
- It is not a slice, so you cannot e.g.
slices.DeleteFunc(sliceNotArena, ...)and clobber all your indexes. - It gives somewhere to document pointer validity returned from Get.
- Done (Add -> Append, Len)
| type arenaMap[K comparable, V any] struct { | ||
| indices map[K]Index | ||
| arena Arena[V] | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
same as with the area - I think this type should be public as must of the time I'll like to create the type and not use the interface
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not done. All fields of this type must be initialized.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We can return the struct as it matches the interface.
It is fine to return a public type through NewXXX and not expect the user to instantiate one on their own - like os.File for example.
It will enable users to use the concrete type, which will perform better than going through an interface.
arielshaqed
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think that IndexedList https://share.google/aimode/SvvBgwdSKbgMJI86J or similar construct can help building the final structure we like to work with. In case we don't want to have the backed map - we can just use slice and Optimize function.
Good point. But if we did that we would need to call Optimize really often - or memory would blow up even faster. And Using small-map / big-map lets us control memory usage nicely. In practice the copy at the top of Optimize is nearly contiguous ordered blocks of memory, so it's even fast.
Thanks, PTAL!
| type Index struct { | ||
| o int | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
But then Index is an int. And then callers can write things like
var i Index = myArena.Add(..)
i++
myArena.Get(i) // blows upor even just
i := Index(17)
myArena.Get(i)I miss Rust and C++ 😭 .
| // Arena holds Ts packed together for fast random access. It is an efficient way to store many | ||
| // Ts. If T is a pointer (or smaller), this does not save anything. | ||
| type Arena[T any] interface { | ||
| // Add copies t into the arena and returns its index. It invalidates results from all | ||
| // previous Gets. The returned Index is valid for the lifetime of the arena. | ||
| Add(t T) Index | ||
| // Get returns the T at index or nil. The returned pointer is valid until the next time | ||
| // the Arena is mutated (New or Add). | ||
| Get(index Index) *T | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
- Done: renamed
arena→SliceArena. FTR, that will give a zero growth factor, which is usually bad. - Why Arena and not just a slice? I think the abstraction is worthwhile!
- It exactly controls growth. Really important, otherwise we sometimes take 2x as much space as needed this is a bad memory/time tradeoff for GC which takes up the majority of process memory.
- It is not a slice, so you cannot e.g.
slices.DeleteFunc(sliceNotArena, ...)and clobber all your indexes. - It gives somewhere to document pointer validity returned from Get.
- Done (Add -> Append, Len)
pkg/arena/arena.go
Outdated
|
|
||
| func newArenaMap[K comparable, V any]() *arenaMap[K, V] { | ||
| return &arenaMap[K, V]{ | ||
| indices: make(map[K]Index, 0), |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done, thanks!
| type arenaMap[K comparable, V any] struct { | ||
| indices map[K]Index | ||
| arena Arena[V] | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not done. All fields of this type must be initialized.
| if index, ok := m.indices[k]; ok { | ||
| ptr := m.arena.Get(index) | ||
| *ptr = v | ||
| return ptr | ||
| } else { | ||
| index = m.arena.Add(v) | ||
| m.indices[k] = index | ||
| return m.arena.Get(index) | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done.
| if index, ok := m.indices[k]; ok { | ||
| return m.arena.Get(index) | ||
| } else { | ||
| return nil | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Early return adopted throughout, thanks!
pkg/arena/arena.go
Outdated
| } | ||
| } | ||
|
|
||
| const KEY_SIZE_BOUND = 16 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
True, sorry: writing contiguous-memory data structures put me into C++/Rust mode.
pkg/arena/arena.go
Outdated
| // NewBoundedKeyMap returns a Map that uses string-like keys of bounded length. Keys are | ||
| // zero-padded, so must not end in zero bytes. This Map is not thread-safe. | ||
| // | ||
| // It to keep keys in an Arena. The map *panics* if it encounters a longer key. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks!
pkg/arena/arena.go
Outdated
| m.bigMap[writeIdx] = m.bigMap[readIdx] | ||
| } else { | ||
| // Same key - keep the later one (at readIdx) | ||
| m.bigMap[writeIdx] = m.bigMap[readIdx] |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We actually need this!
Your way of doing it also works, and has less indentation 👉🏽 better.
| for readIdx := 1; readIdx < len(m.bigMap); readIdx++ { | ||
| if m.compareEntries(m.bigMap[writeIdx], m.bigMap[readIdx]) != 0 { | ||
| writeIdx++ | ||
| m.bigMap[writeIdx] = m.bigMap[readIdx] | ||
| } else { | ||
| // Same key - keep the later one (at readIdx) | ||
| m.bigMap[writeIdx] = m.bigMap[readIdx] | ||
| } | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That doesn't work, this loop needs the invariant writeIdx < readIdx.
Actually I had another look, and this is almost just a poor implementation of slices.CompactFunc -- except we want to keep the last equivalent element where that one keeps the first. We can't use it, but reading its code reminds me to clear the last elements at the end - which is good for GC if V holds any references.
Anyway, simplified the implementation slightly, but added a call to clear; hope it's good now.
pkg/arena/arena.go
Outdated
| } | ||
|
|
||
| type boundedArenaMap[K ~string, V any] struct { | ||
| // bigMap is sorted slice of pairs. Apart from calls to Optimize it is immutable, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
| // bigMap is sorted slice of pairs. Apart from calls to Optimize it is immutable, | |
| // bigMap is sorted slice of pairs. Apart from calls to Optimize it is immutable. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done :-)
| type Index struct { | ||
| o int | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't see an issue - it is an integer.
If you call Get with invalid index you get nil by design.
You can always use Index as pointer to the data which the user can't increment.
arielshaqed
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks! I do want to protect Index, at least a bit, without making it take more space or time. So the alternative of returning an interface type Index which wraps a type index int isn't possible, either. Let's talk F2F.
Relevant references:
- SO "What is the purpose of creating a struct with only one field" - for C but the same point applies even more to Go.
- type CID in github.com/ipfs/go-cid holds a single unexported string.
- Here's the same idea in filecoin
| type Index struct { | ||
| o int | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I do not see the harm in type safety here.
- Incrementing an index does not yield a meaningful new index. I would like indexes to be safe.
- Using a pointer, however, would be unsafe! As soon you add elements and the arena resizes, the pointer goes bad. (This is exactly what happens in the current use case - we keep an Index in the map and then keep on growing the arena.)
pkg/arena/arena.go
Outdated
| } | ||
|
|
||
| type boundedArenaMap[K ~string, V any] struct { | ||
| // bigMap is sorted slice of pairs. Apart from calls to Optimize it is immutable, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done :-)
Summary of f2f discussions with @nopcoder: 1. smallMap is already a map, which is a reference. No need to keep a pointer to it. 1. Deduplication in Optimize _is_ necessary. Sample scenario: ```go m.Put("key", 1) // bigMap: {}, smallMap: {"key": 1} m.Optimize() // bigMap: {"key": 1}, smallMap: {} m.Put("key", 2) // bigMap: {"key": 1}, smallMap: {"key", 2} m.Optimize() // must give bigmap: {"key": 2}, smallMap: {} ``` 1. Deduplication _works_ thanks to stable sort. In the above sample, smallMap is first appended to bigMap giving [{"key": 1}, {"key": 2}] _in that order_. Deduplication takes the last occurrence of "key" and gives the correct result.
arielshaqed
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks! PTAL...
But of course, also verify that each concrete type satisfies its interface type.
nopcoder
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
lgtm.
| type Index struct { | ||
| o int | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
not blocking.
about the pointer - it will keep a reference to the data, but it is not what you want.
| type arenaMap[K comparable, V any] struct { | ||
| indices map[K]Index | ||
| arena Arena[V] | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We can return the struct as it matches the interface.
It is fine to return a public type through NewXXX and not expect the user to instantiate one on their own - like os.File for example.
It will enable users to use the concrete type, which will perform better than going through an interface.
31154db to
581e1c3
Compare
…icient-arena-based-maps
|
Thanks! I merged master into this branch (rebasing was too painful given the number of changes on trunk) and will pull once tests pass. |
Introduce memory-efficient map implementations using arena allocation,
designed to reduce memory overhead for storing large numbers of small
objects.
Components
arena.Arena
Indexreferences (8 bytes) instead of pointersGet(index Index) *Tarena.Map
Put(k K, v V) *Vreturns pointer to stored valueBoundedKeyMap
[16]bytearraysOptimize(): UsessmallMap(regular arena.Map)Optimize(): Moves to sortedbigMapslice with binary searchMemory Savings
Measured with 1000 entries using
github.com/DmitriyVTitov/size:map[string]stringarena.MapBoundedKeyMap(before opt)BoundedKeyMap(after opt)Implementation Details
BoundedKeyMap.Optimize()deduplicates entries and uses stable sort to preserve insertion order for duplicate keyscompareKeyuses byte comparison of zero-padded keys for binary searchWhodunit
it and correct the rest.
Favourite Claude quote from the description (emphasis mine):
Monty Python should sue Anthropic.