diff --git a/ballista/client/tests/common/mod.rs b/ballista/client/tests/common/mod.rs index 93a2214a20..67b63a9114 100644 --- a/ballista/client/tests/common/mod.rs +++ b/ballista/client/tests/common/mod.rs @@ -20,6 +20,7 @@ use std::error::Error; use std::path::PathBuf; use ballista::prelude::{SessionConfigExt, SessionContextExt}; +use ballista_core::config::BALLISTA_ADAPTIVE_PLANNER_ENABLED; use ballista_core::serde::{ BallistaCodec, protobuf::scheduler_grpc_client::SchedulerGrpcClient, }; @@ -257,6 +258,34 @@ pub async fn remote_context_with_state() -> SessionContext { .unwrap() } +/// Session state with the adaptive (AQE) planner enabled, which is off by +/// default. +#[allow(dead_code)] +fn adaptive_planner_state() -> SessionState { + let config = SessionConfig::new_with_ballista() + .set_bool(BALLISTA_ADAPTIVE_PLANNER_ENABLED, true); + SessionStateBuilder::new() + .with_config(config) + .with_default_features() + .build() +} + +#[allow(dead_code)] +pub async fn standalone_context_with_aqe() -> SessionContext { + SessionContext::standalone_with_state(adaptive_planner_state()) + .await + .unwrap() +} + +#[allow(dead_code)] +pub async fn remote_context_with_aqe() -> SessionContext { + let state = adaptive_planner_state(); + let (host, port) = setup_test_cluster_with_state(state.clone()).await; + SessionContext::remote_with_state(&format!("df://{host}:{port}"), state) + .await + .unwrap() +} + #[ctor::ctor(unsafe)] fn init() { // Enable RUST_LOG logging configuration for test diff --git a/ballista/client/tests/context_checks.rs b/ballista/client/tests/context_checks.rs index dc0576d9d6..474fbf07e1 100644 --- a/ballista/client/tests/context_checks.rs +++ b/ballista/client/tests/context_checks.rs @@ -20,8 +20,8 @@ mod common; mod supported { use crate::common::{ - remote_context, remote_context_with_state, standalone_context, - standalone_context_with_state, + remote_context, remote_context_with_aqe, remote_context_with_state, + standalone_context, standalone_context_with_aqe, standalone_context_with_state, }; use ballista_core::config::BallistaConfig; @@ -77,6 +77,53 @@ mod supported { Ok(()) } + // An uncorrelated scalar subquery is executed first and its value inlined + // into the outer plan (rather than decorrelated to a join), so the outer + // query distributes without a `ScalarSubqueryExec`. Runs with the adaptive + // (AQE) planner both off (default) and on. See #1910. + #[rstest] + #[case::standalone(standalone_context())] + #[case::remote(remote_context())] + #[case::standalone_aqe(standalone_context_with_aqe())] + #[case::remote_aqe(remote_context_with_aqe())] + #[tokio::test] + async fn should_execute_uncorrelated_scalar_subquery( + #[future(awt)] + #[case] + ctx: SessionContext, + test_data: String, + ) -> datafusion::error::Result<()> { + ctx.register_parquet( + "test", + &format!("{test_data}/alltypes_plain.parquet"), + Default::default(), + ) + .await?; + + // The optimizer is left free to plan a physical scalar subquery (this is + // the DataFusion default; #1909 disabled it as a workaround, and #1910 + // reverted that). Set it explicitly so the test states what it covers. + ctx.sql( + "SET datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery = true", + ) + .await? + .collect() + .await?; + + // Exactly the row holding the maximum id matches the inlined subquery + // value, so the count is 1. + let result = ctx + .sql("select count(*) as cnt from test where id = (select max(id) from test)") + .await? + .collect() + .await?; + let expected = ["+-----+", "| cnt |", "+-----+", "| 1 |", "+-----+"]; + + assert_batches_eq!(expected, &result); + + Ok(()) + } + // tests if client will collect statistics for // collect/show operation #[rstest] diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs index aaa68c6735..3c455c6c1b 100644 --- a/ballista/core/src/extension.rs +++ b/ballista/core/src/extension.rs @@ -787,22 +787,6 @@ impl SessionConfigHelperExt for SessionConfig { // See https://github.com/apache/datafusion-ballista/issues/1648 .set_bool("datafusion.optimizer.prefer_hash_join", false) // - // DataFusion 54 plans uncorrelated scalar subqueries as a physical - // `ScalarSubqueryExec` wrapping a `ScalarSubqueryExpr` that reads an - // in-process shared results container. That container cannot cross - // process or stage boundaries, and `datafusion-proto` can only - // deserialize the expr inside its surrounding exec, so when Ballista - // splits a plan into stages the expr is serialized without its exec - // and the executor fails to decode it. Disabling this option makes - // the optimizer rewrite uncorrelated scalar subqueries to joins, - // which Ballista distributes correctly. - // - // See https://github.com/apache/datafusion-ballista/issues/1909 - .set_bool( - "datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery", - false, - ) - // // DataFusion's dynamic filters are populated at runtime by an // upstream operator (a hash join build side, a TopK heap, a partial // aggregate) and read by a downstream scan within the same plan. @@ -1107,15 +1091,15 @@ mod test { ); } - // Uncorrelated scalar subqueries must be rewritten to joins rather than - // planned as a physical `ScalarSubqueryExec`, whose `ScalarSubqueryExpr` - // cannot be deserialized once Ballista splits the plan into stages. See - // #1909. + // Ballista must NOT force-disable physical uncorrelated scalar subqueries: + // the client query planner materializes them (executes the subquery first + // and substitutes the value) instead of decorrelating to a join, so the + // optimizer is left to keep them as scalar subqueries. See #1909 / #1910. #[test] - fn should_disable_physical_uncorrelated_scalar_subquery() { + fn should_keep_physical_uncorrelated_scalar_subquery_default() { let config = SessionConfig::new().upgrade_for_ballista(); assert!( - !config + config .options() .optimizer .enable_physical_uncorrelated_scalar_subquery diff --git a/ballista/core/src/planner.rs b/ballista/core/src/planner.rs index ecf41733b8..67f3b2b60a 100644 --- a/ballista/core/src/planner.rs +++ b/ballista/core/src/planner.rs @@ -20,14 +20,19 @@ use crate::execution_plans::{DistributedExplainAnalyzeExec, DistributedQueryExec use crate::serde::BallistaLogicalExtensionCodec; use datafusion::arrow::datatypes::Schema; -use datafusion::common::tree_node::{TreeNode, TreeNodeVisitor}; +use datafusion::common::ScalarValue; +use datafusion::common::tree_node::{ + Transformed, TreeNode, TreeNodeRecursion, TreeNodeVisitor, +}; use datafusion::error::DataFusionError; use datafusion::execution::context::{QueryPlanner, SessionState}; -use datafusion::logical_expr::{LogicalPlan, TableScan}; +use datafusion::logical_expr::{Expr, LogicalPlan, Subquery, TableScan}; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner}; use datafusion_proto::logical_plan::{AsLogicalPlan, LogicalExtensionCodec}; +use futures::future::BoxFuture; +use std::collections::HashMap; use std::marker::PhantomData; use std::sync::Arc; @@ -151,10 +156,20 @@ impl QueryPlanner for BallistaQueryPlanner { _ => { 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 + .materialize_scalar_subqueries( + logical_plan.clone(), + session_state, + ) + .await?; + Ok(Arc::new(DistributedQueryExec::::with_extension( self.scheduler_url.clone(), self.config.clone(), - logical_plan.clone(), + logical_plan, self.extension_codec.clone(), session_state.session_id().to_string(), ))) @@ -164,6 +179,132 @@ impl QueryPlanner for BallistaQueryPlanner { } } +impl BallistaQueryPlanner { + /// 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> { + 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 = 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 { + let data_type = plan.schema().field(0).data_type().clone(); + + let exec: Arc = + Arc::new(DistributedQueryExec::::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?; + let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + if num_rows > 1 { + return Err(DataFusionError::Execution( + "scalar subquery returned more than one row".to_string(), + )); + } + match batches.into_iter().find(|b| b.num_rows() == 1) { + Some(batch) => ScalarValue::try_from_array(batch.column(0), 0), + None => ScalarValue::try_from(&data_type), + } + } +} + +/// Collect uncorrelated scalar subqueries from a plan's expressions, without +/// descending into the subquery plans themselves (those are handled by +/// recursion in [`BallistaQueryPlanner::materialize_scalar_subqueries`]). +fn collect_uncorrelated_scalar_subqueries( + plan: &LogicalPlan, +) -> Result, DataFusionError> { + let mut found = Vec::new(); + let mut stack = vec![plan]; + while let Some(node) = stack.pop() { + for expr in node.expressions() { + expr.apply(|e| { + if let Expr::ScalarSubquery(subquery) = e + && subquery.outer_ref_columns.is_empty() + { + found.push(subquery.clone()); + // Do not descend into the subquery's own plan here. + return Ok(TreeNodeRecursion::Jump); + } + Ok(TreeNodeRecursion::Continue) + })?; + } + stack.extend(node.inputs()); + } + Ok(found) +} + +/// Replace each uncorrelated scalar subquery (matched by the identity of its +/// plan Arc) with its precomputed literal value. +fn substitute_scalar_subqueries( + plan: LogicalPlan, + values: &HashMap, +) -> Result { + plan.transform_up(|node| { + node.map_expressions(|expr| { + expr.transform_up(|e| { + if let Expr::ScalarSubquery(subquery) = &e { + let key = Arc::as_ptr(&subquery.subquery) as usize; + if let Some(value) = values.get(&key) { + return Ok(Transformed::yes(Expr::Literal(value.clone(), None))); + } + } + Ok(Transformed::no(e)) + }) + }) + }) + .map(|transformed| transformed.data) +} + /// A Visitor which detect if query is using local tables, /// such as tables located in `information_schema` and returns true /// only if all scans are in from local tables @@ -316,4 +457,106 @@ mod test { ); Ok(()) } + + fn one_row_subquery() -> std::sync::Arc { + use datafusion::logical_expr::{LogicalPlanBuilder, lit}; + std::sync::Arc::new( + LogicalPlanBuilder::empty(true) + .project(vec![lit(42i64).alias("m")]) + .unwrap() + .build() + .unwrap(), + ) + } + + #[test] + fn collect_finds_uncorrelated_scalar_subqueries() { + use datafusion::logical_expr::LogicalPlanBuilder; + use datafusion::logical_expr::expr_fn::scalar_subquery; + + let subquery = one_row_subquery(); + let plan = LogicalPlanBuilder::empty(true) + .project(vec![ + scalar_subquery(std::sync::Arc::clone(&subquery)).alias("x"), + ]) + .unwrap() + .build() + .unwrap(); + + // The uncorrelated subquery is collected, and its plan is what comes back. + let found = super::collect_uncorrelated_scalar_subqueries(&plan).unwrap(); + assert_eq!(found.len(), 1); + assert_eq!( + found[0].subquery.display_indent().to_string(), + "Projection: Int64(42) AS m\n EmptyRelation: rows=1" + ); + } + + #[test] + fn collect_skips_correlated_scalar_subqueries() { + use datafusion::logical_expr::expr_fn::scalar_subquery; + use datafusion::logical_expr::{Expr, LogicalPlanBuilder, Subquery, col, lit}; + + // An uncorrelated subquery (eligible) sits next to a correlated one that + // references an outer column and must be skipped. + let uncorrelated = one_row_subquery(); + let correlated = Expr::ScalarSubquery(Subquery { + subquery: std::sync::Arc::new( + LogicalPlanBuilder::empty(true) + .project(vec![lit(99i64).alias("m")]) + .unwrap() + .build() + .unwrap(), + ), + outer_ref_columns: vec![col("y")], + spans: Default::default(), + }); + let plan = LogicalPlanBuilder::empty(true) + .project(vec![ + scalar_subquery(std::sync::Arc::clone(&uncorrelated)).alias("a"), + correlated.alias("b"), + ]) + .unwrap() + .build() + .unwrap(); + + // Only the uncorrelated subquery is collected. The plan that comes back + // is the `Int64(42)` one, not the correlated `Int64(99)` subquery. + let found = super::collect_uncorrelated_scalar_subqueries(&plan).unwrap(); + assert_eq!(found.len(), 1); + assert_eq!( + found[0].subquery.display_indent().to_string(), + "Projection: Int64(42) AS m\n EmptyRelation: rows=1" + ); + } + + #[test] + fn substitute_replaces_scalar_subqueries_with_literals() { + use datafusion::common::ScalarValue; + use datafusion::logical_expr::LogicalPlanBuilder; + use datafusion::logical_expr::expr_fn::scalar_subquery; + use std::collections::HashMap; + + let subquery = one_row_subquery(); + let plan = LogicalPlanBuilder::empty(true) + .project(vec![ + scalar_subquery(std::sync::Arc::clone(&subquery)).alias("x"), + ]) + .unwrap() + .build() + .unwrap(); + + let mut values = HashMap::new(); + values.insert( + std::sync::Arc::as_ptr(&subquery) as usize, + ScalarValue::Int64(Some(42)), + ); + + // The scalar subquery is replaced by its literal value in the plan. + let rewritten = super::substitute_scalar_subqueries(plan, &values).unwrap(); + assert_eq!( + rewritten.display_indent().to_string(), + "Projection: Int64(42) AS x\n EmptyRelation: rows=1" + ); + } }