feat: support reference relations in the Substrait consumer

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

@@ -17,6 +17,7 @@
 
 mod expr;
 mod plan;
+mod references;
 mod rel;
 mod substrait_consumer;
 mod types;

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

@@ -18,13 +18,18 @@
 use super::utils::{make_renamed_schema, rename_expressions};
 use super::{DefaultSubstraitConsumer, SubstraitConsumer};
 use crate::extensions::Extensions;
-use datafusion::common::{not_impl_err, plan_err};
+use datafusion::common::{not_impl_err, plan_datafusion_err, plan_err};
 use datafusion::execution::SessionState;
 use datafusion::logical_expr::{Aggregate, LogicalPlan, Projection, col};
 use std::sync::Arc;
 use substrait::proto::{Plan, plan_rel};
 
-/// Convert Substrait Plan to DataFusion LogicalPlan
+/// Convert a Substrait plan to a DataFusion logical plan.
+///
+/// Supports one output root and additional relation trees referenced through
+/// `ReferenceRel`. Each reference is inlined independently; referenced trees are
+/// not materialized or cached. A single bare relation is also accepted.
+/// Referenced trees containing outer references are not supported.
 pub async fn from_substrait_plan(
     state: &SessionState,
     plan: &Plan,
@@ -35,90 +40,103 @@ pub async fn from_substrait_plan(
         return not_impl_err!("Type variation extensions are not supported");
     }
 
-    let consumer = DefaultSubstraitConsumer::new(&extensions, state);
+    let consumer = DefaultSubstraitConsumer::new_with_plan(&extensions, state, plan);
     from_substrait_plan_with_consumer(&consumer, plan).await
 }
 
