perf(encoding): cache complex-all-null levels in RLE run form and drain run-by-run#7564
Conversation
…in run-by-run Complex-all-null pages (nested/list/struct all-null columns) cached their rep/def levels by eagerly expanding RLE to `num_values` u16s. A production heap profile attributed ~79% of in-use bytes to `ComplexAllNullScheduler` holding these expanded levels. Cache the levels lazily in coalesced RLE run form instead: - rle.rs: extract `parse_rle_block_frame` from `BlockDecompressor::decompress` (pure refactor) and add `decode_rle_u16_runs`, which parses an RLE frame into per-run `(values, offsets)` and coalesces adjacent equal-valued runs. The block encoder caps runs at 255, so a logically constant page arrives as `ceil(num_values / 255)` runs; coalescing collapses them to O(distinct levels), making the cached form O(runs) instead of O(num_values). - primitive.rs: introduce `LevelCodec` (Uncompressed | Rle | Block) so only RLE bypasses the eager decompressor, and `LazyLevels` (Dense | Rle) to hold the run form. `initialize` builds `LazyLevels` per level buffer; non-RLE paths still expand to `Dense` exactly as before. - Drain run-by-run with a monotonic `LevelCursor`: `seek_row_start` / `count_le_cursor` advance forward through runs so a full page decode is O(runs), not O(rows), which keeps the run-form representation from making per-row seeks costly. Emitted level slices, visible counts, the `RepDefUnraveler`, and the `AllNull` output are unchanged. Tests: Dense-vs-Rle parity, a 512-case proptest against a brute-force drain reference, run-form compactness, RLE coalescing, and an end-to-end all-null nested round-trip whose levels RLE-compress to a single run. Co-authored-by: Cursor <cursoragent@cursor.com>
Drop comments that restated the code (the extend-in-place branch, the "eager expansion" test label, the extend_into summary), remove the coalescing rationale duplicated between decode_rle_u16_runs' doc and its body, and drop the sentence pinning decode_rle_u16_runs to the complex-all-null caller. No behavior change. Co-authored-by: Cursor <cursoragent@cursor.com>
| Self::Rle { values, offsets } => { | ||
| let mut level = cursor.level; | ||
| let mut run = cursor.run; | ||
| while need > 0 { |
There was a problem hiding this comment.
we're still walking each run here, for a far row on a high-run page we should binary search instead if we cached the per-run cumulative row counts, but not sure if it's worth it given that RLE is only chosen when runs is a lot less than values
There was a problem hiding this comment.
Yeah, I wouldn't be too worried. It seems like something we can optimize later if we realize it is becoming a problem. I don't think anything in here prevents us from taking that path.
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
Pull request overview
This PR reduces memory overhead in the complex-all-null decoding path by caching repetition/definition levels in a compact RLE run form (values + cumulative offsets) instead of eagerly expanding to num_values dense u16 buffers, and then draining using monotonic cursors so decoding work is proportional to the number of runs.
Changes:
- Refactors RLE block-frame parsing and adds
decode_rle_u16_runsto decode/coalesceu16RLE runs without full expansion. - Introduces
LevelCodec,LazyLevels(Dense vs Rle), and monotonic cursor logic inComplexAllNullPageDecoderto drain run-by-run. - Adds unit tests + proptest parity tests + an end-to-end round trip covering constant-def RLE.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| rust/lance-encoding/src/encodings/physical/rle.rs | Adds reusable RLE frame parsing and run-form decoding for u16 levels (with coalescing) plus tests. |
| rust/lance-encoding/src/encodings/logical/primitive.rs | Switches complex-all-null rep/def caching to LazyLevels with LevelCursor-based monotonic draining and adds extensive parity/footprint tests. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| debug_assert!(target_row >= cursor.row); | ||
| let mut need = target_row - cursor.row; |
There was a problem hiding this comment.
We can make it a real error, but this invariant is maintained internally, so the error is not really actionable
| .map(Arc::from); | ||
| // RLE levels are kept lazily in run form; other block | ||
| // compressions keep flowing through the eager decompressor. | ||
| let level_codec = |
There was a problem hiding this comment.
Compression::Rle(_) is mapped directly to LevelCodec::Rle, so malformed metadata that the old decompressor rejected can be accepted and interpreted from payload bytes.
There was a problem hiding this comment.
ah good catch! thanks for the review, will fix!
There was a problem hiding this comment.
seems like RLE block lens can now be u16/u32 as well, I need to catch up on the latest changes
| } | ||
| } | ||
|
|
||
| fn deep_size(&self) -> usize { |
There was a problem hiding this comment.
For RLE blocks with run counts between roughly 20% and 50% of row count, the new run-form cache can exceed dense cache size and conflict with the PR memory reduction goal.
There was a problem hiding this comment.
hmm good point, I guess we can't naively map all RLE blocks to this repr...
Co-authored-by: Cursor <cursoragent@cursor.com>
Select per-page between physical RLE, coalesced runs, and dense levels to minimize retained memory while preserving efficient traversal. Validate compressed metadata and level buffers before caching. Co-authored-by: Cursor <cursoragent@cursor.com>
📝 WalkthroughWalkthroughThe change adds validated physical RLE run decoding, introduces lazy repetition/definition level representations, and rewrites complex-all-null scheduling and decoding to process level runs without eagerly materializing dense buffers. ChangesComplex all-null level decoding
Estimated code review effort: 5 (Critical) | ~100 minutes Sequence Diagram(s)sequenceDiagram
participant ComplexAllNullScheduler
participant LazyLevels
participant ComplexAllNullPageDecoder
participant DecodeComplexAllNullTask
ComplexAllNullScheduler->>LazyLevels: build and validate rep/def levels
ComplexAllNullScheduler->>ComplexAllNullPageDecoder: schedule level cursors
ComplexAllNullPageDecoder->>LazyLevels: seek rows and count visible items
ComplexAllNullPageDecoder->>DecodeComplexAllNullTask: pass level slices and run hints
DecodeComplexAllNullTask->>LazyLevels: extend repetition/definition vectors
LazyLevels-->>DecodeComplexAllNullTask: materialized level vectors
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. 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 1516-1545: Remove the num_values == 0 conditional in the
LevelCodec::Uncompressed branch and always compute expected_bytes via
expected_level_bytes(num_values, level_type), requiring compressed_bytes.len()
to equal it before constructing LazyLevels::Dense. Preserve the existing
invalid-input error context and expected/actual byte reporting.
🪄 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: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 40554288-d8f1-4e3f-8c7e-1d548acc4582
📒 Files selected for processing (3)
rust/lance-encoding/src/compression.rsrust/lance-encoding/src/encodings/logical/primitive.rsrust/lance-encoding/src/encodings/physical/rle.rs
Co-authored-by: Cursor <cursoragent@cursor.com>
Ali2Arslan
left a comment
Ali2Arslan
left a comment
There was a problem hiding this comment.
@Xuanwo mind taking a look again? I believe this should address your concerns (hopefully).
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/physical/rle.rs (1)
1559-1561: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winReject non-empty RLE block payloads when
num_values == 0.The parsed value/length buffers are passed to
decode_data, whose zero-value fast path returns before validating them. A frame with a valid header plus arbitrary trailing bytes is therefore silently accepted as an empty block, hiding corruption. Validate both parsed buffers are empty before this call and add a regression test.Suggested guard
let (values_buffer, lengths_buffer) = parse_rle_block_frame(&data)?; +if num_values == 0 + && (!values_buffer.is_empty() || !lengths_buffer.is_empty()) +{ + return Err(Error::invalid_input_source(format!( + "RLE block has non-empty buffers for zero values: values={} bytes, lengths={} bytes", + values_buffer.len(), + lengths_buffer.len(), + ).into())); +} let data = vec![values_buffer, lengths_buffer];🤖 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/physical/rle.rs` around lines 1559 - 1561, Update the RLE block decode path around parse_rle_block_frame and decode_data to reject non-empty values_buffer or lengths_buffer when num_values is zero, before invoking decode_data. Preserve the existing decoding behavior for nonzero counts, and add a regression test covering a valid header with trailing payload bytes.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/physical/rle.rs`:
- Around line 1559-1561: Update the RLE block decode path around
parse_rle_block_frame and decode_data to reject non-empty values_buffer or
lengths_buffer when num_values is zero, before invoking decode_data. Preserve
the existing decoding behavior for nonzero counts, and add a regression test
covering a valid header with trailing payload bytes.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 692aeb04-c3bc-44c3-aa1b-6fb8366f20eb
📒 Files selected for processing (3)
rust/lance-encoding/src/compression.rsrust/lance-encoding/src/encodings/logical/primitive.rsrust/lance-encoding/src/encodings/physical/rle.rs
Preserve adaptive level-cache test coverage while adopting main's MiniBlockChunkIndex. Co-authored-by: Cursor <cursoragent@cursor.com>
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 (3)
rust/lance-encoding/src/encodings/logical/primitive.rs (3)
8977-8990: 📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy liftAdd a pre-V2_2 complex-all-null decode regression.
This suite covers V2_2 only. Add a checked-in pre-V2_2 read path that exercises a V2_1-compat page where
num_rep_values == 0andnum_def_values == 0, usingcopy_test_data_to_tmpand a version-assertingdatagen.pyas required for compatibility tests.🤖 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 8977 - 8990, Extend test_complex_all_null_constant_def_round_trip with a checked-in pre-V2_2 compatibility read case using copy_test_data_to_tmp. Add the required datagen.py version assertion and fixture generation for a V2_1 page with num_rep_values == 0 and num_def_values == 0, then decode that fixture and verify the expected all-null complex values.Source: Coding guidelines
2178-2184: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftReturn an error instead of truncating the row plan.
This branch leaves
chunk_instructionspartial after only logging a warning. The downstream row-count check is adebug_assert_eq!, so release builds can silently return fewer rows. Propagate aResultand returnError::invalid_input_sourcecontaining the requested range,block_index, andnum_chunks.As per coding guidelines, invalid boundary values must be rejected with descriptive errors rather than silently adjusted.
🤖 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 2178 - 2184, Update the scheduling loop around the block_index >= num_chunks guard to return a Result error instead of logging and breaking. Propagate the error through the enclosing instruction-planning function, using Error::invalid_input_source with the requested range, block_index, and num_chunks, while preserving normal instruction generation for valid boundaries.Source: Coding guidelines
2418-2461: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftValidate metadata-derived counts with checked arithmetic.
items_in_page - countedandcounted += countare unchecked, while invalid log values are rejected only bydebug_assert!. Malformed page metadata can therefore panic in debug builds, wrap in release builds, or produce incorrect chunk boundaries. Return an error from this validation path, use checked arithmetic, and include the chunk count, item count, accumulated count, and offending log value in the error.As per coding guidelines, on-disk inputs must be validated and counters must use checked arithmetic.
🤖 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 2418 - 2461, Update analyze_value_counts and flat_value_counts_iter to validate metadata at runtime and return a validation error instead of relying on debug_assert!. Use checked shifts, additions, and subtraction for log-derived and accumulated counts; reject invalid logs and cases where counted values exceed items_in_page, reporting chunk count, item count, accumulated count, and offending log. Propagate the resulting error through callers and preserve valid-page chunk iteration behavior.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 8977-8990: Extend test_complex_all_null_constant_def_round_trip
with a checked-in pre-V2_2 compatibility read case using copy_test_data_to_tmp.
Add the required datagen.py version assertion and fixture generation for a V2_1
page with num_rep_values == 0 and num_def_values == 0, then decode that fixture
and verify the expected all-null complex values.
- Around line 2178-2184: Update the scheduling loop around the block_index >=
num_chunks guard to return a Result error instead of logging and breaking.
Propagate the error through the enclosing instruction-planning function, using
Error::invalid_input_source with the requested range, block_index, and
num_chunks, while preserving normal instruction generation for valid boundaries.
- Around line 2418-2461: Update analyze_value_counts and flat_value_counts_iter
to validate metadata at runtime and return a validation error instead of relying
on debug_assert!. Use checked shifts, additions, and subtraction for log-derived
and accumulated counts; reject invalid logs and cases where counted values
exceed items_in_page, reporting chunk count, item count, accumulated count, and
offending log. Propagate the resulting error through callers and preserve
valid-page chunk iteration behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: d5cde324-876f-4bb5-8f7b-db90f58fa0b3
📒 Files selected for processing (1)
rust/lance-encoding/src/encodings/logical/primitive.rs
|
Thanks a lot for working on this! One last thing I want to confirm is maybe we should gate this change in 2.3? The old reader might can't decode newly written file. |
This PR doesn't touch the writer side, so the bytes written should be same for both 2.2- and 2.3+ and newer RLE encodings already seem to have a 2.3 gate, so it should be fine? |
|
Great, thanks for the explanation. Let's move! |
Summary
Complex-all-null pages (nested/list/struct all-null columns) cached their rep/def levels by eagerly expanding RLE to
num_valuesu16s. A production heap profile attributed ~79% of in-use bytes toComplexAllNullSchedulerholding these expanded levels.This caches the levels lazily in coalesced RLE run form and drains them run-by-run:
rle.rs: extractparse_rle_block_framefromBlockDecompressor::decompress(pure refactor), and adddecode_rle_u16_runs, which parses an RLE frame into per-run(values, offsets)and coalesces adjacent equal-valued runs. The block encoder caps runs at 255, so a logically constant page arrives asceil(num_values / 255)runs; coalescing collapses them to O(distinct levels), making the cached form O(runs) instead of O(num_values).primitive.rs: introduceLevelCodec(Uncompressed | Rle | Block) so only RLE bypasses the eager decompressor, andLazyLevels(Dense | Rle) to hold the run form.initializebuildsLazyLevelsper level buffer; non-RLE paths still expand toDenseexactly as before.LevelCursor(seek_row_start/count_le_cursoradvance forward through runs), so a full page decode is O(runs) not O(rows) — this keeps the run-form representation from making per-row seeks costly. Emitted level slices, visible counts, theRepDefUnraveler, and theAllNulloutput are unchanged.The non-RLE path (uncompressed / other block compression) is untouched.
Test plan
cargo fmt --allcargo clippy -p lance-encoding --all-targets -- -D warnings(clean)cargo test -p lance-encoding— 402 passed, 0 failed, 5 ignoredNew tests:
decode_rle_u16_runs_matches_eager/_empty/_coalesces_equal_runslazy_levels_rle_matches_dense(Dense-vs-Rle parity)complex_all_null_drain_parity::drain_matches_reference(512-case proptest vs a brute-force drain reference, both Dense and Rle forms)lazy_levels_rle_is_compact(run-form footprint independent of logical length)test_complex_all_null_constant_def_round_trip(end-to-end all-null nested round-trip whose levels RLE-compress to a single run)Made with Cursor