Pass execution props to scalar subqueries#23622
Conversation
Add an ignored regression test showing that per-plan scalar subquery state is not available when an extension planner lowers its expressions through PhysicalPlanner. AI Disclosure: This code was written in part by an AI agent.
Wrap the PhysicalPlanner passed to extension callbacks so expression lowering uses the current plan's ExecutionProps, including scalar subquery state. Enable the extension scalar-subquery regression test. AI Disclosure: This code was written in part by an AI agent.
| #[tokio::test] | ||
| async fn scalar_subquery_in_extension_expr_plans() -> Result<()> { | ||
| let subquery = LogicalPlanBuilder::empty(true) | ||
| .project(vec![lit(42_i32)])? | ||
| .build()?; | ||
| let logical_plan = LogicalPlan::Extension(Extension { | ||
| node: Arc::new(NoOpExtensionNode { | ||
| expressions: vec![scalar_subquery(Arc::new(subquery))], | ||
| ..Default::default() | ||
| }), | ||
| }); | ||
| let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new( | ||
| ExpressionExtensionPlanner, | ||
| )]); | ||
|
|
||
| let plan = planner | ||
| .create_initial_plan(&logical_plan, &make_session_state()) | ||
| .await?; | ||
|
|
||
| assert_contains!(format!("{plan:?}"), "ScalarSubqueryExec"); | ||
| Ok(()) | ||
| } |
There was a problem hiding this comment.
If you run this test without the other changes, you will get a failure.
| fn create_physical_expr( | ||
| &self, | ||
| expr: &Expr, | ||
| input_dfschema: &DFSchema, | ||
| _session_state: &SessionState, | ||
| ) -> Result<Arc<dyn PhysicalExpr>> { | ||
| create_physical_expr(expr, input_dfschema, self.execution_props) | ||
| } |
There was a problem hiding this comment.
This is the key to this PR. By intercepting the calls to create_physical_expr and making sure they get the execution props we support the scalar subqueries instead of relying on getting them from the session, which in this PR no longer has the correct props.
There was a problem hiding this comment.
🤔 I didn't realize we'd also need this. This starts to feel a bit more hacky...
timsaucer
changed the title
| maybe_plan = extension_planner | ||
| .plan_extension( | ||
| self, | ||
| &planner, | ||
| node.as_ref(), | ||
| &logical_input, | ||
| &children, | ||
| session_state, |
There was a problem hiding this comment.
This might actually be brittle... we are passing contradicting information to this function:
- The
ExecutionPropspresent in theplanner - The
ExecutionPropspresent in thesession_state
Both contain different info (one has uncorrelated subquery information, and the other doesn't).
This .plan_extension method will call people's code, and I imagine this has chances of blowing up in weird ways in their faces.
There was a problem hiding this comment.
I'm having some second thoughts about this approach... and realizing that it might actually be better to go with your original PR, even if it's a bigger change, it should be more future-proof (#22340)
There was a problem hiding this comment.
Alternatively, how hard would it be to have something like:
trait Session {
fn with_execution_props(&self, props: ExecutionProps) -> Self;
}And implementations are responsible for cloning as necessary?
There was a problem hiding this comment.
Or do we make pub trait Session: Send + Sync + Clone and just punt on all this work?
There was a problem hiding this comment.
If you do that, you'd still need to add some methods to trait Session that allow &mut access to the ExecutionProps no?
There was a problem hiding this comment.
It's worse than that. Clone makes it no longer dyn compatible, as would making a method that returns -> Self.
One agent was suggesting adding a wrapper
struct SessionWithProps<'a> {
inner: &'a dyn Session,
execution_props: ExecutionProps,
}
and then passing that everywhere, but this seems like a terrible approach to me. You'd implement Session on it and delegate almost all of the methods to inner except the execution_props. That would be the way to keep a minimal code change, but honestly it feels like it's adding a ton of technical debt.
I think our options are:
- This PR merges into yours and then yours into main. It's a small-ish change and targeted with minimal/no churn on other projects. The big down side is it does create some brittleness to the code.
- Add in the wrapper type suggested above, send it around everywhere, but again that feels like it's just adding technical debt.
- Go back to my original, large PR. It does need a correction for this user defined expression issue, but I can work on that.
There was a problem hiding this comment.
damn... this is hard indeed.
There was a problem hiding this comment.
I'm leaning towards just going back to your original PR. I'll take a closer look with a fresh set of eyes, but it does seem like the most robust thing to do now
|
Closing in favor of #23649 |
2 participants
This is an update to #23418 to include a failing test for user defined expressions and also a correction for the test.