[DISCUSS] Representing Shared State / `ExecutionPlan::reset_state`

[alamb]

Is your feature request related to a problem or challenge?

I brought this out into its own discussion so it wasn't lost on a PR comment

@Dandandan observes on #21351 (comment):

I don't like the additions in some recent PRs of reset_state on plans / sources etc.

In my opinion we should avoid setting mutable state on ExecutionPlans (and related data structures) & move that to the inner streams.

In an ideal scenario:

During any time we can clone a plan and share / run it without expecting any side effects (e.g. partially executed plans)
Concurrent runs of a query (ideallyexecutes) should result in the same result, even while a query is running.
In my PR #20481 I combined the shared state using TaskContext (which, in normal usages, will be created on each new query run), which I think still not is 100% the desired end state but avoids a lot of scenario's.

https://github.com/apache/datafusion/pull/20481/files#diff-e3975f93f1d4599fe6f19905e81f02574d65402e604e39c45b87448b3892d1adR70

Describe the solution you'd like

No response

Describe alternatives you've considered

No response

Additional context

No response

[alamb]

In my opinion we should avoid setting mutable state on ExecutionPlans (and related data structures) & move that to the inner streams.

I think actually most of the state is already on the inner streams. I believe the core usecase for mutable state on ExecutionPlans is making connections between the inner streams.

Sometimes this connection is for streams that come from the same ExecutionPlan -- e.g. the CollectLeft in joins

Sometimes the connection is for streams for different ExecutionPlans -- e.g dynamic filters need to source from a TopK operator

[alamb]

I combined the shared state using TaskContext (which, in normal usages, will be created on each new query run), which I think still not is 100% the desired end state but avoids a lot of scenario's.

I really like this idea of using TaskContext to store the shared state somehow. It maps quite nicely to existing patterns and I think would not be super disruptive. I think if we got the general purpose API right, that could also solve a lot of challenges (like how to serialize dynamic filter references, for example) -- e.g. the dynamic filters could get some sort of id and then the actual runtime connection could be made by looking up in the TaskContext

[alamb]

Here are some examples (found with codex, reviewed by me) of this shared state:

Details

Shared recursive CTE work table

This seems like the clearest example of explicit shared mutable operator/runtime state.

• WorkTable itself is mutable shared state (Mutex<Option<ReservedBatches>>)
• WorkTableExec stores Arc<WorkTable>

Hash join shared build-side state, especially CollectLeft

• HashJoinExec stores a shared left_fut: Arc<OnceAsync<JoinLeftData>> and optional dynamic_filter
• JoinLeftData contains shared runtime artifacts: hash map, visited bitmap, bounds, membership, and atomics
• In execute, PartitionMode::CollectLeft shares one left_fut across all probe partitions

Dynamic filter information on joins

• HashJoinExecDynamicFilter holds the shared DynamicFilterPhysicalExpr plus OnceLock<Arc<SharedBuildAccumulator>>
• The dynamic filter is also exposed as part of the plan expressions in apply_expressions
• In execute, the dynamic filter path lazily creates a shared SharedBuildAccumulator
• SharedBuildAccumulator is the cross-partition shared state used to accumulate bounds / membership and update the filter
• Filter updates happen in SharedBuildAccumulator when all build partitions report in

Other join operators with shared build-side state

These follow the same general pattern: plan owns a OnceAsync<_> that multiple output partitions/streams wait on.

CrossJoinExec

• CrossJoinExec has shared left_fut: OnceAsync<JoinLeftData>
• reset_state explicitly recreates left_fut
• execute shares the buffered left side through left_fut.try_once(...)

NestedLoopJoinExec

• NestedLoopJoinExec has shared build_side_data: OnceAsync<JoinLeftData>

PiecewiseMergeJoinExec

• PiecewiseMergeJoinExec has shared buffered_fut: OnceAsync<BufferedSideData>

SortExec

• SortExec::reset_state recreates the TopK dynamic filter and resets metrics
• The TopK self-filter is pushed down from operator-owned state in gather_filters_for_pushdown

AggregateExec

• AggregateExec::init_dynamic_filter can attach a dynamic_filter owned by the plan

[alamb]

I think @askalt @adriangb @gabotechs @xudong963 @geoffreyclaude @gene-bordegaray and others may have some thoughts about this topic as well

[geoffreyclaude]

I think @askalt @adriangb @gabotechs @xudong963 @geoffreyclaude @gene-bordegaray and others may have some thoughts about this topic as well

The main reason I worked on this was datafusion-tracing: recursive queries broke once plans were wrapped, because the runtime wiring depended on concrete WorkTableExec handling. That led to with_new_state in #16469 and the follow-up in #18803.

So I’m sympathetic to moving shared state off plans, but whatever replaces it should still support wrapper/custom ExecutionPlans without concrete type checks. TaskContext seems like a plausible direction if it preserves that.

[gabotechs]

In general, not baking shared state into execution plans is greatly beneficial to any system the relies on plan [de]serialization to work, for example FFI-based wrappers (https://github.com/apache/datafusion-python) or distributed engines (https://github.com/apache/datafusion-ballista or https://github.com/datafusion-contrib/datafusion-distributed).

Indeed, one of the main challenges for integrating dynamic filters with a distributed system is the collision with the assumption that data can simply be shared through Arc-based shared references.

I actually gave it a try some time ago in gabotechs#7, but abandoned the project as a #18739 was a prerequisite but was not very well received.

I can imagine how not relying on shared state inside can add some complexity if not done with care, but IMO it's something worth pursuing.