fix(engine): bound hydration's Lance fetch before it allocates

[ragnorc]

Why

HydrateExec (from #887) fetches the return-only columns of the rows a limit kept. It planned each fetch window from the widest row seen so far, then took the window's distinct rows with one TakeBuilder::execute(). That call allocates the whole result, and Lance reads ahead up to its default I/O buffer (32 MiB per I/O thread, 2 GiB on a cloud store), all before the query pool sees a byte. When rows widen after a narrow run, a window planned from the narrow rows overshoots the chunk bound by the whole window. The new test below reproduced this: a top-k of 150 under a 16 MiB pool held 3,278,868 bytes in one fetch against a 2 MiB bound. The #887 review thread named this as the remaining limit.

What changes

The fetch is bounded by design, not by the widths already seen.

• Row-by-row decode, charged as it goes. Hydration reads through FileFragment::open and FragmentReader::take(offsets, 1, ..), one row per decode task. The reader merges overlays and applies the deletion vector, as TakeBuilder does. Each decoded row is copied at its own size through the pool's checked take, then charged before the fetch waits for the next row. The copy matters because a Lance row can slice a buffer of the whole read, and the pool charges a buffer by its allocation's capacity. At most 8 rows decode at once.
• Stop and replan. When the decoded rows pass half of the chunk's share, the fetch stops and drops the read. The window is then planned again from the widest row decoded.
• Read-ahead inside the bound. An engine ScanScheduler caps Lance's read-ahead at a quarter of hydrate_chunk_bytes. That quarter is reserved in the pool for the operator's life as hydrate read-ahead.
• The bound. The operator holds at most hydrate_chunk_bytes (the declared retained_limit) plus the row that crosses it. Lance treats its I/O buffer as a soft limit: the request at the head of its queue always goes, and Lance can coalesce it from many rows' ranges, up to its 16 MiB maximum I/O size. So the gauge bounds what hydration holds, not every buffer in Lance or the process. The review observed 2.69 MB read before the first wide row reached hydration, with a 512 KiB read-ahead share.
• Checked interleave. Copies to output rows go through a new checked WorkMemory::interleave_once, which admits one copy per pick before Arrow builds it. take_bytes is gone; hydration was its only caller.
• Fallback for inherited files. Lance 11 ignores a passed scheduler for a data file reached through a base_id, opening its own max-bandwidth scheduler instead (open_current_file_reader). A Lance branch reaches its parent's files this way. In OmniGraph only table forks created before v11 hold such files; graph branches stage detached commits in the table's own base. A fragment with such a file is read one row per request and counted in a new single_row_reads operator metric.
• lance-io becomes a regular engine dependency. lance already links it.

Evidence

• engine_v2_memory.rs::a_hydrated_fetch_stops_at_the_chunk_bound_when_rows_widen: red on main (one fetch held 3,278,868 bytes against a 2,097,152-byte bound), green here. To reproduce the red run, the old code needs one observer line, plan.peak.set_max(fetched_bytes);, just before its if fetched_bytes > hard check. Without it, the old gauge never records the fetches it discards, and the test passes on main. The narrow rows sit in their own fragment. In a shared one, Lance's shared-buffer measurement makes the narrow seed look wide, which hides the overshoot.

• lance_surface_guards.rs::fragment_reads_bypass_the_passed_scheduler_for_inherited_files pins the Lance behaviour behind the fallback. It goes red at the Lance release that honours the passed scheduler for every base, and then the fallback can be removed.

• In-source operators::hydrate::tests::fragment_rows_reads_inherited_and_own_fragments_by_address exercises both read shapes against a real Lance branch.

• engine_v2_memory.rs::hydration_on_a_graph_branch_reads_in_coalesced_requests: on a graph branch, single_row_reads stays 0.

• The memory.rs unit test now also covers the checked interleave: per-pick admission, and refusal before Arrow allocates.

• The GQT case hydrate_return_columns_above_the_limit.gqt gains a step that hydrates on a branch.

• Request counts are unchanged. On the slow wide-column case, run on the in-memory test store with --measure: limit 1 takes 202 store requests and the top-101 takes 204, on main and on this branch alike.

• Local, before merging main:

  • fmt, workspace clippy with -D warnings, check-docs.py and typos are clean.
  • All 62 Lance guards, the engine suites, the 461 lib tests and the DST suite (51 scenarios, including the cost golden) pass.
  • The planner and GQT-core suites pass, as do both slow cases.
  • The full GQT corpus didn't run clean on the shared machine (load average 30–300). Every failure was a 10 s timeout in a case unrelated to hydration, with a different set each run. Each re-ran alone in 2–6 s.

• After merging main (52beecfe):

  • fmt, clippy, check-docs.py and typos are clean.
  • All 64 Lance guards, forbidden_apis, engine_v2, engine_v2_memory, engine_v2_plan_replay, end_to_end, search, branching and the 481 lib tests pass.
  • The GQT hydration case and both slow cases pass.
  • The full GQT corpus at default parallelism passed 299 of 308; all 9 failures were timeouts, and each passes alone in about 5.5 s.
  • runner_dispatch::concurrent_block_refusals failed under load. Its wall-clock budget is known to be load-sensitive.

Follow-ups

• Declared residency for every query node (plan-time bounds) still needs an RFC.

[ragnorc]

Recommendation: Approve once the remaining workspace CI job passes, with the non-blocking documentation note below. I found no blocking correctness defect at 3fc6676d9be87686afc30f4278891cd10992fdfe.

This PR changes how a limited query loads its return-only columns. The old path chose a window from earlier row widths, then allocated that whole window before charging the query pool. A narrow run followed by wide rows could exceed the chunk budget before the retry. The new path decodes one row per task, copies it into checked memory, and stops when the fetched rows cross the target. It then replans the window. Users keep the same rows, order, duplicates and branch snapshot, with less decoded data retained at once.

The important design change is in fetch_window: past widths now guide performance; they no longer decide how much decoded data must be retained before checking the bound. Checked interleave admits copies per output pick, including repeated join rows. I traced the pinned snapshot, physical addresses, null addresses, duplicate picks, the merged fusion caller and branch paths. In pinned Lance 11.0.0, the fragment reader applies overlays and deletions. Its inherited-file open path does bypass the supplied scheduler. The fallback and its upstream guard address a real substrate constraint. Current graph branches retain the coalesced path.

The tradeoff is more small decode tasks, a row copy, and possible repeated reads after replanning. The change removes correctness dependence on a width estimate and removes take_bytes. It adds direct scheduler coupling, two concurrency constants, checked interleave and a version-specific fallback. I judge that a net reduction in liability for decoded-row retention. There is no new persistent state or GQ API. The fallback has a concrete removal condition in its Lance guard. Five similar fixes should reuse the checked memory kernels and substrate tests, keeping allocation policy at the operator that owns it.

Local validation used Rust 1.97.1, CARGO_TARGET_DIR=/tmp/review-928/target, and RUSTFLAGS='--cfg tokio_unstable --cfg tokio_unstable'. These passed at the reviewed head:

• cargo test --locked -p omnigraph-engine --features failpoints -j 3 --test engine_v2_memory hydrat -- --nocapture — 4 tests.
• cargo test --locked -p omnigraph-engine --features failpoints -j 3 --test lance_surface_guards fragment_reads_bypass_the_passed_scheduler_for_inherited_files -- --exact --nocapture — 1 test.
• cargo test --locked -p omnigraph-engine --features failpoints -j 3 --lib engine::operators::hydrate::tests:: -- --nocapture — 1 test.
• cargo test --locked -p omnigraph-engine --features failpoints -j 3 --lib engine::operators::memory::tests:: -- --nocapture — 13 tests.
• cargo test --locked -p omnigraph-engine --features failpoints -j 3 --test engine_v2_plan_replay — 33 tests.
• After cargo build --locked -p omnigraph-gqt --features omnigraph/failpoints --bin omnigraph-gqt -j 3, /tmp/review-928/target/debug/omnigraph-gqt cases/v2/planner/hydrate_return_columns_above_the_limit.gqt, from the GQT crate — passed. /tmp/review-928/target/debug/omnigraph-gqt cases/v2/planner/hydrate_return_columns_above_a_fusion.gqt also passed, covering the caller added by the main merge.

The before/after proof needs one detail in the reproduction instructions. On base 52beecfefc2d8b5d0617448d82b0d5e715215976, copying the new widening test alone passes: the old gauge omits discarded fetches. I added only plan.peak.set_max(fetched_bytes); immediately before the old if fetched_bytes > hard check, to observe those fetches without changing their behavior. Then this command failed at the predicted assertion:

cargo test --locked -p omnigraph-engine --features failpoints -j 3 --test engine_v2_memory a_hydrated_fetch_stops_at_the_chunk_bound_when_rows_widen -- --exact --nocapture

It reported 3,278,868 bytes against a 2,097,152-byte bound. The same assertion passes at the reviewed head, including after all temporary edits were restored. Please include that observer in the red-reproduction instructions. The existing Rust owner is appropriate: GQT cannot set this per-query pool or inspect the gauge. The added GQT step checks branch results; it is not the memory regression proof.

The storage limit remains soft. Lance can load a coalesced request before its rows are charged, and its scheduler refunds the reservation when it delivers the bytes. In a temporary variant at the earlier 17dbf1e8 head (the hydration code is unchanged) with distinct, less compressible payloads, the scheduler recorded 2,692,011 bytes read before the first wide row reached hydration, with a 524,288-byte read-ahead setting. This is an I/O observation, not an RSS measurement. The chunk gauge does not establish a hard bound on all Lance buffers or the process. This agrees with the stated head-request exception and upstream memory guidance.

Separately, storage/cloud checks, GQT CI and DST CI are green for this exact head. The Test Workspace job is still running; the overall CI gate is not yet confirmed. I did not repeat cloud tests locally or verify the latency/request-count performance claims. The checkout is clean; no source or test changes were pushed.