perf(filtered-read): reuse plan-time fragment metadata for the read stream#7792
perf(filtered-read): reuse plan-time fragment metadata for the read stream#7792LuQQiu wants to merge 1 commit into
Conversation
📝 WalkthroughWalkthroughFiltered reads now retain loaded fragment metadata, calculate fragment-planning statistics, and conditionally consolidate sparse take-shaped batches. Stream construction reuses cached fragments when available, while dense and non-indexed paths retain existing batching behavior. ChangesTake-shaped filtered reads
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: 🚥 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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rust/lance/src/io/exec/filtered_read.rs (1)
1978-2050: 🚀 Performance & Scalability | 🔵 TrivialConfirm the streaming-latency trade-off for sparse masked reads is acceptable.
When
is_sparse_planis true and total planned rows are belowtarget,coalesce_batchesonly emits once the inner stream is exhausted (ready_to_emitrequiresexhausted && !buffered.is_empty()), soconsolidated_streamwithholds all output until the entire masked scan across all touched fragments completes. Previously, per-fragment batches would stream out incrementally as each fragment finished. For a caller relying on early termination (e.g. aLimitExecsitting above a read whosescan_range_after_filtercouldn't be pushed down because of a recheck filter), this increases time-to-first-batch even though it doesn't change total I/O volume.This is likely intentional for the take-shaped scenario this PR targets, but worth confirming it doesn't regress latency for interactive/streaming consumers.
🤖 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/src/io/exec/filtered_read.rs` around lines 1978 - 2050, Review the sparse masked-read path in the stream construction around is_sparse_plan and consolidated_stream, and confirm that withholding batches until all touched fragments finish is acceptable for interactive consumers and upstream early termination such as LimitExec. If the latency regression is not acceptable, constrain consolidation to the intended take-shaped scenario or preserve incremental per-fragment emission for other masked reads; otherwise document or validate the intentional trade-off without changing unrelated batch-size behavior.
🤖 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/src/io/exec/filtered_read.rs`:
- Around line 2639-2709: The test_take_shaped_mask_consolidation coverage only
uses the default threading mode; extend it with a
FilteredReadThreadingMode::MultiplePartitions case. Configure the read options
with MultiplePartitions(n), execute the same take-shaped mask, and verify
consolidation preserves correct, disjoint row sets and expected batch results
across partitions without relying on single-partition ordering assumptions.
---
Outside diff comments:
In `@rust/lance/src/io/exec/filtered_read.rs`:
- Around line 1978-2050: Review the sparse masked-read path in the stream
construction around is_sparse_plan and consolidated_stream, and confirm that
withholding batches until all touched fragments finish is acceptable for
interactive consumers and upstream early termination such as LimitExec. If the
latency regression is not acceptable, constrain consolidation to the intended
take-shaped scenario or preserve incremental per-fragment emission for other
masked reads; otherwise document or validate the intentional trade-off without
changing unrelated batch-size behavior.
🪄 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: 7cd1a99c-22d0-4c8b-9d69-5cadc0d469f7
📒 Files selected for processing (1)
rust/lance/src/io/exec/filtered_read.rs
| /// Take-shaped masked reads consolidate their tiny per-fragment batches; | ||
| /// few-fragment and dense masked reads keep per-fragment boundaries. | ||
| #[test_log::test(tokio::test)] | ||
| async fn test_take_shaped_mask_consolidation() { | ||
| // 20 fragments x 2000 rows, value = global row number | ||
| let tmp_path = TempStrDir::default(); | ||
| let data = gen_batch() | ||
| .col("value", array::step::<UInt32Type>()) | ||
| .into_reader_rows(RowCount::from(2000), BatchCount::from(20)); | ||
| let dataset = Arc::new( | ||
| Dataset::write( | ||
| data, | ||
| tmp_path.as_str(), | ||
| Some(WriteParams { | ||
| max_rows_per_file: 2000, | ||
| ..Default::default() | ||
| }), | ||
| ) | ||
| .await | ||
| .unwrap(), | ||
| ); | ||
|
|
||
| let mask_input = |addrs: Vec<u64>| -> Arc<dyn ExecutionPlan> { | ||
| let covered: RoaringBitmap = dataset.fragments().iter().map(|f| f.id as u32).collect(); | ||
| let batch = | ||
| IndexExprResult::exact(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter(addrs))) | ||
| .serialize(&covered, IndexExprResultWireFormat::default()) | ||
| .unwrap(); | ||
| let schema = batch.schema(); | ||
| let stream = futures::stream::once(async move { Ok(batch) }); | ||
| Arc::new(OneShotExec::new(Box::pin(RecordBatchStreamAdapter::new( | ||
| schema, stream, | ||
| )))) | ||
| }; | ||
| let run = |input: Arc<dyn ExecutionPlan>| { | ||
| let dataset = dataset.clone(); | ||
| async move { | ||
| // Pin the batch size so batch-count assertions don't depend | ||
| // on LANCE_DEFAULT_BATCH_SIZE | ||
| let options = FilteredReadOptions::basic_full_read(&dataset).with_batch_size(2000); | ||
| let plan = | ||
| FilteredReadExec::try_new(dataset.clone(), options, Some(input)).unwrap(); | ||
| let stream = plan.execute(0, Arc::new(TaskContext::default())).unwrap(); | ||
| stream.try_collect::<Vec<_>>().await.unwrap() | ||
| } | ||
| }; | ||
| let addr = |frag: u32, offset: u32| u64::from(RowAddress::new_from_parts(frag, offset)); | ||
|
|
||
| // Take shape: 20 fragments, 2 rows each -> one consolidated batch, | ||
| // rows in fragment order | ||
| let addrs: Vec<u64> = (0..20u32).flat_map(|f| [addr(f, 3), addr(f, 7)]).collect(); | ||
| let batches = run(mask_input(addrs)).await; | ||
| assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 40); | ||
| assert_eq!(batches.len(), 1); | ||
| let expected = | ||
| UInt32Array::from_iter_values((0..20u32).flat_map(|f| [f * 2000 + 3, f * 2000 + 7])); | ||
| assert_eq!(batches[0].column(0).as_ref(), &expected); | ||
|
|
||
| // Too few fragments -> inline path, one batch per fragment | ||
| let batches = run(mask_input(vec![addr(0, 3), addr(1, 7)])).await; | ||
| assert_eq!(batches.len(), 2); | ||
|
|
||
| // Dense (2000 planned rows per fragment) -> inline path | ||
| let addrs: Vec<u64> = (0..8u32) | ||
| .flat_map(|f| (0..2000u32).map(move |o| addr(f, o))) | ||
| .collect(); | ||
| let batches = run(mask_input(addrs)).await; | ||
| assert_eq!(batches.len(), 8); | ||
| assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 16000); | ||
| } | ||
|
|
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Extend consolidation test coverage to MultiplePartitions threading mode.
test_take_shaped_mask_consolidation only exercises the default OnePartitionMultipleThreads mode. Each partition wraps its own inner stream independently in obtain_stream, but under FilteredReadThreadingMode::MultiplePartitions the underlying task_stream is shared and pulled by whichever partition wins the lock race, so per-partition ordering/content assumptions differ from the single-partition case. Consider adding a case with threading_mode: FilteredReadThreadingMode::MultiplePartitions(n) to confirm consolidation still produces correct, disjoint row sets per partition.
🤖 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/src/io/exec/filtered_read.rs` around lines 2639 - 2709, The
test_take_shaped_mask_consolidation coverage only uses the default threading
mode; extend it with a FilteredReadThreadingMode::MultiplePartitions case.
Configure the read options with MultiplePartitions(n), execute the same
take-shaped mask, and verify consolidation preserves correct, disjoint row sets
and expected batch results across partitions without relying on single-partition
ordering assumptions.
…anch Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…tream Masked reads loaded every fragment's metadata twice per query: once in get_or_create_plan_impl to compute the plan and again in FilteredReadStream::try_new to build the read. The plan cell now keeps the loaded fragments alongside the plan and stream construction reuses them. The distributed path (with_plan, plan computed on another node) carries no fragments and loads them at stream construction as before. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
LuQQiu
force-pushed
the
lu/take-single-load
branch
from
e955296 to
8b5c856
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/src/io/exec/filtered_read.rs`:
- Around line 2015-2033: Add a log::debug! inside the sparse-plan branch where
is_sparse_plan is computed, logging touched_fragments, planned_rows, and the
resolved batch_target_rows when the heuristic evaluates true. Keep consolidation
behavior unchanged and ensure the message identifies that masked-read
consolidation was triggered.
- Around line 337-446: Extract the consolidation utilities from filtered_read.rs
into a dedicated consolidate submodule, including consolidated_stream,
coalesce_batches, and the CONSOLIDATE_MIN_FRAGMENTS and
CONSOLIDATE_MAX_AVG_PLANNED_ROWS_PER_FRAGMENT constants. Update their callers
and visibility/imports so planning and execution behavior remains unchanged
while the main module retains only the wiring.
🪄 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: 4ccf83f8-cf9c-453e-bafd-c102cffd94be
📒 Files selected for processing (1)
rust/lance/src/io/exec/filtered_read.rs
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 2
🤖 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/src/io/exec/filtered_read.rs`:
- Around line 2015-2033: Add a log::debug! inside the sparse-plan branch where
is_sparse_plan is computed, logging touched_fragments, planned_rows, and the
resolved batch_target_rows when the heuristic evaluates true. Keep consolidation
behavior unchanged and ensure the message identifies that masked-read
consolidation was triggered.
- Around line 337-446: Extract the consolidation utilities from filtered_read.rs
into a dedicated consolidate submodule, including consolidated_stream,
coalesce_batches, and the CONSOLIDATE_MIN_FRAGMENTS and
CONSOLIDATE_MAX_AVG_PLANNED_ROWS_PER_FRAGMENT constants. Update their callers
and visibility/imports so planning and execution behavior remains unchanged
while the main module retains only the wiring.
🪄 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: 4ccf83f8-cf9c-453e-bafd-c102cffd94be
📒 Files selected for processing (1)
rust/lance/src/io/exec/filtered_read.rs
🛑 Comments failed to post (2)
rust/lance/src/io/exec/filtered_read.rs (2)
337-446: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Consider extracting the consolidation utilities into a dedicated submodule.
consolidated_stream/coalesce_batches(plus the two threshold constants) are ~110 lines of self-contained, new algorithmic logic added directly into an already very large file (filtered_read.rsspans thousands of lines). Traced the state machine by hand (partial-buffer emission on oversized incoming batches, single-item passthrough, exhaustion handling) — the logic itself is correct and well-commented.As per coding guidelines, "Extract substantial new logic such as bin packing or scheduling into dedicated submodules instead of inlining it into large files." Moving this to e.g.
filtered_read/consolidate.rswould improve discoverability and keep the main file focused on planning/execution wiring.🤖 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/src/io/exec/filtered_read.rs` around lines 337 - 446, Extract the consolidation utilities from filtered_read.rs into a dedicated consolidate submodule, including consolidated_stream, coalesce_batches, and the CONSOLIDATE_MIN_FRAGMENTS and CONSOLIDATE_MAX_AVG_PLANNED_ROWS_PER_FRAGMENT constants. Update their callers and visibility/imports so planning and execution behavior remains unchanged while the main module retains only the wiring.Source: Coding guidelines
2015-2033: 📐 Maintainability & Code Quality | 🔵 Trivial
Consider a debug log when the sparse-plan consolidation heuristic fires.
The
is_sparse_plandecision (touched_fragments/planned_rows vs. thresholds) is computed once per query but never logged anywhere, making it hard to confirm in production whether a given masked read actually triggered consolidation. Alog::debug!withtouched_fragments,planned_rows, and the resolvedbatch_target_rowswould help diagnose unexpected batch shapes without needing to re-derive the heuristic from metrics.🤖 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/src/io/exec/filtered_read.rs` around lines 2015 - 2033, Add a log::debug! inside the sparse-plan branch where is_sparse_plan is computed, logging touched_fragments, planned_rows, and the resolved batch_target_rows when the heuristic evaluates true. Keep consolidation behavior unchanged and ensure the message identifies that masked-read consolidation was triggered.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
1 participant
Problem
Masked reads (
FilteredReadExecwith aRowSetinput) load every fragment's metadata twice per query:get_or_create_plan_implloads all fragments to compute the plan (plan_scanneeds deletion vectors / physical row counts), then throws the loaded set away;FilteredReadStream::try_newimmediately loads the same fragments again to build the read stream.Where it was found and how big it is
Found while stage-timing take-shaped masked reads (an exact row-id mask selecting ~100 scattered rows) on a 100M-row, 100-fragment dataset, 64 concurrent queries, page-cached NVMe. Per-query planning cost breaks down as:
The two loads together are 0.74 ms — ~64% of planning cost, all of it running in the first poll of the stream (i.e., on the caller's task). For a take-shaped query whose whole read is only a couple of milliseconds this is a significant fixed tax; it scales linearly with fragment count and query rate. Long scans amortize it to noise, which is why it went unnoticed.
Fix
The plan cell now keeps the loaded fragments alongside the resolved plan, and stream construction reuses them. The distributed path (
with_plan, plan computed on another node, e.g. a query node planning for executor nodes) carries no fragment metadata and loads at stream construction exactly as before — the plan phase stays independently callable.Tests
filtered_readsuite (28, single-threaded) andscannersuite (169) green; fmt + clippy clean. No behavior change — the same loaded set is used, just not loaded twice.🤖 Generated with Claude Code
https://claude.ai/code/session_015RpPYGnQ9JYaP9B2BbujBn