feat: execute uncorrelated scalar subqueries first instead of rewriting to a join

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

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]

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