From ccbcea562b3c8112760e67ba8a65d237b2037e7c Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Fri, 23 Oct 2026 15:22:59 -0400 Subject: [PATCH 1/4] feat(substrait): add a hook for consuming relation inputs --- .../src/logical_plan/consumer/rel/join_rel.rs | 25 +- .../src/logical_plan/consumer/rel/set_rel.rs | 25 +- .../consumer/substrait_consumer.rs | 225 ++++++++++++++++++ 3 files changed, 259 insertions(+), 16 deletions(-) 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..4843b114eb510 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/join_rel.rs +++ b/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)?; 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..434c785426e83 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs @@ -17,7 +17,9 @@ 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, }; @@ -58,16 +60,21 @@ async fn union_rels( rels: &[Rel], 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 inputs = rels.iter().collect::>(); + let plans = consumer.consume_rel_inputs(&inputs).await?; + if plans.len() != rels.len() { + return internal_err!("SetRel consumer returned the wrong number of inputs"); + } + let mut plans = plans.into_iter(); + let Some(first) = plans.next() else { + return substrait_err!("Set operation requires at least two inputs"); + }; + let mut union_builder = Ok(LogicalPlanBuilder::from(first)); + for plan in plans { union_builder = if is_all { - union_builder?.union(rel_plan) + union_builder?.union(plan) } else { - union_builder?.union_distinct(rel_plan) + union_builder?.union_distinct(plan) }; } union_builder?.build() diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index 36ef39561e7c3..56a4f109f8093 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 JOIN or UNION, 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, @@ -897,4 +915,211 @@ mod tests { assert_eq!(result.fields()[0].name(), "a"); assert!(consumer.get_outer_schema(2).is_none()); } + + // A custom consumer exercises the trait default independently of the default + // consumer's scheduling policy. + struct InputConsumer { + extensions: Extensions, + state: SessionState, + reads: std::sync::Mutex>, + ready: tokio::sync::Semaphore, + } + + #[async_trait] + impl SubstraitConsumer for InputConsumer { + fn get_extensions(&self) -> &Extensions { + &self.extensions + } + + fn get_function_registry(&self) -> &impl FunctionRegistry { + &self.state + } + + async fn resolve_table_ref( + &self, + _: &TableReference, + ) -> datafusion::common::Result>> { + unreachable!("test consumes reads directly") + } + + async fn consume_rel( + &self, + rel: &Rel, + ) -> datafusion::common::Result { + let Some(proto::rel::RelType::Read(read)) = &rel.rel_type else { + panic!("expected a read") + }; + let Some(proto::read_rel::ReadType::NamedTable(table)) = &read.read_type + else { + panic!("expected a named table") + }; + self.reads.lock().unwrap().push(table.names[0].clone()); + self.ready.acquire().await.unwrap().forget(); + if table.names[0] == "fail" { + return substrait_err!("input failed"); + } + let provider = + datafusion::catalog::empty::EmptyTable::new(Arc::new(Schema::new(vec![ + Field::new("x", DataType::Int64, false), + ]))); + datafusion::logical_expr::LogicalPlanBuilder::scan( + table.names[0].clone(), + datafusion::datasource::provider_as_source(Arc::new(provider)), + None, + )? + .build() + } + } + + fn input_consumer() -> InputConsumer { + InputConsumer { + extensions: Extensions::default(), + state: datafusion::execution::SessionStateBuilder::new_with_default_features( + ) + .build(), + reads: std::sync::Mutex::default(), + ready: tokio::sync::Semaphore::new(0), + } + } + + fn input_read(name: &str) -> Rel { + Rel { + rel_type: Some(proto::rel::RelType::Read(Box::new(ReadRel { + read_type: Some(proto::read_rel::ReadType::NamedTable( + proto::read_rel::NamedTable { + names: vec![name.into()], + ..Default::default() + }, + )), + ..Default::default() + }))), + } + } + + #[tokio::test] + async fn input_hook_default_is_sequential_and_stops_on_error() { + use std::future::{Future, poll_fn}; + use std::task::Poll; + let consumer = input_consumer(); + let rels = [input_read("a"), input_read("fail"), input_read("unreached")]; + let inputs = rels.iter().collect::>(); + let mut pending = Box::pin(consumer.consume_rel_inputs(&inputs)); + poll_fn(|cx| { + assert!(pending.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + assert_eq!(*consumer.reads.lock().unwrap(), ["a"]); + consumer.ready.add_permits(3); + assert!( + pending + .await + .unwrap_err() + .to_string() + .contains("input failed") + ); + assert_eq!(*consumer.reads.lock().unwrap(), ["a", "fail"]); + } + + struct HookConsumer { + inner: InputConsumer, + calls: std::sync::Mutex>, + truncate: bool, + } + + #[async_trait] + impl SubstraitConsumer for HookConsumer { + fn get_extensions(&self) -> &Extensions { + &self.inner.extensions + } + fn get_function_registry(&self) -> &impl FunctionRegistry { + &self.inner.state + } + async fn resolve_table_ref( + &self, + _: &TableReference, + ) -> datafusion::common::Result>> { + unreachable!("test consumes reads directly") + } + async fn consume_rel_inputs( + &self, + rels: &[&Rel], + ) -> datafusion::common::Result> { + self.calls.lock().unwrap().push(rels.len()); + let mut plans = self.inner.consume_rel_inputs(rels).await?; + if self.truncate { + plans.pop(); + } + Ok(plans) + } + } + + #[tokio::test] + async fn join_and_union_use_the_input_hook() { + for truncate in [false, true] { + let consumer = HookConsumer { + inner: input_consumer(), + calls: std::sync::Mutex::default(), + truncate, + }; + consumer.inner.ready.add_permits(6); + let join = JoinRel { + left: Some(Box::new(input_read("a"))), + right: Some(Box::new(input_read("b"))), + r#type: proto::join_rel::JoinType::Inner as i32, + ..Default::default() + }; + let result = consumer.consume_join(&join).await; + if truncate { + assert!( + result + .unwrap_err() + .to_string() + .contains("wrong number of inputs") + ); + } else { + let LogicalPlan::Join(join) = result.unwrap() else { + panic!("expected join") + }; + assert!( + join.left + .display_indent() + .to_string() + .contains("TableScan: a") + ); + assert!( + join.right + .display_indent() + .to_string() + .contains("TableScan: b") + ); + } + for op in [ + proto::set_rel::SetOp::UnionAll, + proto::set_rel::SetOp::UnionDistinct, + ] { + let set = SetRel { + inputs: vec![input_read("a"), input_read("b")], + op: op as i32, + ..Default::default() + }; + let result = consumer.consume_set(&set).await; + if truncate { + assert!( + result + .unwrap_err() + .to_string() + .contains("wrong number of inputs") + ); + } else { + let plan = result.unwrap().display_indent().to_string(); + assert!( + plan.find("TableScan: a").unwrap() + < plan.find("TableScan: b").unwrap() + ); + } + } + assert_eq!(*consumer.calls.lock().unwrap(), [2, 2, 2]); + } + } } From 45a41e7772843542d29306ede856b858f02fe815 Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Thu, 29 Oct 2026 16:57:36 -0400 Subject: [PATCH 2/4] feat(substrait): route all independent relation inputs through the hook --- .../logical_plan/consumer/rel/cross_rel.rs | 22 +++- .../src/logical_plan/consumer/rel/set_rel.rs | 97 +++++++-------- .../consumer/substrait_consumer.rs | 116 +++++++++--------- .../tests/cases/roundtrip_logical_plan.rs | 34 ++--- 4 files changed, 137 insertions(+), 132 deletions(-) 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..1981cb23183bb 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/cross_rel.rs +++ b/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 { - 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() } 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 434c785426e83..801e1c474b0a5 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs @@ -24,73 +24,66 @@ 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 { 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::>(); + 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, is_all: bool, ) -> datafusion::common::Result { - let inputs = rels.iter().collect::>(); - let plans = consumer.consume_rel_inputs(&inputs).await?; - if plans.len() != rels.len() { - return internal_err!("SetRel consumer returned the wrong number of inputs"); - } - let mut plans = plans.into_iter(); - let Some(first) = plans.next() else { - return substrait_err!("Set operation requires at least two inputs"); - }; - let mut union_builder = Ok(LogicalPlanBuilder::from(first)); + let mut builder = LogicalPlanBuilder::from(first); for plan in plans { - union_builder = if is_all { - union_builder?.union(plan) + builder = if is_all { + builder.union(plan)? } else { - union_builder?.union_distinct(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, is_all: bool, ) -> datafusion::common::Result { - 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) } @@ -219,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, is_all: bool, ) -> datafusion::common::Result { - 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) } diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index 56a4f109f8093..bb5a96c4d4e98 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -31,7 +31,7 @@ use datafusion::arrow::datatypes::{DataType, FieldRef}; use datafusion::catalog::TableProvider; use datafusion::common::datatype::FieldExt; use datafusion::common::{ - DFSchema, ScalarValue, TableReference, not_impl_err, substrait_err, + DFSchema, ScalarValue, TableReference, internal_err, not_impl_err, substrait_err, }; use datafusion::execution::{FunctionRegistry, SessionState}; use datafusion::logical_expr::expr::LambdaVariable; @@ -215,7 +215,7 @@ pub trait SubstraitConsumer: Send + Sync + Sized { from_substrait_rel(self, rel).await } - /// Consumes independent inputs of a JOIN or UNION, returning plans in input order. + /// 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. @@ -493,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, @@ -697,10 +698,12 @@ 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?; + if inputs.len() != input_rels.len() { + return internal_err!( + "ExtensionMultiRel consumer returned the wrong number of inputs" + ); } let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?; Ok(LogicalPlan::Extension(Extension { node: plan })) @@ -846,6 +849,7 @@ mod tests { use super::*; use crate::logical_plan::consumer::utils::tests::test_consumer; use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::logical_expr::logical_plan::builder::table_scan; fn make_schema(fields: &[(&str, DataType)]) -> Arc { let arrow_fields: Vec = fields @@ -958,13 +962,9 @@ mod tests { if table.names[0] == "fail" { return substrait_err!("input failed"); } - let provider = - datafusion::catalog::empty::EmptyTable::new(Arc::new(Schema::new(vec![ - Field::new("x", DataType::Int64, false), - ]))); - datafusion::logical_expr::LogicalPlanBuilder::scan( - table.names[0].clone(), - datafusion::datasource::provider_as_source(Arc::new(provider)), + table_scan( + Some(table.names[0].as_str()), + &Schema::new(vec![Field::new("x", DataType::Int64, false)]), None, )? .build() @@ -1055,55 +1055,50 @@ mod tests { } #[tokio::test] - async fn join_and_union_use_the_input_hook() { - for truncate in [false, true] { - let consumer = HookConsumer { - inner: input_consumer(), - calls: std::sync::Mutex::default(), - truncate, - }; - consumer.inner.ready.add_permits(6); - let join = JoinRel { + async fn multi_input_relations_use_the_input_hook() { + use proto::rel::RelType; + use proto::set_rel::SetOp; + + let mut relations = vec![ + RelType::Join(Box::new(JoinRel { left: Some(Box::new(input_read("a"))), right: Some(Box::new(input_read("b"))), r#type: proto::join_rel::JoinType::Inner as i32, ..Default::default() - }; - let result = consumer.consume_join(&join).await; - if truncate { - assert!( - result - .unwrap_err() - .to_string() - .contains("wrong number of inputs") - ); - } else { - let LogicalPlan::Join(join) = result.unwrap() else { - panic!("expected join") + })), + RelType::Cross(Box::new(CrossRel { + left: Some(Box::new(input_read("a"))), + right: Some(Box::new(input_read("b"))), + ..Default::default() + })), + ]; + for op in [ + SetOp::UnionAll, + SetOp::UnionDistinct, + SetOp::IntersectionPrimary, + SetOp::IntersectionMultiset, + SetOp::IntersectionMultisetAll, + SetOp::MinusPrimary, + SetOp::MinusPrimaryAll, + ] { + relations.push(RelType::Set(SetRel { + inputs: vec![input_read("a"), input_read("b"), input_read("c")], + op: op as i32, + ..Default::default() + })); + } + for truncate in [false, true] { + for rel_type in &relations { + let consumer = HookConsumer { + inner: input_consumer(), + calls: std::sync::Mutex::default(), + truncate, }; - assert!( - join.left - .display_indent() - .to_string() - .contains("TableScan: a") - ); - assert!( - join.right - .display_indent() - .to_string() - .contains("TableScan: b") - ); - } - for op in [ - proto::set_rel::SetOp::UnionAll, - proto::set_rel::SetOp::UnionDistinct, - ] { - let set = SetRel { - inputs: vec![input_read("a"), input_read("b")], - op: op as i32, - ..Default::default() + consumer.inner.ready.add_permits(3); + let rel = Rel { + rel_type: Some(rel_type.clone()), }; - let result = consumer.consume_set(&set).await; + let result = consumer.consume_rel(&rel).await; if truncate { assert!( result @@ -1117,9 +1112,16 @@ mod tests { plan.find("TableScan: a").unwrap() < plan.find("TableScan: b").unwrap() ); + if let RelType::Set(_) = rel_type { + assert!( + plan.find("TableScan: b").unwrap() + < plan.find("TableScan: c").unwrap() + ); + } } + let expected_inputs = if let RelType::Set(_) = rel_type { 3 } else { 2 }; + assert_eq!(*consumer.calls.lock().unwrap(), [expected_inputs]); } - assert_eq!(*consumer.calls.lock().unwrap(), [2, 2, 2]); } } } diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index bcc7becf0af04..13d4590ff9830 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -2283,21 +2283,25 @@ async fn new_test_grammar() -> Result<()> { #[tokio::test] async fn extension_logical_plan() -> Result<()> { let ctx = create_context().await?; - let validation_bytes = b"MockUserDefinedLogicalPlan".to_vec(); - let ext_plan = LogicalPlan::Extension(Extension { - node: Arc::new(MockUserDefinedLogicalPlan { - validation_bytes, - inputs: vec![], - empty_schema: Arc::new(DFSchema::empty()), - }), - }); - - let proto = to_substrait_plan(&ext_plan, &ctx.state())?; - let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; - - let plan1str = format!("{ext_plan}"); - let plan2str = format!("{plan2}"); - assert_eq!(plan1str, plan2str); + let inputs = [ + ctx.table("data").await?.into_unoptimized_plan(), + ctx.table("data2").await?.into_unoptimized_plan(), + ]; + for input_count in 0..=2 { + let ext_plan = LogicalPlan::Extension(Extension { + node: Arc::new(MockUserDefinedLogicalPlan { + validation_bytes: b"MockUserDefinedLogicalPlan".to_vec(), + inputs: inputs[..input_count].to_vec(), + empty_schema: Arc::new(DFSchema::empty()), + }), + }); + let proto = to_substrait_plan(&ext_plan, &ctx.state())?; + let roundtrip = from_substrait_plan(&ctx.state(), &proto).await?; + assert_eq!( + ext_plan.display_indent().to_string(), + roundtrip.display_indent().to_string() + ); + } Ok(()) } From 6a36a7e9582d727e3a0fac5a02e84d38ba475eb0 Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Thu, 29 Oct 2026 18:08:05 -0400 Subject: [PATCH 3/4] test(substrait): rely on existing relation coverage for the input hook --- .../consumer/substrait_consumer.rs | 206 ------------------ .../tests/cases/roundtrip_logical_plan.rs | 34 ++- 2 files changed, 15 insertions(+), 225 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index bb5a96c4d4e98..6dececff014bc 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -849,7 +849,6 @@ mod tests { use super::*; use crate::logical_plan::consumer::utils::tests::test_consumer; use datafusion::arrow::datatypes::{Field, Schema}; - use datafusion::logical_expr::logical_plan::builder::table_scan; fn make_schema(fields: &[(&str, DataType)]) -> Arc { let arrow_fields: Vec = fields @@ -919,209 +918,4 @@ mod tests { assert_eq!(result.fields()[0].name(), "a"); assert!(consumer.get_outer_schema(2).is_none()); } - - // A custom consumer exercises the trait default independently of the default - // consumer's scheduling policy. - struct InputConsumer { - extensions: Extensions, - state: SessionState, - reads: std::sync::Mutex>, - ready: tokio::sync::Semaphore, - } - - #[async_trait] - impl SubstraitConsumer for InputConsumer { - fn get_extensions(&self) -> &Extensions { - &self.extensions - } - - fn get_function_registry(&self) -> &impl FunctionRegistry { - &self.state - } - - async fn resolve_table_ref( - &self, - _: &TableReference, - ) -> datafusion::common::Result>> { - unreachable!("test consumes reads directly") - } - - async fn consume_rel( - &self, - rel: &Rel, - ) -> datafusion::common::Result { - let Some(proto::rel::RelType::Read(read)) = &rel.rel_type else { - panic!("expected a read") - }; - let Some(proto::read_rel::ReadType::NamedTable(table)) = &read.read_type - else { - panic!("expected a named table") - }; - self.reads.lock().unwrap().push(table.names[0].clone()); - self.ready.acquire().await.unwrap().forget(); - if table.names[0] == "fail" { - return substrait_err!("input failed"); - } - table_scan( - Some(table.names[0].as_str()), - &Schema::new(vec![Field::new("x", DataType::Int64, false)]), - None, - )? - .build() - } - } - - fn input_consumer() -> InputConsumer { - InputConsumer { - extensions: Extensions::default(), - state: datafusion::execution::SessionStateBuilder::new_with_default_features( - ) - .build(), - reads: std::sync::Mutex::default(), - ready: tokio::sync::Semaphore::new(0), - } - } - - fn input_read(name: &str) -> Rel { - Rel { - rel_type: Some(proto::rel::RelType::Read(Box::new(ReadRel { - read_type: Some(proto::read_rel::ReadType::NamedTable( - proto::read_rel::NamedTable { - names: vec![name.into()], - ..Default::default() - }, - )), - ..Default::default() - }))), - } - } - - #[tokio::test] - async fn input_hook_default_is_sequential_and_stops_on_error() { - use std::future::{Future, poll_fn}; - use std::task::Poll; - let consumer = input_consumer(); - let rels = [input_read("a"), input_read("fail"), input_read("unreached")]; - let inputs = rels.iter().collect::>(); - let mut pending = Box::pin(consumer.consume_rel_inputs(&inputs)); - poll_fn(|cx| { - assert!(pending.as_mut().poll(cx).is_pending()); - Poll::Ready(()) - }) - .await; - assert_eq!(*consumer.reads.lock().unwrap(), ["a"]); - consumer.ready.add_permits(3); - assert!( - pending - .await - .unwrap_err() - .to_string() - .contains("input failed") - ); - assert_eq!(*consumer.reads.lock().unwrap(), ["a", "fail"]); - } - - struct HookConsumer { - inner: InputConsumer, - calls: std::sync::Mutex>, - truncate: bool, - } - - #[async_trait] - impl SubstraitConsumer for HookConsumer { - fn get_extensions(&self) -> &Extensions { - &self.inner.extensions - } - fn get_function_registry(&self) -> &impl FunctionRegistry { - &self.inner.state - } - async fn resolve_table_ref( - &self, - _: &TableReference, - ) -> datafusion::common::Result>> { - unreachable!("test consumes reads directly") - } - async fn consume_rel_inputs( - &self, - rels: &[&Rel], - ) -> datafusion::common::Result> { - self.calls.lock().unwrap().push(rels.len()); - let mut plans = self.inner.consume_rel_inputs(rels).await?; - if self.truncate { - plans.pop(); - } - Ok(plans) - } - } - - #[tokio::test] - async fn multi_input_relations_use_the_input_hook() { - use proto::rel::RelType; - use proto::set_rel::SetOp; - - let mut relations = vec![ - RelType::Join(Box::new(JoinRel { - left: Some(Box::new(input_read("a"))), - right: Some(Box::new(input_read("b"))), - r#type: proto::join_rel::JoinType::Inner as i32, - ..Default::default() - })), - RelType::Cross(Box::new(CrossRel { - left: Some(Box::new(input_read("a"))), - right: Some(Box::new(input_read("b"))), - ..Default::default() - })), - ]; - for op in [ - SetOp::UnionAll, - SetOp::UnionDistinct, - SetOp::IntersectionPrimary, - SetOp::IntersectionMultiset, - SetOp::IntersectionMultisetAll, - SetOp::MinusPrimary, - SetOp::MinusPrimaryAll, - ] { - relations.push(RelType::Set(SetRel { - inputs: vec![input_read("a"), input_read("b"), input_read("c")], - op: op as i32, - ..Default::default() - })); - } - for truncate in [false, true] { - for rel_type in &relations { - let consumer = HookConsumer { - inner: input_consumer(), - calls: std::sync::Mutex::default(), - truncate, - }; - consumer.inner.ready.add_permits(3); - let rel = Rel { - rel_type: Some(rel_type.clone()), - }; - let result = consumer.consume_rel(&rel).await; - if truncate { - assert!( - result - .unwrap_err() - .to_string() - .contains("wrong number of inputs") - ); - } else { - let plan = result.unwrap().display_indent().to_string(); - assert!( - plan.find("TableScan: a").unwrap() - < plan.find("TableScan: b").unwrap() - ); - if let RelType::Set(_) = rel_type { - assert!( - plan.find("TableScan: b").unwrap() - < plan.find("TableScan: c").unwrap() - ); - } - } - let expected_inputs = if let RelType::Set(_) = rel_type { 3 } else { 2 }; - assert_eq!(*consumer.calls.lock().unwrap(), [expected_inputs]); - } - } - } } diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 13d4590ff9830..bcc7becf0af04 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -2283,25 +2283,21 @@ async fn new_test_grammar() -> Result<()> { #[tokio::test] async fn extension_logical_plan() -> Result<()> { let ctx = create_context().await?; - let inputs = [ - ctx.table("data").await?.into_unoptimized_plan(), - ctx.table("data2").await?.into_unoptimized_plan(), - ]; - for input_count in 0..=2 { - let ext_plan = LogicalPlan::Extension(Extension { - node: Arc::new(MockUserDefinedLogicalPlan { - validation_bytes: b"MockUserDefinedLogicalPlan".to_vec(), - inputs: inputs[..input_count].to_vec(), - empty_schema: Arc::new(DFSchema::empty()), - }), - }); - let proto = to_substrait_plan(&ext_plan, &ctx.state())?; - let roundtrip = from_substrait_plan(&ctx.state(), &proto).await?; - assert_eq!( - ext_plan.display_indent().to_string(), - roundtrip.display_indent().to_string() - ); - } + let validation_bytes = b"MockUserDefinedLogicalPlan".to_vec(); + let ext_plan = LogicalPlan::Extension(Extension { + node: Arc::new(MockUserDefinedLogicalPlan { + validation_bytes, + inputs: vec![], + empty_schema: Arc::new(DFSchema::empty()), + }), + }); + + let proto = to_substrait_plan(&ext_plan, &ctx.state())?; + let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; + + let plan1str = format!("{ext_plan}"); + let plan2str = format!("{plan2}"); + assert_eq!(plan1str, plan2str); Ok(()) } From ab9a1a78344a20c083693c7cfc5b6958ad14c01b Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Fri, 30 Oct 2026 14:06:33 -0400 Subject: [PATCH 4/4] refactor(substrait): minimize input hook call-site changes --- .../logical_plan/consumer/rel/cross_rel.rs | 16 ++-- .../src/logical_plan/consumer/rel/join_rel.rs | 18 +---- .../src/logical_plan/consumer/rel/set_rel.rs | 77 +++++++++---------- .../consumer/substrait_consumer.rs | 7 +- 4 files changed, 46 insertions(+), 72 deletions(-) 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 1981cb23183bb..9f70cc5e4e63d 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/cross_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/cross_rel.rs @@ -16,7 +16,6 @@ // 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; @@ -27,19 +26,14 @@ pub async fn from_cross_rel( consumer: &impl SubstraitConsumer, cross: &CrossRel, ) -> datafusion::common::Result { - 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]) + .consume_rel_inputs(&[ + cross.left.as_ref().unwrap(), + cross.right.as_ref().unwrap(), + ]) .await? .try_into() - .map_err(|_| { - internal_datafusion_err!( - "CrossRel consumer returned the wrong number of inputs" - ) - })?; + .unwrap(); let left = LogicalPlanBuilder::from(left); let right = LogicalPlanBuilder::from(right); let (left, right, _requalified) = requalify_sides_if_needed(left, right)?; 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 4843b114eb510..9f3b8f430ed30 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/join_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/join_rel.rs @@ -16,10 +16,7 @@ // under the License. use crate::logical_plan::consumer::SubstraitConsumer; -use datafusion::common::{ - Column, JoinType, NullEquality, internal_datafusion_err, not_impl_err, plan_err, - substrait_err, -}; +use datafusion::common::{Column, JoinType, NullEquality, not_impl_err, plan_err}; use datafusion::logical_expr::requalify_sides_if_needed; use datafusion::logical_expr::utils::split_conjunction_owned; use datafusion::logical_expr::{ @@ -36,19 +33,12 @@ pub async fn from_join_rel( return not_impl_err!("JoinRel with post_join_filter is not yet supported"); } - 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]) + .consume_rel_inputs(&[join.left.as_ref().unwrap(), join.right.as_ref().unwrap()]) .await? .try_into() - .map_err(|_| { - internal_datafusion_err!( - "JoinRel consumer returned the wrong number of inputs" - ) - })?; - let left = LogicalPlanBuilder::from(left); + .unwrap(); + let left: LogicalPlanBuilder = LogicalPlanBuilder::from(left); let right = LogicalPlanBuilder::from(right); let (left, right, _requalified) = requalify_sides_if_needed(left, right)?; 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 801e1c474b0a5..7633295f0d394 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs @@ -17,9 +17,7 @@ use crate::logical_plan::consumer::SubstraitConsumer; use datafusion::arrow::datatypes::Field; -use datafusion::common::{ - JoinType, NullEquality, internal_err, not_impl_err, substrait_err, -}; +use datafusion::common::{JoinType, NullEquality, not_impl_err, substrait_err}; use datafusion::logical_expr::{ Cast, Expr, LogicalPlan, LogicalPlanBuilder, requalify_sides_if_needed, }; @@ -32,58 +30,53 @@ pub async fn from_set_rel( set: &SetRel, ) -> datafusion::common::Result { if set.inputs.len() < 2 { - 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::>(); - 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) + 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(rels, true), + SetOp::UnionDistinct => union_rels(rels, false), + SetOp::IntersectionPrimary => { + intersect_rel(rels.next().unwrap(), union_rels(rels, true)?, false) + } + 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:?}"), } - 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:?}"), } } fn union_rels( - first: LogicalPlan, - plans: impl Iterator, + mut rels: impl Iterator, is_all: bool, ) -> datafusion::common::Result { - let mut builder = LogicalPlanBuilder::from(first); - for plan in plans { - builder = if is_all { - builder.union(plan)? + 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 { - builder.union_distinct(plan)? + union_builder?.union_distinct(rel_plan) }; } - builder.build() + union_builder?.build() } fn intersect_rels( - mut rel: LogicalPlan, - inputs: impl Iterator, + mut rels: impl Iterator, is_all: bool, ) -> datafusion::common::Result { - for input in inputs { + let mut rel = rels.next().unwrap(); + + for input in rels { rel = intersect_rel(rel, input, is_all)?; } + Ok(rel) } @@ -213,12 +206,14 @@ fn intersect_rel( } fn except_rels( - mut rel: LogicalPlan, - inputs: impl Iterator, + mut rels: impl Iterator, is_all: bool, ) -> datafusion::common::Result { - for input in inputs { + let mut rel = rels.next().unwrap(); + + 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 6dececff014bc..a0b1fc7a5b105 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -31,7 +31,7 @@ use datafusion::arrow::datatypes::{DataType, FieldRef}; use datafusion::catalog::TableProvider; use datafusion::common::datatype::FieldExt; use datafusion::common::{ - DFSchema, ScalarValue, TableReference, internal_err, not_impl_err, substrait_err, + DFSchema, ScalarValue, TableReference, not_impl_err, substrait_err, }; use datafusion::execution::{FunctionRegistry, SessionState}; use datafusion::logical_expr::expr::LambdaVariable; @@ -700,11 +700,6 @@ impl SubstraitConsumer for DefaultSubstraitConsumer<'_> { .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?; let input_rels = rel.inputs.iter().collect::>(); let inputs = self.consume_rel_inputs(&input_rels).await?; - if inputs.len() != input_rels.len() { - return internal_err!( - "ExtensionMultiRel consumer returned the wrong number of inputs" - ); - } let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?; Ok(LogicalPlan::Extension(Extension { node: plan })) }