fix: make with_new_state a trait method for ExecutionPlan

[geoffreyclaude]

Which issue does this PR close?

This PR addresses the extensibility limitation that causes recursive queries to fail when used with the datafusion-tracing crate, as reported in datafusion-contrib/datafusion-tracing#5.

While there isn't a specific DataFusion issue filed, this change enables external crates like datafusion-tracing to properly support recursive queries by implementing work table injection in their wrapper nodes.

Rationale for this change

Currently, recursive query execution in RecursiveQueryExec uses downcasting to find WorkTableExec nodes and inject work tables. This approach fails when execution plans are wrapped by custom execution nodes, as is done by the datafusion-tracing crate, causing recursive queries to fail with "Unexpected empty work table" errors.

The fundamental issue is that RecursiveQueryExec assumes direct access to concrete WorkTableExec types, but wrapper nodes (like those used for tracing) break this assumption. This limits DataFusion's extensibility for recursive queries and prevents external crates from properly supporting this feature.

What changes are included in this PR?

Commit 1

• Add with_work_table as a public trait method on ExecutionPlan with a default implementation returning None.
• Update WorkTableExec to implement the trait method, creating a new instance with the provided work table.
• Modify assign_work_table logic in RecursiveQueryExec to use the trait method instead of down-casting, enabling support for wrapper nodes.
• Publicly re-export WorkTable so external implementers can access it without knowing internal module structure.
• Add comprehensive documentation explaining the purpose and usage of with_work_table.

Commit 2 (follow-up/generalisation)

• Replace the specialised with_work_table API with a generic with_new_state(&self, state: Arc<dyn Any + Send + Sync>) on ExecutionPlan.
  The default implementation still returns None.
• Implement with_new_state in WorkTableExec by down-casting the supplied state to an Arc<WorkTable>.
• Refactor RecursiveQueryExec::assign_work_table to use the new generic API.
• Update documentation throughout to reflect the generic nature of state injection and keep WorkTableExec as an illustrative example.
• Maintain the public re-export of WorkTable from the root of physical-plan so custom external nodes can downcast Arc<dyn Any + Send + Sync> into Arc<WorkTable>

Are these changes tested?

The changes are covered by existing tests:

• All existing recursive query tests continue to pass, ensuring backward compatibility
• The WorkTableExec implementation is tested through existing recursive query execution paths
• The trait method follows the same logic as the previous direct implementation

This change enables functionality that was previously broken (recursive queries with instrumentation), so it fixes existing test scenarios in external crates rather than requiring new DataFusion-specific tests.

Compiling the datafusion-tracing crate against this branch locally enables full instrumentation of recursive queries, with one "span group" per recursion, as demonstrated by the Jaeger screenshot below for the following basic recursive SQL:

WITH RECURSIVE numbers(n) AS (
    SELECT 1 AS n
    UNION ALL
    SELECT n + 1 FROM numbers WHERE n < 3
) SELECT n FROM numbers

Are there any user-facing changes?

API Addition (Non-breaking):

• Adds with_new_state method to the ExecutionPlan trait with a default None-returning implementation
• External crates can down-cast any late-bound state they need (e.g. WorkTable) inside their own ExecutionPlan implementations.

Benefits for Users:

• External crates can now implement with_new_state in their custom execution plans to support recursive queries
• Instrumentation and wrapper nodes can properly participate in recursive query execution
• Enables recursive queries to work with tracing, monitoring, and other execution plan decorators

No Breaking Changes:

• All existing code continues to work unchanged
• The new trait method has a default implementation, so existing ExecutionPlan implementations don't need updates
• Internal behavior for WorkTableExec remains identical

This change primarily enables extensibility that was previously impossible, resolving compatibility issues between recursive queries and external crates.

[alamb]

Thanks @geoffreyclaude -- the problem and solution make sense to me. However, I do wonder if we should try and make the API a little more general.

What do you think?

[alamb]

As this comment notes, this method is quite specific to WorkTableExec

I wonder if it might be a good idea to make a slightly more general API for updating state

Maybe something like

fn with_updated_state(&self, state: Arc<dyn Any>) ->  Option<Arc<dyn ExecutionPlan>> 

This would allow any arbitrary state to be updated and would make it eaiser for wrapper nodes like that in Datafusion tracing to pass through state without having to know anything about what is in it 🤔

[geoffreyclaude]

Hi @alamb,
I'm glad you mention the Any solution! I was initially going with something similar, but it seemed to me only the WorkTableExec needed this, so I didn't want to over-generalize.
However, moving to the Any solution anytime in the future (maybe because this late update pattern becomes more frequent for other concrete nodes) will mean breaking the specific WorkTableExec implementation.
So for the sake of forward compatibility, I very much agree with your suggested signature! I'll update the PR and mention you once done.

[geoffreyclaude]

@alamb PR's ready for a second review 👀

[alamb]

Makes sense to me @geoffreyclaude -- and thank you for the very clear and thorough comments