feat(substrait): add a hook for consuming relation inputs

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

@@ -26,12 +26,16 @@ 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 [left, right] = consumer
+        .consume_rel_inputs(&[
+            cross.left.as_ref().unwrap(),
+            cross.right.as_ref().unwrap(),
+        ])
+        .await?
+        .try_into()
+        .unwrap();
+    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

@@ -33,12 +33,13 @@ 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 [left, right] = consumer
+        .consume_rel_inputs(&[join.left.as_ref().unwrap(), join.right.as_ref().unwrap()])
+        .await?
+        .try_into()
+        .unwrap();
+    let left: LogicalPlanBuilder = 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

@@ -22,8 +22,8 @@ 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,
@@ -32,38 +32,32 @@ pub async fn from_set_rel(
     if set.inputs.len() < 2 {
         substrait_err!("Set operation requires at least two inputs")
     } else {
+        if let SetOp::Unspecified | SetOp::MinusMultiset = set.op() {
+            return not_impl_err!("Unsupported set operator: {:?}", set.op());
+        }
+        let inputs = set.inputs.iter().collect::<Vec<_>>();
+        let mut rels = consumer.consume_rel_inputs(&inputs).await?.into_iter();
         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::UnionAll => union_rels(rels, true),
+            SetOp::UnionDistinct => union_rels(rels, false),
+            SetOp::IntersectionPrimary => {
+                intersect_rel(rels.next().unwrap(), union_rels(rels, true)?, false)
             }
-            SetOp::MinusPrimary => except_rels(consumer, &set.inputs, false).await,
-            SetOp::MinusPrimaryAll => except_rels(consumer, &set.inputs, true).await,
+            SetOp::IntersectionMultiset => intersect_rels(rels, false),
+            SetOp::IntersectionMultisetAll => intersect_rels(rels, true),
+            SetOp::MinusPrimary => except_rels(rels, false),
+            SetOp::MinusPrimaryAll => except_rels(rels, true),
             set_op => not_impl_err!("Unsupported set operator: {set_op:?}"),
         }
     }
 }
 
-async fn union_rels(
-    consumer: &impl SubstraitConsumer,
-    rels: &[Rel],
+fn union_rels(
+    mut rels: 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?;
-
+    let mut union_builder = Ok(LogicalPlanBuilder::from(rels.next().unwrap()));
+    for rel_plan in rels {
         union_builder = if is_all {
             union_builder?.union(rel_plan)
         } else {
@@ -73,15 +67,14 @@ async fn union_rels(
     union_builder?.build()
 }
 
-async fn intersect_rels(
-    consumer: &impl SubstraitConsumer,
-    rels: &[Rel],
+fn intersect_rels(
+    mut rels: impl Iterator<Item = LogicalPlan>,
     is_all: bool,
 ) -> datafusion::common::Result<LogicalPlan> {
-    let mut rel = consumer.consume_rel(&rels[0]).await?;
+    let mut rel = rels.next().unwrap();
 
-    for input in &rels[1..] {
-        rel = intersect_rel(rel, consumer.consume_rel(input).await?, is_all)?;
+    for input in rels {
+        rel = intersect_rel(rel, input, is_all)?;
     }
 
     Ok(rel)
@@ -212,16 +205,14 @@ fn intersect_rel(
     LogicalPlanBuilder::from(joined).project(exprs)?.build()
 }
 
-async fn except_rels(
-    consumer: &impl SubstraitConsumer,
-    rels: &[Rel],
+fn except_rels(
+    mut rels: impl Iterator<Item = LogicalPlan>,
     is_all: bool,
 ) -> datafusion::common::Result<LogicalPlan> {
-    let mut rel = consumer.consume_rel(&rels[0]).await?;
+    let mut rel = rels.next().unwrap();
 
-    for input in &rels[1..] {
-        rel =
-            LogicalPlanBuilder::except(rel, consumer.consume_rel(input).await?, is_all)?;
+    for input in rels {
+        rel = LogicalPlanBuilder::except(rel, input, is_all)?;
     }
 
     Ok(rel)

datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs

@@ -215,6 +215,24 @@ pub trait SubstraitConsumer: Send + Sync + Sized {
         from_substrait_rel(self, rel).await
     }
 
+    /// Consumes independent inputs of a relation, returning plans in input order.
+    ///
+    /// Implementations must return exactly one plan per input, in the same order.
+    /// On error, implementations must drop any unfinished input futures.
+    ///
+    /// The default preserves sequential consumption. Implementations may override it
+    /// when they can isolate mutable traversal state between inputs.
+    async fn consume_rel_inputs(
+        &self,
+        rels: &[&Rel],
+    ) -> datafusion::common::Result<Vec<LogicalPlan>> {
+        let mut plans = Vec::with_capacity(rels.len());
+        for rel in rels {
+            plans.push(self.consume_rel(rel).await?);
+        }
+        Ok(plans)
+    }
+
     async fn consume_read(
         &self,
         rel: &ReadRel,
@@ -475,6 +493,7 @@ pub trait SubstraitConsumer: Send + Sync + Sized {
         substrait_err!("Missing handler for ExtensionSingleRel")
     }
 
+    /// Custom handlers can use [`Self::consume_rel_inputs`] to consume independent inputs.
     async fn consume_extension_multi(
         &self,
         rel: &ExtensionMultiRel,
@@ -679,11 +698,8 @@ impl SubstraitConsumer for DefaultSubstraitConsumer<'_> {
             .state
             .serializer_registry()
             .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?;
-        let mut inputs = Vec::with_capacity(rel.inputs.len());
-        for input in &rel.inputs {
-            let input_plan = self.consume_rel(input).await?;
-            inputs.push(input_plan);
-        }
+        let input_rels = rel.inputs.iter().collect::<Vec<_>>();
+        let inputs = self.consume_rel_inputs(&input_rels).await?;
         let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?;
         Ok(LogicalPlan::Extension(Extension { node: plan }))
     }