-/// Convert Substrait Plan to DataFusion LogicalPlan using the given consumer
+/// Convert Substrait Plan to DataFusion LogicalPlan using the given consumer.
+///
+/// The consumer must resolve references against this plan. For the default
+/// consumer, construct it with [`DefaultSubstraitConsumer::new_with_plan`].
+/// Custom consumers can implement [`SubstraitConsumer::consume_reference`].
 pub async fn from_substrait_plan_with_consumer(
     consumer: &impl SubstraitConsumer,
     plan: &Plan,
 ) -> datafusion::common::Result<LogicalPlan> {
-    match plan.relations.len() {
-        1 => {
-            match plan.relations[0].rel_type.as_ref() {
-                Some(rt) => match rt {
-                    plan_rel::RelType::Rel(rel) => Ok(consumer.consume_rel(rel).await?),
-                    plan_rel::RelType::Root(root) => {
-                        let plan =
-                            consumer.consume_rel(root.input.as_ref().unwrap()).await?;
-                        if root.names.is_empty() {
-                            // Backwards compatibility for plans missing names
-                            return Ok(plan);
-                        }
-                        let renamed_schema =
-                            make_renamed_schema(plan.schema(), &root.names)?;
-                        if renamed_schema
-                            .has_equivalent_names_and_types(plan.schema())
-                            .is_ok()
-                        {
-                            // Nothing to do if the schema is already equivalent
-                            return Ok(plan);
-                        }
-                        match plan {
-                            // If the last node of the plan produces expressions, bake the renames into those expressions.
-                            // This isn't necessary for correctness, but helps with roundtrip tests.
-                            LogicalPlan::Projection(p) => {
-                                Ok(LogicalPlan::Projection(Projection::try_new(
-                                    rename_expressions(
-                                        p.expr,
-                                        p.input.schema(),
-                                        renamed_schema.fields(),
-                                    )?,
-                                    p.input,
-                                )?))
-                            }
-                            LogicalPlan::Aggregate(a) => {
-                                let (group_fields, expr_fields) =
-                                    renamed_schema.fields().split_at(a.group_expr.len());
-                                let new_group_exprs = rename_expressions(
-                                    a.group_expr,
-                                    a.input.schema(),
-                                    group_fields,
-                                )?;
-                                let new_aggr_exprs = rename_expressions(
-                                    a.aggr_expr,
-                                    a.input.schema(),
-                                    expr_fields,
-                                )?;
-                                Ok(LogicalPlan::Aggregate(Aggregate::try_new(
-                                    a.input,
-                                    new_group_exprs,
-                                    new_aggr_exprs,
-                                )?))
-                            }
-                            // There are probably more plans where we could bake things in, can add them later as needed.
-                            // Otherwise, add a new Project to handle the renaming.
-                            _ => Ok(LogicalPlan::Projection(Projection::try_new(
-                                rename_expressions(
-                                    plan.schema()
-                                        .columns()
-                                        .iter()
-                                        .map(|c| col(c.to_owned())),
-                                    plan.schema(),
-                                    renamed_schema.fields(),
-                                )?,
-                                Arc::new(plan),
-                            )?)),
-                        }
+    // Keep accepting a single bare relation for backwards compatibility. Multiple
+    // trees need an explicit, unique output root; other trees are reference targets.
+    let output = if plan.relations.len() == 1 {
+        &plan.relations[0]
+    } else {
+        let mut roots = plan.relations.iter().filter(|relation| {
+            matches!(relation.rel_type, Some(plan_rel::RelType::Root(_)))
+        });
+        let Some(root) = roots.next() else {
+            return plan_err!("Substrait plan must have exactly one output root");
+        };
+        if roots.next().is_some() {
+            return not_impl_err!(
+                "Substrait plans with multiple output roots are not supported"
+            );
+        }
+        root
+    };
+    match output.rel_type.as_ref() {
+        Some(rt) => match rt {
+            plan_rel::RelType::Rel(rel) => Ok(consumer.consume_rel(rel).await?),
+            plan_rel::RelType::Root(root) => {
+                let plan = consumer
+                    .consume_rel(root.input.as_ref().ok_or_else(|| {
+                        plan_datafusion_err!("Substrait root is missing its input")
+                    })?)
+                    .await?;
+                if root.names.is_empty() {
+                    // Backwards compatibility for plans missing names
+                    return Ok(plan);
+                }
+                let renamed_schema = make_renamed_schema(plan.schema(), &root.names)?;
+                if renamed_schema
+                    .has_equivalent_names_and_types(plan.schema())
+                    .is_ok()
+                {
+                    // Nothing to do if the schema is already equivalent
+                    return Ok(plan);
+                }
+                match plan {
+                    // If the last node of the plan produces expressions, bake the renames into those expressions.
+                    // This isn't necessary for correctness, but helps with roundtrip tests.
+                    LogicalPlan::Projection(p) => {
+                        Ok(LogicalPlan::Projection(Projection::try_new(
+                            rename_expressions(
+                                p.expr,
+                                p.input.schema(),
+                                renamed_schema.fields(),
+                            )?,
+                            p.input,
+                        )?))
+                    }
+                    LogicalPlan::Aggregate(a) => {
+                        let (group_fields, expr_fields) =
+                            renamed_schema.fields().split_at(a.group_expr.len());
+                        let new_group_exprs = rename_expressions(
+                            a.group_expr,
+                            a.input.schema(),
+                            group_fields,
+                        )?;
+                        let new_aggr_exprs = rename_expressions(
+                            a.aggr_expr,
+                            a.input.schema(),
+                            expr_fields,
+                        )?;
+                        Ok(LogicalPlan::Aggregate(Aggregate::try_new(
+                            a.input,
+                            new_group_exprs,
+                            new_aggr_exprs,
+                        )?))
                     }
-                },
-                None => plan_err!("Cannot parse plan relation: None"),
+                    // There are probably more plans where we could bake things in, can add them later as needed.
+                    // Otherwise, add a new Project to handle the renaming.
+                    _ => Ok(LogicalPlan::Projection(Projection::try_new(
+                        rename_expressions(
+                            plan.schema().columns().iter().map(|c| col(c.to_owned())),
+                            plan.schema(),
+                            renamed_schema.fields(),
+                        )?,
+                        Arc::new(plan),
+                    )?)),
+                }
             }
-        }
-        _ => not_impl_err!(
-            "Substrait plan with more than 1 relation trees not supported. Number of relation trees: {:?}",
-            plan.relations.len()
-        ),
+        },
+        None => plan_err!("Cannot parse plan relation: None"),
     }
 }

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

@@ -0,0 +1,92 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Resolves references against the enclosing plan and rejects cyclic references.
+
+use super::SubstraitConsumer;
+use datafusion::common::{Result, not_impl_err, substrait_datafusion_err, substrait_err};
+use datafusion::logical_expr::LogicalPlan;
+use std::collections::HashSet;
+use std::sync::Mutex;
+use substrait::proto::{Plan, ReferenceRel, plan_rel};
+
+pub(super) struct References<'a> {
+    plan: &'a Plan,
+    visiting: Mutex<HashSet<i32>>,
+}
+
+impl<'a> References<'a> {
+    pub(super) fn new(plan: &'a Plan) -> Self {
+        Self {
+            plan,
+            visiting: Mutex::new(HashSet::new()),
+        }
+    }
+
+    pub(super) async fn consume(
+        &self,
+        consumer: &impl SubstraitConsumer,
+        reference: &ReferenceRel,
+    ) -> Result<LogicalPlan> {
+        let ordinal = reference.subtree_ordinal;
+        let relation = usize::try_from(ordinal)
+            .ok()
+            .and_then(|index| self.plan.relations.get(index))
+            .ok_or_else(|| {
+                substrait_datafusion_err!("Invalid subtree ordinal {ordinal}")
+            })?;
+        let input = match relation.rel_type.as_ref() {
+            Some(plan_rel::RelType::Rel(rel)) => rel,
+            Some(plan_rel::RelType::Root(root)) => {
+                // ReferenceRel ignores the output names of a referenced root.
+                root.input.as_ref().ok_or_else(|| {
+                    substrait_datafusion_err!(
+                        "Referenced root {ordinal} is missing its input"
+                    )
+                })?
+            }
+            None => return substrait_err!("Referenced relation {ordinal} is empty"),
+        };
+        if !self.visiting.lock().unwrap().insert(ordinal) {
+            return substrait_err!("Cyclic Substrait reference at subtree {ordinal}");
+        }
+        // Keep only the active recursion path, so repeated references are valid.
+        // The guard also removes the visit when conversion errors or is cancelled.
+        let _visit = ReferenceVisit {
+            visiting: &self.visiting,
+            ordinal,
+        };
+        let body = consumer.consume_rel(input).await?;
+        if body.contains_outer_reference() {
+            return not_impl_err!(
+                "Correlated shared Substrait relations are not supported"
+            );
+        }
+        Ok(body)
+    }
+}
+
+struct ReferenceVisit<'a> {
+    visiting: &'a Mutex<HashSet<i32>>,
+    ordinal: i32,
+}
+
+impl Drop for ReferenceVisit<'_> {
+    fn drop(&mut self) {
+        self.visiting.lock().unwrap().remove(&self.ordinal);
+    }
+}

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

