phillipleblanc commented on code in PR #1912:
URL:
https://github.com/apache/datafusion-ballista/pull/1912#discussion_r3717551501
##########
ballista/core/src/planner.rs:
##########
@@ -151,10 +156,20 @@ impl<T: 'static + AsLogicalPlan> QueryPlanner for
BallistaQueryPlanner<T> {
_ => {
log::debug!("create_physical_plan - handling general
statement");
+ // Execute any uncorrelated scalar subqueries first and
+ // substitute their values, so the plan sent to the
scheduler
+ // is subquery-free. See #1910.
+ let logical_plan = self
Review Comment:
I believe this will eagerly execute for `EXPLAIN` plans which shouldn't
perform any execution, not a huge deal but something worth calling out.
##########
ballista/core/src/planner.rs:
##########
@@ -164,6 +179,132 @@ impl<T: 'static + AsLogicalPlan> QueryPlanner for
BallistaQueryPlanner<T> {
}
}
+impl<T: 'static + AsLogicalPlan> BallistaQueryPlanner<T> {
+ /// Execute every uncorrelated scalar subquery in `plan` and substitute its
+ /// single value as a literal, returning a subquery-free plan.
+ ///
+ /// DataFusion 54 plans an uncorrelated scalar subquery as a physical
+ /// `ScalarSubqueryExec` whose `ScalarSubqueryExpr` cannot be deserialized
+ /// once Ballista splits the plan into stages. Rather than decorrelating
to a
+ /// join (correct but adds join work for what is logically a constant), the
+ /// subquery is run first and its value is inlined. See #1910.
+ fn materialize_scalar_subqueries<'a>(
+ &'a self,
+ plan: LogicalPlan,
+ session_state: &'a SessionState,
+ ) -> BoxFuture<'a, Result<LogicalPlan, DataFusionError>> {
+ Box::pin(async move {
+ let subqueries = collect_uncorrelated_scalar_subqueries(&plan)?;
+ if subqueries.is_empty() {
+ return Ok(plan);
+ }
+
+ // Map each subquery (keyed by the identity of its plan Arc) to its
+ // computed value.
+ let mut values: HashMap<usize, ScalarValue> = HashMap::new();
+ for subquery in subqueries {
+ let key = Arc::as_ptr(&subquery.subquery) as usize;
+ if values.contains_key(&key) {
+ continue;
+ }
+ // A subquery may itself contain nested scalar subqueries;
inline
+ // those before running it.
+ let inner = self
+ .materialize_scalar_subqueries(
+ subquery.subquery.as_ref().clone(),
+ session_state,
+ )
+ .await?;
+ let value = self.execute_scalar_subquery(inner,
session_state).await?;
+ values.insert(key, value);
+ }
+
+ substitute_scalar_subqueries(plan, &values)
+ })
+ }
+
+ /// Run a subquery plan as its own distributed query and reduce its result
to
+ /// a single [`ScalarValue`]: 0 rows -> a typed null, 1 row -> the value,
+ /// more than one row -> an error.
+ async fn execute_scalar_subquery(
+ &self,
+ plan: LogicalPlan,
+ session_state: &SessionState,
+ ) -> Result<ScalarValue, DataFusionError> {
+ let data_type = plan.schema().field(0).data_type().clone();
+
+ let exec: Arc<dyn ExecutionPlan> =
+ Arc::new(DistributedQueryExec::<T>::with_extension(
+ self.scheduler_url.clone(),
+ self.config.clone(),
+ plan,
+ self.extension_codec.clone(),
+ session_state.session_id().to_string(),
+ ));
+
+ let batches =
+ datafusion::physical_plan::collect(exec,
session_state.task_ctx()).await?;
Review Comment:
nit: This will collect all of the results into memory before asserting it
only has a single row. It should stream the results and stop as soon as it gets
more than one row.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
