perf(encoding): compact mini-block chunk index for cached page state#7651
perf(encoding): compact mini-block chunk index for cached page state#7651Ali2Arslan wants to merge 9 commits into
Conversation
The mini-block scheduler cached each page as a `Vec<ChunkMeta>` plus a repetition index of per-chunk blocks (~48 bytes/chunk), most of which is derivable from a few cumulative quantities. Replace both with a compact `MiniBlockChunkIndex` that stores only the non-redundant data and derives the rest: prefix-sum byte offsets (narrowed to `u32` when the page fits) and a row mapping chosen per page shape -- `UniformFlat` (arithmetic, no allocation) for fixed-width/bitpacked pages, `Flat` for non-uniform flat pages (RLE/FSST), and `Nested` (row prefix sums + a trailer bitmap) when a repetition index is present. Row->chunk lookups stay O(1)/O(log n) and the scheduler's multi-range merge now runs in place. The `DeepSizeOf` impl accounts for every retained buffer so cache weighing reflects the smaller footprint. Adds unit tests for page-shape detection, byte-range/row/item axes, uniform-vs-flat scheduler parity, and per-variant size bounds. Co-authored-by: Cursor <cursoragent@cursor.com>
Trim the docs and comments added in the previous commit to the project convention: explain why not what, drop pure narration, and keep each to 2-3 lines. Comment/doc only; no behavior change. Co-authored-by: Cursor <cursoragent@cursor.com>
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
| impl DeepSizeOf for MiniBlockCacheableState { | ||
| fn deep_size_of_children(&self, context: &mut Context) -> usize { | ||
| self.rep_index.deep_size_of_children(context) | ||
| self.chunk_index.deep_size_of_children(context) |
There was a problem hiding this comment.
this also fixes the under-reporting (wasn't accounting for chunk_meta before)
| // are _adjacent_ (i.e. don't merge "take first row of chunk 0" and "take third row of chunk 0" into "take 2 | ||
| // rows of chunk 0 starting at 0") | ||
| if user_ranges.len() > 1 { | ||
| // TODO: Could probably optimize this allocation away |
There was a problem hiding this comment.
bit of a drive-by, but addressed this as well
westonpace
left a comment
westonpace
left a comment
There was a problem hiding this comment.
Thanks for looking at this. That schedule_instructions method still gives me nightmares. I'm glad we have LLMs now to deal with that kind of complexity.
The optimization is great, I think this is a good balance of performance and space. Do you have any benchmark results you can share?
Just a few nits around commenting. I'm mostly relying on the unit tests (the existing tests should be fairly complete) for correctness verification.
| //! Replaces two parallel per-chunk arrays (`Vec<ChunkMeta>` + a repetition index | ||
| //! of blocks, ~48 bytes/chunk) with the non-redundant cumulative quantities | ||
| //! alone, deriving the rest. The scheduler looks chunks up by index (byte range, | ||
| //! leaf value count) and by row (which chunk holds a row). |
There was a problem hiding this comment.
Minor nit: this comment will not age well as we have removed ChunkMeta so future readers won't know what this is talking about. Maybe:
The chunk index is stored on disk in an extremely compressed form that
requires a lot of CPU to work with. However, extracting it out to its
full width can be RAM-intensive. As a compromise we extract into a
prefix-sum array that we fit into u32 if possible and we avoid storing
per-block row counts when those are redundant.
| /// Index of the chunk whose half-open span `[get(i), get(i+1))` contains | ||
| /// `value`. On an exact hit against a chunk start, returns the *first* chunk | ||
| /// with that start (chunks can share a start row). | ||
| pub fn find(&self, value: u64) -> usize { | ||
| // Match the width once, then binary-search only the starts (not the | ||
| // trailing total). `partition_point` already yields the first of any | ||
| // duplicated starts; the `idx - 1` fallback is safe since `get(0) == 0`. | ||
| match self { | ||
| Self::U32(values) => { | ||
| let starts = &values[..values.len() - 1]; | ||
| let idx = starts.partition_point(|&start| (start as u64) < value); | ||
| if idx < starts.len() && starts[idx] as u64 == value { | ||
| idx | ||
| } else { | ||
| idx - 1 | ||
| } | ||
| } | ||
| Self::U64(values) => { | ||
| let starts = &values[..values.len() - 1]; | ||
| let idx = starts.partition_point(|&start| start < value); | ||
| if idx < starts.len() && starts[idx] == value { | ||
| idx | ||
| } else { | ||
| idx - 1 | ||
| } | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
Is there any reason we can't use rust's builtin binary search function (https://doc.rust-lang.org/std/vec/struct.Vec.html#method.binary_search)?
There was a problem hiding this comment.
seems like Vec's binary search impl can return an arbitrary match (among equal entries), and we rely on the first of the duplicates semantic. The old code did use binary search but then walked back to find the first duplicate, I can do that instead, if you would like?
| impl DeepSizeOf for PrefixSums { | ||
| fn deep_size_of_children(&self, context: &mut Context) -> usize { | ||
| match self { | ||
| Self::U32(values) => values.deep_size_of_children(context), | ||
| Self::U64(values) => values.deep_size_of_children(context), | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
Minor nit: you should be able to add the DeepSize derive and avoid this explicit impl
| /// `log2` of the value count of each chunk; the last chunk is handled via | ||
| /// `last_chunk_values`. |
There was a problem hiding this comment.
Maybe mention in the comment that we use log2 here for efficiency as this must be cached in RAM?
| } | ||
| } | ||
|
|
||
| impl DeepSizeOf for ItemCounts { |
There was a problem hiding this comment.
Again, prefer derived DeepSize if possible
| /// Compact per-page chunk index that replaces the previous `Vec<ChunkMeta>` | ||
| /// plus repetition-index representation. See the module docs for the layout. |
There was a problem hiding this comment.
| /// Compact per-page chunk index that replaces the previous `Vec<ChunkMeta>` | |
| /// plus repetition-index representation. See the module docs for the layout. | |
| /// Compact per-page chunk index that avoids fully materializing the repetition | |
| /// index into u64's to save RAM. See the module docs for the layout. |
Thanks for taking a look Weston! I'll fix the comments! |
Reword the module and struct docs to describe the CPU/RAM tradeoff instead of the now-removed `ChunkMeta`, note that `ItemCounts::PerChunkLog` stores log2 per chunk to stay compact in RAM, and replace the hand-written `DeepSizeOf` impls for `PrefixSums` and `ItemCounts` with the derive. No behavior change. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: QUIET Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughMiniBlock structural decoding now builds compact ChangesMiniBlock chunk indexing
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested labels: Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant MiniBlockScheduler
participant MiniBlockChunkIndex
participant ChunkInstructions
participant DataBuffer
MiniBlockScheduler->>MiniBlockChunkIndex: lookup chunk byte ranges and item counts
ChunkInstructions->>MiniBlockChunkIndex: locate chunks and row boundaries
ChunkInstructions->>MiniBlockChunkIndex: read preamble and trailer flags
ChunkInstructions->>DataBuffer: schedule merged chunk reads
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/lance-encoding/src/encodings/logical/primitive.rs`:
- Around line 1803-1806: Make build_chunk_index propagate Result and replace the
repetition-index length assert with Error::corrupt_file including the buffer
length and stride. In rust/lance-encoding/src/encodings/logical/primitive.rs
lines 1803-1806, update the rep_index_bytes path to return the checked error. In
rust/lance-encoding/src/encodings/logical/primitive/chunk_index.rs lines
393-431, make parse_nested_rep fallible, replace slice unwraps with error
handling, and validate cumulative subtraction so malformed bytes return
Error::corrupt_file instead of panicking.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: c749cf27-70d9-4a63-9ecf-9a3672f24740
📒 Files selected for processing (2)
rust/lance-encoding/src/encodings/logical/primitive.rsrust/lance-encoding/src/encodings/logical/primitive/chunk_index.rs
| let rows = if let Some(rep_index_data) = rep_index_bytes { | ||
| assert!(rep_index_data.len() % 8 == 0); | ||
| let stride = repetition_index_depth as usize + 1; | ||
| let (row_starts, has_trailer) = parse_nested_rep(rep_index_data, stride); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm build_chunk_index is only reached from a Result-returning context and locate the Error constructors available.
rg -nP '\bbuild_chunk_index\s*\(' rust/lance-encoding/src
rg -nP 'fn corrupt_file' rust/lance-core/src
rg -nP '\bparse_nested_rep\s*\(' rust/lance-encoding/srcRepository: lance-format/lance
Length of output: 2080
Return Error::corrupt_file for malformed repetition-index bytes
rust/lance-encoding/src/encodings/logical/primitive.rs#L1803-L1806: replaceassert!(rep_index_data.len() % 8 == 0)with a checked error path and threadResultthroughbuild_chunk_indexso the failure can include the buffer length and stride.rust/lance-encoding/src/encodings/logical/primitive/chunk_index.rs#L393-L431: makeparse_nested_repfallible too; the slice.unwrap()and cumulative subtraction can panic on malformed bytes.
📍 Affects 2 files
rust/lance-encoding/src/encodings/logical/primitive.rs#L1803-L1806(this comment)rust/lance-encoding/src/encodings/logical/primitive/chunk_index.rs#L393-L431
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance-encoding/src/encodings/logical/primitive.rs` around lines 1803 -
1806, Make build_chunk_index propagate Result and replace the repetition-index
length assert with Error::corrupt_file including the buffer length and stride.
In rust/lance-encoding/src/encodings/logical/primitive.rs lines 1803-1806,
update the rep_index_bytes path to return the checked error. In
rust/lance-encoding/src/encodings/logical/primitive/chunk_index.rs lines
393-431, make parse_nested_rep fallible, replace slice unwraps with error
handling, and validate cumulative subtraction so malformed bytes return
Error::corrupt_file instead of panicking.
Source: Learnings
There was a problem hiding this comment.
I can do this, if we want @westonpace , though it's not technically related to this PR (it was like this before), but happy to fix it (either here - or as a separate follow-up)
There was a problem hiding this comment.
Seems like the humans are having a chat. I'll hop back into my burrow for now. If you need me again, just tag @coderabbitai in a new comment, and I'll come hopping out!
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rust/lance-encoding/src/encodings/logical/primitive.rs (1)
1733-1774: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftValidate decoded chunk counts before subtracting.
items_in_page - countedis unchecked, whilecountedis derived from on-disk metadata. Corrupt metadata can therefore panic in debug builds or wrap in release builds, producing an invalidMiniBlockChunkIndexand potentially incorrect or out-of-bounds reads. The surroundingdebug_assert!checks do not protect release builds.Make the count-analysis helpers fallible, use checked arithmetic, return
Error::corrupt_filewith the chunk index and count values, propagate the error throughbuild_chunk_index, and add malformed-metadata tests.As per coding guidelines, validate inputs at API boundaries, use checked arithmetic, include full error context, and add tests for bug fixes.
Also applies to: 1884-1892
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance-encoding/src/encodings/logical/primitive.rs` around lines 1733 - 1774, Make analyze_value_counts and flat_value_counts_iter fallible, validating each decoded chunk count before subtraction and using checked arithmetic for shifts, accumulation, and remaining-count calculations. Return Error::corrupt_file with the chunk index and relevant count values on invalid metadata, then propagate the failure through build_chunk_index instead of constructing an invalid MiniBlockChunkIndex. Add malformed-metadata tests covering underflow and invalid decoded counts.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@rust/lance-encoding/src/encodings/logical/primitive.rs`:
- Around line 1733-1774: Make analyze_value_counts and flat_value_counts_iter
fallible, validating each decoded chunk count before subtraction and using
checked arithmetic for shifts, accumulation, and remaining-count calculations.
Return Error::corrupt_file with the chunk index and relevant count values on
invalid metadata, then propagate the failure through build_chunk_index instead
of constructing an invalid MiniBlockChunkIndex. Add malformed-metadata tests
covering underflow and invalid decoded counts.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 714dcdbc-3ca5-408e-8d65-c8365a0eee7f
📒 Files selected for processing (1)
rust/lance-encoding/src/encodings/logical/primitive.rs
westonpace
left a comment
westonpace
left a comment
There was a problem hiding this comment.
Looks like you addressed the review comments. Is this ready to merge?
Btw, in future PRs, if you hit re-request review then it will go back into my queue which makes it easier to keep track:
Yup ready to merge, and I see, will hit re-review in the future, thank you!!!! |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rust/lance-encoding/src/encodings/logical/primitive.rs (1)
3327-3336: 🎯 Functional Correctness | 🔴 Critical | ⚡ Quick winPass an owned
DataTypehereSparseStructuralScheduler::try_newtakesDataType, buttarget_field.data_type()yields&DataType; clone it at this call site.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance-encoding/src/encodings/logical/primitive.rs` around lines 3327 - 3336, Update the SparseLayout branch in the scheduler construction to pass an owned DataType to SparseStructuralScheduler::try_new by cloning the value returned from target_field.data_type(). Leave the surrounding arguments and error propagation unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@rust/lance-encoding/src/encodings/logical/primitive.rs`:
- Around line 3327-3336: Update the SparseLayout branch in the scheduler
construction to pass an owned DataType to SparseStructuralScheduler::try_new by
cloning the value returned from target_field.data_type(). Leave the surrounding
arguments and error propagation unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 25a9ea19-e851-4e86-9cda-0ae3ecea8302
📒 Files selected for processing (1)
rust/lance-encoding/src/encodings/logical/primitive.rs
The main merge pulled in the sparse structural encoding module (primitive/sparse.rs), which still relies on the `ChunkMeta` struct this branch had removed from the mini-block path. Re-add the struct so the new consumer compiles; the mini-block scheduler continues to use the compact `MiniBlockChunkIndex`. Co-authored-by: Cursor <cursoragent@cursor.com>
2 participants
Summary
The mini-block structural scheduler caches per-page state that is consulted on
every read. Today that state is a
Vec<ChunkMeta>(num_values, size, offset perchunk) plus a repetition index of per-chunk blocks (first_row,
starts_including_trailer, has_preamble, has_trailer) — ~48 bytes/chunk, almost
all of it derivable from a handful of cumulative quantities.
This PR replaces both with a compact
MiniBlockChunkIndex(newprimitive/chunk_index.rssubmodule) that stores only the non-redundant dataand derives the rest:
byte_starts: PrefixSums— cumulative chunk byte sizes, stored asu32when the page's data buffer fits in 4 GiB and
u64otherwise.rows: RowMappingselected by page shape:UniformFlat— fixed-width / bitpacked pages where every non-last chunkholds the same number of values; row↔chunk is pure arithmetic with no
heap allocation.
Flat— non-uniform flat pages (RLE / FSST); a single cumulative valuearray.
Nested— repetition present; cumulative row starts + ahas_trailerbitmap (preamble derived from the previous chunk's trailer) + leaf item
counts.
byte_range/items_in_chunk/find_chunk/first_row/rows_in_chunk/has_preamble/has_trailerpreserve the exact semantics the scheduler reliedon (
find_chunkstill returns the first of duplicated start rows). Row→chunklookups remain O(1) (uniform flat) or O(log n) (binary search), and the
multi-range instruction merge in
schedule_instructionsnow runs in placeinstead of allocating a second
Vec.MiniBlockChunkIndeximplementsDeepSizeOf(composing the innerPrefixSums/ItemCounts/RowMapping), so cache weighing reflects the smallerfootprint.
This is an idiomatic port of an internal change to current upstream: it uses
lance's
DeepSizeOftrait (rather than a bespoke size method), threads theexisting
initializeIO structure, and drops fork-only test scaffolding thathas no counterpart here.
Test plan
cargo test -p lance-encoding --lib— 422 passed, 0 failed, 5 ignored(covers flat, non-uniform-flat, and nested/list miniblock round trips,
e.g.
test_sparse_large_string_list::...MINIBLOCK,test_sparse_boolean_list_uses_miniblock,test_nested_strings).test_flat_detectioncases),nested row/item axes (
test_nested_detection_and_axes), uniform-vs-Flatscheduler/lookup parity (
test_uniform_flat_matches_prefix_sum_flat),per-variant heap-size bounds vs the legacy layout
(
test_deep_size_per_variant_below_legacy), and aparse_nested_reporacle test against the previous per-block decode.
cargo fmt -p lance-encodingcargo clippy -p lance-encoding --tests --benches -- -D warningsMade with Cursor