feat: make the index cache size (in bytes) available

[westonpace]

This PR makes extensive use of the deepsize crate for tracking the deep size of standard structures (Vec, BTreeMap, etc.). This crate does appear to be abandoned (no updates in 3 years). However, I'm unable to find anything similar, the code is fairly small should we need to vendor it in the future, and it is unlikely to be doing anything dangerous.

It turns out we store many many things in the cache (transitively) though most of them are probably small.

[westonpace]

I'm leaving this in draft for now because I don't quite think we are doing the right thing for arrow array memory sizes. Some of the index structures may share the same arrow array and I believe this implementation (using get_array_memory_size) will double-count. We will need to create our own version of get_array_memory_size that somehow keeps track of whether or not we've seen an array before. I'd like to at least spend a bit of time investigating if that is possible.

On the other hand, overcounting might be better than having 0 stats available, so I'm not opposed to merging this in as it is if we want the capability sooner.

[westonpace]

@BubbleCal do you think visited_generator_queue will occupy a lot of space in memory? I wasn't sure if it was safe to skip that here.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 5.17928% with 238 lines in your changes are missing coverage. Please review.

Project coverage is 79.68%. Comparing base (83ecc01) to head (8731ae1).

Files	Patch %	Lines
rust/lance-index/src/scalar/btree.rs	0.00%	49 Missing ⚠️
rust/lance-index/src/scalar/flat.rs	0.00%	18 Missing ⚠️
rust/lance/src/index/cache.rs	0.00%	18 Missing ⚠️
rust/lance/src/index/vector/pq.rs	0.00%	13 Missing ⚠️
rust/lance-core/src/cache.rs	50.00%	12 Missing ⚠️
rust/lance-index/src/vector/pq/storage.rs	0.00%	12 Missing ⚠️
rust/lance-table/src/format/index.rs	0.00%	11 Missing ⚠️
rust/lance-index/src/vector/hnsw/builder.rs	0.00%	8 Missing and 1 partial ⚠️
rust/lance-io/src/object_store.rs	0.00%	9 Missing ⚠️
rust/lance-index/src/vector/ivf/storage.rs	0.00%	8 Missing ⚠️
... and 25 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2381      +/-   ##
==========================================
- Coverage   80.04%   79.68%   -0.36%     
==========================================
  Files         199      199              
  Lines       54469    54713     +244     
  Branches    54469    54713     +244     
==========================================
+ Hits        43598    43600       +2     
- Misses       8349     8586     +237     
- Partials     2522     2527       +5     

Flag	Coverage Δ
unittests	79.68% <5.17%> (-0.36%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[westonpace]

I can't figure out a great way to get the Arrow memory size more reliably. The way deepsize avoids double counting Arc is by storing a set of Arc pointers that it has already visited. Arrow's Buffer does not expose its inner Arc<Bytes> and so we can't tie into deepsize's Context.

There is as_ptr but deepsize doesn't have a way to manually track pointers. We could create our own Context but this would mean we have to replace deepsize with our own thing which is a bit more involved than I want to do without knowing there is a problem.

So if two structures have a cloned ArrayRef we are fine. If two structures have a cloned Array we may overcount due to double counting.

[wjones127]

So if two structures have a cloned ArrayRef we are fine. If two structures have a cloned Array we may overcount due to double counting.

Honestly, these seems acceptable to me for now.

[wjones127]

I wonder if we should have some sort of session handle instead. It might be too early, but would be nice to have something users could pass around and re-use across datasets.

[westonpace]

I've added a new session object. It can only report the size right now, so you can't use it to pass between datasets, but it will help avoid breaking changes in the future when we do add that capability.

[wjones127]

Why make this zero? Also, could we use ScalarValue::size()? https://docs.rs/datafusion/latest/datafusion/common/enum.ScalarValue.html#method.size

[westonpace]

Can't believe I missed that 🤦 . Yes, that will work. It's 0 because deepsize does the std::mem::size_of already so this is only supposed to account for anything above and beyond that.

[wjones127]

For clarity, should it be cache_size_bytes()?

[westonpace]

I think, if we have any other stateful objects in the session, we would want to capture that here as well. In other words, it would make more sense to rename Session to DatasetCache than it would to rename size_bytes to cache_size_bytes.

[wjones127]

I see. For context, I'm thinking other methods we might set on the context will be execution args, like:

class Session:
    def set_cpu_pool_size(num_cpus: int): ...
    def set_io_pool_size(num_cpus: int): ...
    def set_memory_limit(memory_limit: int): ...
    def set_spill_location(path: Path): ...

So I'm thinking of the Session object as containing more state than the cache.

[westonpace]

So I'm thinking of the Session object as containing more state than the cache.

Sorry, my comment was probably confusing. I think this is good, my point should have been "it doesn't make sense to rename Session to DatasetCache and therefore doesn't make sense to rename size_bytes to cache_size_bytes".

For example, maybe DF starts to need some "thread local temp space" or "slab allocator working area" or whatever (I'm probably just making terms up at this point). This would go in the session, would not be trivial in size, and would not be considered a "cache", but we would still want to make sure that size_bytes includes it.