feat(substrait): fork consumers for concurrent relation inputs

datafusion/substrait/Cargo.toml

@@ -38,6 +38,7 @@ async-recursion = "1.2"
 async-trait = { workspace = true }
 chrono = { workspace = true }
 datafusion = { workspace = true, features = ["sql"] }
+futures = { workspace = true }
 half = { workspace = true }
 itertools = { workspace = true }
 object_store = { workspace = true }

datafusion/substrait/src/logical_plan/consumer/rel/cross_rel.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use crate::logical_plan::consumer::SubstraitConsumer;
+use datafusion::common::{internal_datafusion_err, substrait_err};
 use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder};
 
 use datafusion::logical_expr::requalify_sides_if_needed;
@@ -26,12 +27,21 @@ pub async fn from_cross_rel(
     consumer: &impl SubstraitConsumer,
     cross: &CrossRel,
 ) -> datafusion::common::Result<LogicalPlan> {
-    let left = LogicalPlanBuilder::from(
-        consumer.consume_rel(cross.left.as_ref().unwrap()).await?,
-    );
-    let right = LogicalPlanBuilder::from(
-        consumer.consume_rel(cross.right.as_ref().unwrap()).await?,
-    );
+    let (Some(left), Some(right)) = (cross.left.as_deref(), cross.right.as_deref())
+    else {
+        return substrait_err!("CrossRel requires left and right inputs");
+    };
+    let [left, right] = consumer
+        .consume_rel_inputs(&[left, right])
+        .await?
+        .try_into()
+        .map_err(|_| {
+            internal_datafusion_err!(
+                "CrossRel consumer returned the wrong number of inputs"
+            )
+        })?;
+    let left = LogicalPlanBuilder::from(left);
+    let right = LogicalPlanBuilder::from(right);
     let (left, right, _requalified) = requalify_sides_if_needed(left, right)?;
     left.cross_join(right.build()?)?.build()
 }

datafusion/substrait/src/logical_plan/consumer/rel/join_rel.rs

@@ -16,7 +16,10 @@
 // under the License.
 
 use crate::logical_plan::consumer::SubstraitConsumer;
