fix(substrait): roundtrip scalar subqueries in projections

datafusion/substrait/src/logical_plan/producer/rel/join.rs

@@ -18,11 +18,11 @@
 use crate::logical_plan::producer::SubstraitProducer;
 use datafusion::common::{JoinConstraint, JoinType, NullEquality, not_impl_err};
 use datafusion::logical_expr::utils::conjunction;
-use datafusion::logical_expr::{Expr, Join, Operator};
+use datafusion::logical_expr::{Expr, Join, Operator, lit};
 use datafusion::prelude::binary_expr;
 use std::sync::Arc;
 use substrait::proto::rel::RelType;
-use substrait::proto::{JoinRel, Rel, join_rel};
+use substrait::proto::{CrossRel, JoinRel, Rel, join_rel};
 
 pub fn from_join(
     producer: &mut impl SubstraitProducer,
@@ -38,24 +38,34 @@ pub fn from_join(
     let right = producer.handle_plan(join.right.as_ref())?;
     let join_type = to_substrait_jointype(join.join_type);
 
-    let join_expr =
-        to_substrait_join_expr(join.on.clone(), join.null_equality, join.filter.clone());
-    let join_expression = match join_expr {
-        Some(expr) => {
-            let in_join_schema = Arc::new(join.left.schema().join(join.right.schema())?);
-            let expression = producer.handle_expr(&expr, &in_join_schema)?;
-            Some(Box::new(expression))
+    let join_expr = match to_substrait_join_expr(
+        join.on.clone(),
+        join.null_equality,
+        join.filter.clone(),
+    ) {
+        Some(expr) => expr,
+        None if join.join_type == JoinType::Inner => {
+            return Ok(Box::new(Rel {
+                rel_type: Some(RelType::Cross(Box::new(CrossRel {
+                    common: None,
+                    left: Some(left),
+                    right: Some(right),
+                    advanced_extension: None,
+                }))),
+            }));
         }
-        None => None,
+        None => lit(true),
     };
+    let in_join_schema = Arc::new(join.left.schema().join(join.right.schema())?);
+    let join_expression = producer.handle_expr(&join_expr, &in_join_schema)?;
 
     Ok(Box::new(Rel {
         rel_type: Some(RelType::Join(Box::new(JoinRel {
             common: None,
             left: Some(left),
             right: Some(right),
             r#type: join_type as i32,
-            expression: join_expression,
+            expression: Some(Box::new(join_expression)),
             post_join_filter: None,
             advanced_extension: None,
         }))),

datafusion/substrait/tests/cases/roundtrip_logical_plan.rs

@@ -731,6 +731,50 @@ async fn roundtrip_scalar_subquery_substrait() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn roundtrip_scalar_subquery_projection_with_join_rewrite() -> Result<()> {
+    let ctx = create_context().await?;
+    ctx.sql(
+        "SET datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery = false",
+    )
+    .await?
+    .collect()
+    .await?;
+
+    let plan = generate_plan_from_sql_with_ctx(
+        "SELECT a, (SELECT MAX(b) FROM data2) AS max_b FROM data",
+        false,
+        true,
+        &ctx,
+    )
+    .await?;
+
+    assert_snapshot!(
+        plan,
+        @r"
+    Projection: data.a, max(data2.b) AS max_b
+      Left Join:
+        TableScan: data projection=[a]
+        Aggregate: groupBy=[[]], aggr=[[max(data2.b)]]
+          TableScan: data2 projection=[b], partial_filters=[Boolean(true)]
+    "
+    );
+
+    let results = DataFrame::new(ctx.state(), plan).collect().await?;
+    datafusion::assert_batches_sorted_eq!(
+        [
+            "+---+-------+",
+            "| a | max_b |",
+            "+---+-------+",
+            "| 1 | 4.5   |",
+            "| 3 | 4.5   |",
+            "+---+-------+",
+        ],
+        &results
+    );
+    Ok(())
+}
+
 #[tokio::test]
 async fn roundtrip_exists_substrait() -> Result<()> {
     let ctx = create_context().await?;