diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/cross_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/cross_rel.rs index 25c66a8e22972..9f70cc5e4e63d 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/cross_rel.rs +++ b/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 { - 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() } diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/join_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/join_rel.rs index 7850dbea797fb..9f3b8f430ed30 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/join_rel.rs +++ b/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)?; diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs index 10e78a98eef5b..7633295f0d394 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs +++ b/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::>(); + 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, is_all: bool, ) -> datafusion::common::Result { - 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, is_all: bool, ) -> datafusion::common::Result { - 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, is_all: bool, ) -> datafusion::common::Result { - 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) diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index 36ef39561e7c3..a0b1fc7a5b105 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/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> { + 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::>(); + 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 })) }