-use datafusion::common::{Column, JoinType, NullEquality, not_impl_err, plan_err};
+use datafusion::common::{
+    Column, JoinType, NullEquality, internal_datafusion_err, not_impl_err, plan_err,
+    substrait_err,
+};
 use datafusion::logical_expr::requalify_sides_if_needed;
 use datafusion::logical_expr::utils::split_conjunction_owned;
 use datafusion::logical_expr::{
@@ -33,12 +36,20 @@ pub async fn from_join_rel(
         return not_impl_err!("JoinRel with post_join_filter is not yet supported");
     }
 
-    let left: LogicalPlanBuilder = LogicalPlanBuilder::from(
-        consumer.consume_rel(join.left.as_ref().unwrap()).await?,
-    );
-    let right = LogicalPlanBuilder::from(
-        consumer.consume_rel(join.right.as_ref().unwrap()).await?,
-    );
+    let (Some(left), Some(right)) = (join.left.as_deref(), join.right.as_deref()) else {
+        return substrait_err!("JoinRel requires left and right inputs");
+    };
+    let [left, right] = consumer
+        .consume_rel_inputs(&[left, right])
+        .await?
+        .try_into()
+        .map_err(|_| {
+            internal_datafusion_err!(
+                "JoinRel consumer returned the wrong number of inputs"
+            )
+        })?;
+    let left = LogicalPlanBuilder::from(left);
+    let right = LogicalPlanBuilder::from(right);
     let (left, right, _requalified) = requalify_sides_if_needed(left, right)?;
 
     let join_type = from_substrait_jointype(join.r#type)?;

datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs

@@ -17,73 +17,73 @@
 
 use crate::logical_plan::consumer::SubstraitConsumer;
 use datafusion::arrow::datatypes::Field;
-use datafusion::common::{JoinType, NullEquality, not_impl_err, substrait_err};
+use datafusion::common::{
+    JoinType, NullEquality, internal_err, not_impl_err, substrait_err,
+};
 use datafusion::logical_expr::{
     Cast, Expr, LogicalPlan, LogicalPlanBuilder, requalify_sides_if_needed,
 };
 use std::sync::Arc;
+use substrait::proto::SetRel;
 use substrait::proto::set_rel::SetOp;
-use substrait::proto::{Rel, SetRel};
 
 pub async fn from_set_rel(
     consumer: &impl SubstraitConsumer,
     set: &SetRel,
 ) -> datafusion::common::Result<LogicalPlan> {
     if set.inputs.len() < 2 {
-        substrait_err!("Set operation requires at least two inputs")
-    } else {
-        match set.op() {
-            SetOp::UnionAll => union_rels(consumer, &set.inputs, true).await,
-            SetOp::UnionDistinct => union_rels(consumer, &set.inputs, false).await,
-            SetOp::IntersectionPrimary => intersect_rel(
-                consumer.consume_rel(&set.inputs[0]).await?,
-                union_rels(consumer, &set.inputs[1..], true).await?,
-                false,
-            ),
-            SetOp::IntersectionMultiset => {
-                intersect_rels(consumer, &set.inputs, false).await
-            }
-            SetOp::IntersectionMultisetAll => {
-                intersect_rels(consumer, &set.inputs, true).await
-            }
-            SetOp::MinusPrimary => except_rels(consumer, &set.inputs, false).await,
-            SetOp::MinusPrimaryAll => except_rels(consumer, &set.inputs, true).await,
-            set_op => not_impl_err!("Unsupported set operator: {set_op:?}"),
+        return substrait_err!("Set operation requires at least two inputs");
+    }
+    let op = set.op();
+    if let SetOp::Unspecified | SetOp::MinusMultiset = op {
+        return not_impl_err!("Unsupported set operator: {op:?}");
+    }
+    let inputs = set.inputs.iter().collect::<Vec<_>>();
+    let plans = consumer.consume_rel_inputs(&inputs).await?;
+    if plans.len() != inputs.len() {
+        return internal_err!("SetRel consumer returned the wrong number of inputs");
+    }
+    let mut plans = plans.into_iter();
+    let first = plans.next().expect("validated at least two inputs");
+    match op {
+        SetOp::UnionAll => union_rels(first, plans, true),
+        SetOp::UnionDistinct => union_rels(first, plans, false),
+        SetOp::IntersectionPrimary => {
+            let second = plans.next().expect("validated at least two inputs");
+            intersect_rel(first, union_rels(second, plans, true)?, false)
         }
+        SetOp::IntersectionMultiset => intersect_rels(first, plans, false),
+        SetOp::IntersectionMultisetAll => intersect_rels(first, plans, true),
+        SetOp::MinusPrimary => except_rels(first, plans, false),
+        SetOp::MinusPrimaryAll => except_rels(first, plans, true),
+        set_op => not_impl_err!("Unsupported set operator: {set_op:?}"),
     }
 }
 
-async fn union_rels(
-    consumer: &impl SubstraitConsumer,
-    rels: &[Rel],
+fn union_rels(
+    first: LogicalPlan,
+    plans: impl Iterator<Item = LogicalPlan>,
     is_all: bool,
 ) -> datafusion::common::Result<LogicalPlan> {
-    let mut union_builder = Ok(LogicalPlanBuilder::from(
-        consumer.consume_rel(&rels[0]).await?,
-    ));
-    for input in &rels[1..] {
-        let rel_plan = consumer.consume_rel(input).await?;
-
-        union_builder = if is_all {
-            union_builder?.union(rel_plan)
+    let mut builder = LogicalPlanBuilder::from(first);
+    for plan in plans {
+        builder = if is_all {
+            builder.union(plan)?
         } else {
-            union_builder?.union_distinct(rel_plan)
+            builder.union_distinct(plan)?
         };
     }
-    union_builder?.build()
+    builder.build()
 }
 
-async fn intersect_rels(
-    consumer: &impl SubstraitConsumer,
-    rels: &[Rel],
+fn intersect_rels(
+    mut rel: LogicalPlan,
+    inputs: impl Iterator<Item = LogicalPlan>,
     is_all: bool,
 ) -> datafusion::common::Result<LogicalPlan> {
-    let mut rel = consumer.consume_rel(&rels[0]).await?;
-
-    for input in &rels[1..] {
-        rel = intersect_rel(rel, consumer.consume_rel(input).await?, is_all)?;
+    for input in inputs {
+        rel = intersect_rel(rel, input, is_all)?;
     }
-
     Ok(rel)
 }
 
@@ -212,17 +212,13 @@ fn intersect_rel(
     LogicalPlanBuilder::from(joined).project(exprs)?.build()
 }
 
-async fn except_rels(
-    consumer: &impl SubstraitConsumer,
-    rels: &[Rel],
+fn except_rels(
+    mut rel: LogicalPlan,
+    inputs: impl Iterator<Item = LogicalPlan>,
     is_all: bool,
 ) -> datafusion::common::Result<LogicalPlan> {
-    let mut rel = consumer.consume_rel(&rels[0]).await?;
-
-    for input in &rels[1..] {
-        rel =
-            LogicalPlanBuilder::except(rel, consumer.consume_rel(input).await?, is_all)?;
+    for input in inputs {
+        rel = LogicalPlanBuilder::except(rel, input, is_all)?;
     }
-
     Ok(rel)
 }