fix: remove WorkTableExec special case in reset_plan_states
#18803
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Remove the concrete type check in `reset_plan_states` that prevented `WorkTableExec` from being reset, enabling wrapper and custom nodes to participate correctly in recursive query execution. Motivation: Commit 921f4a0 introduced `with_new_state()` as a generic mechanism for injecting run-time state into execution plans, enabling any custom node to participate in recursive query execution. However, `reset_plan_states()` retained a concrete type check `plan.as_any().is::<WorkTableExec>()` that only recognized unwrapped `WorkTableExec` instances. This prevented wrapper or third-party nodes from being properly reset across iterations, undermining the extensibility goals of commit 921f4a0. Changes: - Remove the special case for `WorkTableExec` in `reset_plan_states()`, resetting all nodes uniformly via `reset_state()`. - Simplify the function from 9 lines to 4 lines, eliminating the type check and conditional logic. - Preserve `WorkTable` state correctly through the `Arc<WorkTable>` reference established by `with_new_state()` during initialization. - Rely on `WorkTableExec::with_new_children()` returning `Arc::clone(&self)` to avoid unnecessary recreation of the underlying plan. This change restores full support for wrapper nodes in recursive query execution, allowing external crates and custom execution plans to wrap `WorkTableExec` without breaking state management. The removed special case was a micro-optimization that saved one `Arc::clone()` per iteration but compromised correctness for wrapper nodes.
geoffreyclaude
force-pushed
the
fix/recursive_query
branch
from
213f5da to
b2e7681
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
This PR fixes an oversight in #16469 that broke wrapper nodes in recursive queries.
When
with_new_state()was added as a generic state injection mechanism,reset_plan_states()kept a concrete type check forWorkTableExec. This means wrapper nodes that delegateas_any()to their inner node won't be recognized, causing them to keep stale state across iterations. This breaks recursive queries for external crates that wrap execution plans (like tracing or monitoring tools).Rationale for this change
The
reset_plan_states()function usesplan.as_any().is::<WorkTableExec>()to decide which nodes to skip resetting. This works fine for bareWorkTableExecbut fails when it's wrapped by custom nodes as the type check can't see through the wrapper.This defeats the whole point of
with_new_state(), which was designed to let wrapper and third-party nodes participate in recursive queries without concrete type checks.The special case was added to save one
Arc::clone()per iteration, but it's not worth it sinceWorkTableExec::with_new_children()already just returnsArc::clone(&self)anyway. The sharedWorkTablestate is preserved through theArc<WorkTable>reference, so the optimization doesn't buy us much while breaking wrapper nodes.What changes are included in this PR?
WorkTableExectype check inreset_plan_states()- just reset all nodes uniformly viareset_state().WorkTablestate stays correct becauseWorkTableExec::with_new_children()returnsArc::clone(&self).Are these changes tested?
Yes, covered by existing tests:
WorkTableExecis identical.Since this restores functionality for wrapper nodes rather than changing core DataFusion behavior, it fixes test failures in external crates (like
datafusion-tracing) without needing new tests here.Are there any user-facing changes?
No Breaking Changes:
Benefits:
WorkTableExecand implementwith_new_state()without breaking recursive queries.with_new_statea trait method forExecutionPlan#16469 was designed to provide.