@@ -56,6 +56,7 @@ pub async fn from_substrait_rel(
 ) -> datafusion::common::Result<LogicalPlan> {
     let plan: datafusion::common::Result<LogicalPlan> = match &relation.rel_type {
         Some(rel_type) => match rel_type {
+            RelType::Reference(rel) => consumer.consume_reference(rel).await,
             RelType::Read(rel) => consumer.consume_read(rel).await,
             RelType::Filter(rel) => consumer.consume_filter(rel).await,
             RelType::Fetch(rel) => consumer.consume_fetch(rel).await,

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

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use super::references::References;
 use super::{
     from_aggregate_rel, from_cast, from_cross_rel, from_exchange_rel, from_fetch_rel,
     from_field_reference, from_filter_rel, from_if_then, from_join_rel, from_literal,
@@ -47,7 +48,8 @@ use substrait::proto::{self, Type};
 use substrait::proto::{
     AggregateRel, ConsistentPartitionWindowRel, CrossRel, DynamicParameter, ExchangeRel,
     Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
-    FilterRel, JoinRel, ProjectRel, ReadRel, Rel, SetRel, SortRel, r#type,
+    FilterRel, JoinRel, Plan, ProjectRel, ReadRel, ReferenceRel, Rel, SetRel, SortRel,
+    r#type,
 };
 
 #[async_trait]
@@ -215,6 +217,17 @@ pub trait SubstraitConsumer: Send + Sync + Sized {
         from_substrait_rel(self, rel).await
     }
 
+    /// Resolve a tree in the enclosing Substrait plan. Custom consumers can
+    /// override this to customize reference resolution.
+    async fn consume_reference(
+        &self,
+        _rel: &ReferenceRel,
+    ) -> datafusion::common::Result<LogicalPlan> {
+        not_impl_err!(
+            "ReferenceRel requires a consumer with access to the enclosing plan"
+        )
+    }
+
     async fn consume_read(
         &self,
         rel: &ReadRel,
@@ -578,15 +591,30 @@ pub trait SubstraitConsumer: Send + Sync + Sized {
 ///
 /// Used as the consumer in [crate::logical_plan::consumer::from_substrait_plan]
 pub struct DefaultSubstraitConsumer<'a> {
+    pub(super) references: Option<References<'a>>,
     pub(super) extensions: &'a Extensions,
     pub(super) state: &'a SessionState,
     outer_schemas: RwLock<Vec<Arc<DFSchema>>>,
     lambda_consumer: DefaultSubstraitLambdaConsumer,
 }
 
 impl<'a> DefaultSubstraitConsumer<'a> {
+    /// Create a consumer that can resolve `ReferenceRel` against `plan`.
+    /// References are inlined in the resulting logical plan. Use a separate
+    /// consumer for each concurrent conversion.
+    pub fn new_with_plan(
+        extensions: &'a Extensions,
+        state: &'a SessionState,
+        plan: &'a Plan,
+    ) -> Self {
+        let mut consumer = Self::new(extensions, state);
+        consumer.references = Some(References::new(plan));
+        consumer
+    }
+
     pub fn new(extensions: &'a Extensions, state: &'a SessionState) -> Self {
         DefaultSubstraitConsumer {
+            references: None,
             extensions,
             state,
             outer_schemas: RwLock::new(Vec::new()),
@@ -597,6 +625,18 @@ impl<'a> DefaultSubstraitConsumer<'a> {
 
 #[async_trait]
 impl SubstraitConsumer for DefaultSubstraitConsumer<'_> {
+    async fn consume_reference(
+        &self,
+        rel: &ReferenceRel,
+    ) -> datafusion::common::Result<LogicalPlan> {
+        let Some(references) = &self.references else {
+            return substrait_err!(
+                "ReferenceRel requires DefaultSubstraitConsumer::new_with_plan"
+            );
+        };
+        references.consume(self, rel).await
+    }
+
     async fn resolve_table_ref(
         &self,
         table_ref: &TableReference,