fix: let analyze_plan run under a caller-provided TaskContext - #8980
fix: let analyze_plan run under a caller-provided TaskContext#8980LuQQiu wants to merge 2 commits into
Conversation
`analyze_plan` executed the plan under a TaskContext derived solely from `LanceExecutionOptions`. Nodes that read a session-config extension at execute time (e.g. distributed routing identity) therefore never saw it, errored during `execute`, and `AnalyzeExec` absorbed that per-partition failure into an empty, unexecuted plan tree — the analyze output showed the node with no metrics and no children, silently hiding the real error. Add `analyze_plan_with_context`, which runs the plan under a caller- supplied `TaskContext` when provided (falling back to the options-derived one otherwise). `analyze_plan` keeps its existing signature and behavior. Adds a regression test: a node that requires a session-config extension produces an empty tree via the context-less path and executes correctly when the extension is supplied through the new entry point.
|
The failed tests are not related to my change |
![lance-gatekeeper[bot]](https://avatars.githubusercontent.com/in/4433141?s=60&v=4){kind=small}
There was a problem hiding this comment.
❌ Gate recommendation: request changes.
The helper forwards the caller-provided context as intended, but the regression coverage does not enforce that contract. Make successful child execution observable so removing context forwarding fails the test.
| let report = analyze_plan_with_context(plan, options, Some(task_ctx)) | ||
| .await | ||
| .expect("analyze should succeed when the extension is present"); | ||
| assert!(report.contains("NeedsExtensionExec")); |
There was a problem hiding this comment.
This assertion only verifies the static plan text; the missing-context report contains the same NeedsExtensionExec node. I temporarily changed the helper to ignore the supplied context and reran cargo test -p lance-datafusion test_analyze_plan_uses_provided_task_context -- --nocapture; the test still passed. Please make execution observable—for example, store an Arc<AtomicBool> in NeedsExtensionExec, set it after the extension check, and assert false after the context-less call and true after the supplied-context call—so this regression test fails if forwarding is removed.
There was a problem hiding this comment.
Fixed in 80ba7dfce: the regression now observes successful node execution, and I verified that the test fails when the supplied TaskContext is not forwarded.
lance-gatekeeper
Bot
added
the
K-changes
The prior assertion only checked static plan text, which is identical whether or not the context is forwarded (both reports contain the node). Add an AtomicBool the node sets only after passing the extension check, and assert it is false after the context-less call and true after the supplied-context call, so dropping context forwarding fails the test.
|
Good catch — the static-text assertion didn't enforce the contract. Fixed in 80ba7df: |
lance-gatekeeper
Bot
removed
the
K-changes
There was a problem hiding this comment.
✅ Gate recommendation: approve.
The updated regression now observes successful node execution and fails when the supplied context is not forwarded. This addresses the prior finding; the caller-provided TaskContext path is covered while the existing context-less API remains unchanged.
lance-gatekeeper
Bot
added
the
K-approved
| let report = analyze_plan_with_context(plan, options, Some(task_ctx)) | ||
| .await | ||
| .expect("analyze should succeed when the extension is present"); | ||
| assert!(report.contains("NeedsExtensionExec")); |
There was a problem hiding this comment.
isn't this trivially true?
3 participants
Problem
analyze_planexecutes the plan under aTaskContextderived solely fromLanceExecutionOptions. Execution nodes that read a session-config extension at execute time (for example, a distributed routing identity) never see it, so they error insideexecute.AnalyzeExecabsorbs that per-partition failure and returns an empty, unexecuted plan tree — the analyze output shows the node withmetrics=[]and no children, silently hiding the real error rather than surfacing it.Concretely, a distributed FTS/ANN/scalar-index query executes and returns results normally, but
analyze_planon the same query returns a stub tree with the remote subtree unexecuted (all-zero metrics, no leaf query node).Fix
Add
analyze_plan_with_context(plan, options, task_context), which runs the plan under a caller-suppliedTaskContextwhenSome, falling back to the options-derived context otherwise.analyze_plankeeps its exact signature and behavior (it now delegates withNone).Callers that attach execution-time session-config extensions can then hand the same enriched context to analyze that they use for the real read path, so the analyzed plan executes identically to the query.
Test
Adds a regression test with a node that requires a session-config extension:
NeedsExtensionExec, metrics=[]);analyze_plan_with_contextwith the extension present, it executes and appears in the report.