feat(substrait): serialize correlated subqueries with OuterReference field references

datafusion/substrait/src/logical_plan/producer/expr/field_reference.rs

@@ -15,11 +15,13 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::logical_plan::producer::SubstraitProducer;
+use datafusion::arrow::datatypes::FieldRef;
 use datafusion::common::{Column, DFSchemaRef, substrait_err};
 use datafusion::logical_expr::Expr;
 use substrait::proto::Expression;
 use substrait::proto::expression::field_reference::{
-    ReferenceType, RootReference, RootType,
+    OuterReference, ReferenceType, RootReference, RootType,
 };
 use substrait::proto::expression::{
     FieldReference, ReferenceSegment, RexType, reference_segment,
@@ -35,6 +37,13 @@ pub fn from_column(
 
 pub(crate) fn substrait_field_ref(
     index: usize,
+) -> datafusion::common::Result<Expression> {
+    substrait_field_ref_with_root(index, RootType::RootReference(RootReference {}))
+}
+
+fn substrait_field_ref_with_root(
+    index: usize,
+    root_type: RootType,
 ) -> datafusion::common::Result<Expression> {
     Ok(Expression {
         rex_type: Some(RexType::Selection(Box::new(FieldReference {
@@ -46,7 +55,7 @@ pub(crate) fn substrait_field_ref(
                     }),
                 )),
             })),
-            root_type: Some(RootType::RootReference(RootReference {})),
+            root_type: Some(root_type),
         }))),
     })
 }
@@ -76,20 +85,37 @@ pub(crate) fn try_to_substrait_field_reference(
     }
 }
 
-/// Convert an outer reference column to a Substrait field reference.
-/// Outer reference columns reference columns from an outer query scope in correlated subqueries.
-/// We convert them the same way as regular columns since the subquery plan will be
-/// reconstructed with the proper schema context during consumption.
+/// Convert an outer reference column to a Substrait field reference with an
+/// `OuterReference` root type.
+///
+/// Outer reference columns reference columns from an enclosing query scope in
+/// correlated subqueries. The column is resolved against the producer's stack
+/// of outer schemas (pushed at each subquery boundary), innermost first, and
+/// the resulting `steps_out` records how many query boundaries the reference
+/// crosses (`steps_out = 1` is the immediately enclosing query).
 pub fn from_outer_reference_column(
+    producer: &mut impl SubstraitProducer,
+    _field: &FieldRef,
     col: &Column,
-    schema: &DFSchemaRef,
+    _schema: &DFSchemaRef,
 ) -> datafusion::common::Result<Expression> {
-    // OuterReferenceColumn is converted similarly to a regular column reference.
-    // The schema provided should be the schema context in which the outer reference
-    // column appears. During Substrait round-trip, the consumer will reconstruct
-    // the outer reference based on the subquery context.
-    let index = schema.index_of_column(col)?;
-    substrait_field_ref(index)
+    let mut steps_out = 1;
+    while let Some(outer_schema) = producer.get_outer_schema(steps_out) {
+        if let Some(index) = outer_schema.maybe_index_of_column(col) {
+            return substrait_field_ref_with_root(
+                index,
+                RootType::OuterReference(OuterReference {
+                    steps_out: steps_out as u32,
+                }),
+            );
+        }
+        steps_out += 1;
+    }
+    substrait_err!(
+        "Outer reference column '{col}' could not be resolved against any outer \
+         query schema. If using a custom SubstraitProducer, ensure it maintains \
+         the outer schema stack (push_outer_schema/pop_outer_schema/get_outer_schema)"
+    )
 }
 
 #[cfg(test)]

datafusion/substrait/src/logical_plan/producer/expr/mod.rs

@@ -149,10 +149,8 @@ pub fn to_substrait_rex(
         Expr::Wildcard { .. } => not_impl_err!("Cannot convert {expr:?} to Substrait"),
         Expr::GroupingSet(expr) => not_impl_err!("Cannot convert {expr:?} to Substrait"),
         Expr::Placeholder(expr) => producer.handle_placeholder(expr, schema),
-        Expr::OuterReferenceColumn(_, _) => {
-            // OuterReferenceColumn requires tracking outer query schema context for correlated
-            // subqueries. This is a complex feature that is not yet implemented.
-            not_impl_err!("Cannot convert {expr:?} to Substrait")
+        Expr::OuterReferenceColumn(field, col) => {
+            producer.handle_outer_reference_column(field, col, schema)
         }
         Expr::Unnest(expr) => not_impl_err!("Cannot convert {expr:?} to Substrait"),
         Expr::HigherOrderFunction(expr) => {

datafusion/substrait/src/logical_plan/producer/expr/subquery.rs

@@ -18,11 +18,29 @@
 use crate::logical_plan::producer::{SubstraitProducer, negate};
 use datafusion::common::{DFSchemaRef, substrait_err};
 use datafusion::logical_expr::expr::{Exists, InSubquery, SetComparison, SetQuantifier};
-use datafusion::logical_expr::{Operator, Subquery};
-use substrait::proto::Expression;
+use datafusion::logical_expr::{LogicalPlan, Operator, Subquery};
+use std::sync::Arc;
 use substrait::proto::expression::RexType;
 use substrait::proto::expression::subquery::set_comparison::{ComparisonOp, ReductionOp};
 use substrait::proto::expression::subquery::{InPredicate, Scalar, SetPredicate};
+use substrait::proto::{Expression, Rel};
+
+/// Serialize a subquery plan, making the enclosing query's schema available
+/// for resolving correlated column references.
+///
+/// Substrait represents correlated references using `OuterReference` field
+/// references with a `steps_out` depth. To produce these, the producer
+/// maintains a stack of outer schemas.
+fn produce_subquery_rel(
+    producer: &mut impl SubstraitProducer,
+    plan: &LogicalPlan,
+    outer_schema: &DFSchemaRef,
+) -> datafusion::common::Result<Box<Rel>> {
+    producer.push_outer_schema(Arc::clone(outer_schema));
+    let result = producer.handle_plan(plan);
+    producer.pop_outer_schema();
+    result
+}
 
 pub fn from_in_subquery(
     producer: &mut impl SubstraitProducer,
@@ -36,7 +54,8 @@ pub fn from_in_subquery(
     } = subquery;
     let substrait_expr = producer.handle_expr(expr, schema)?;
 
-    let subquery_plan = producer.handle_plan(subquery.subquery.as_ref())?;
+    let subquery_plan =
+        produce_subquery_rel(producer, subquery.subquery.as_ref(), schema)?;
 
     let substrait_subquery = Expression {
         rex_type: Some(RexType::Subquery(Box::new(
@@ -88,8 +107,11 @@ pub fn from_set_comparison(
     let comparison_op = comparison_op_to_proto(&set_comparison.op)? as i32;
     let reduction_op = reduction_op_to_proto(&set_comparison.quantifier)? as i32;
     let left = producer.handle_expr(set_comparison.expr.as_ref(), schema)?;
-    let subquery_plan =
-        producer.handle_plan(set_comparison.subquery.subquery.as_ref())?;
+    let subquery_plan = produce_subquery_rel(
+        producer,
+        set_comparison.subquery.subquery.as_ref(),
+        schema,
+    )?;
 
     Ok(Expression {
         rex_type: Some(RexType::Subquery(Box::new(
@@ -113,9 +135,10 @@ pub fn from_set_comparison(
 pub fn from_scalar_subquery(
     producer: &mut impl SubstraitProducer,
     subquery: &Subquery,
-    _schema: &DFSchemaRef,
+    schema: &DFSchemaRef,
 ) -> datafusion::common::Result<Expression> {
-    let subquery_plan = producer.handle_plan(subquery.subquery.as_ref())?;
+    let subquery_plan =
+        produce_subquery_rel(producer, subquery.subquery.as_ref(), schema)?;
 
     Ok(Expression {
         rex_type: Some(RexType::Subquery(Box::new(
@@ -136,9 +159,10 @@ pub fn from_scalar_subquery(
 pub fn from_exists(
     producer: &mut impl SubstraitProducer,
     exists: &Exists,
-    _schema: &DFSchemaRef,
+    schema: &DFSchemaRef,
 ) -> datafusion::common::Result<Expression> {
-    let subquery_plan = producer.handle_plan(exists.subquery.subquery.as_ref())?;
+    let subquery_plan =
+        produce_subquery_rel(producer, exists.subquery.subquery.as_ref(), schema)?;
 
     let substrait_exists = Expression {
         rex_type: Some(RexType::Subquery(Box::new(

datafusion/substrait/src/logical_plan/producer/substrait_producer.rs

@@ -21,10 +21,10 @@ use crate::logical_plan::producer::{
     from_case, from_cast, from_column, from_distinct, from_empty_relation, from_exists,
     from_filter, from_higher_order_function, from_in_list, from_in_subquery, from_join,
     from_lambda, from_lambda_variable, from_like, from_limit, from_literal,
-    from_placeholder, from_projection, from_repartition, from_scalar_function,
-    from_scalar_subquery, from_set_comparison, from_sort, from_subquery_alias,
-    from_table_scan, from_try_cast, from_unary_expr, from_union, from_values,
-    from_window, from_window_function, to_substrait_rel, to_substrait_rex,
+    from_outer_reference_column, from_placeholder, from_projection, from_repartition,
+    from_scalar_function, from_scalar_subquery, from_set_comparison, from_sort,
+    from_subquery_alias, from_table_scan, from_try_cast, from_unary_expr, from_union,
+    from_values, from_window, from_window_function, to_substrait_rel, to_substrait_rex,
     to_substrait_type_from_field,
 };
 use datafusion::arrow::datatypes::FieldRef;
@@ -433,6 +433,15 @@ pub trait SubstraitProducer: Send + Sync + Sized {
         from_exists(self, exists, schema)
     }
 
+    fn handle_outer_reference_column(
+        &mut self,
+        field: &FieldRef,
+        column: &Column,
+        schema: &DFSchemaRef,
+    ) -> datafusion::common::Result<Expression> {
+        from_outer_reference_column(self, field, column, schema)
+    }
+
     fn handle_placeholder(
         &mut self,
         placeholder: &Placeholder,
@@ -441,6 +450,30 @@ pub trait SubstraitProducer: Send + Sync + Sized {
         from_placeholder(self, placeholder)
     }
 
+    // Outer Schema management API.
+    //
+    // These methods manage a stack of outer schemas for correlated subquery support
+    // such as when entering a subquery, the enclosing query's schema is pushed onto
+    // the stack.
+    //
+    // Serializing an Expr::OuterReferenceColumn uses these to resolve the column
+    // against the correct enclosing query and emit an OuterReference field reference
+    // with the corresponding `steps_out`.
+
+    /// Push an outer schema onto the stack when entering a subquery.
+    fn push_outer_schema(&mut self, _schema: DFSchemaRef) {}
+
+    /// Pop an outer schema from the stack when leaving a subquery.
+    fn pop_outer_schema(&mut self) {}
+
+    /// Get the outer schema at the given nesting depth.
+    /// `steps_out = 1` is the immediately enclosing query, `steps_out = 2`
+    /// is two levels out, etc. Returns `None` if `steps_out` is 0 or
+    /// exceeds the current nesting depth.
+    fn get_outer_schema(&self, _steps_out: usize) -> Option<DFSchemaRef> {
+        None
+    }
+
     fn handle_lambda(
         &mut self,
         lambda: &Lambda,
@@ -499,6 +532,7 @@ pub struct DefaultSubstraitProducer<'a> {
     extensions: Extensions,
     serializer_registry: &'a dyn SerializerRegistry,
     lambda_producer: DefaultSubstraitLambdaProducer,
+    outer_schemas: Vec<DFSchemaRef>,
 }
 
 impl<'a> DefaultSubstraitProducer<'a> {
@@ -507,6 +541,7 @@ impl<'a> DefaultSubstraitProducer<'a> {
             extensions: Extensions::default(),
             serializer_registry: state.serializer_registry().as_ref(),
             lambda_producer: DefaultSubstraitLambdaProducer::new(),
+            outer_schemas: Vec::new(),
         }
     }
 }
@@ -562,6 +597,23 @@ impl SubstraitProducer for DefaultSubstraitProducer<'_> {
         }))
     }
 
+    fn push_outer_schema(&mut self, schema: DFSchemaRef) {
+        self.outer_schemas.push(schema);
+    }
+
+    fn pop_outer_schema(&mut self) {
+        self.outer_schemas.pop();
+    }
+
+    fn get_outer_schema(&self, steps_out: usize) -> Option<DFSchemaRef> {
+        // steps_out=1 → last element, steps_out=2 → second-to-last, etc.
+        // Returns None for steps_out=0 or steps_out > stack depth.
+        self.outer_schemas
+            .len()
+            .checked_sub(steps_out)
+            .and_then(|idx| self.outer_schemas.get(idx).cloned())
+    }
+
     fn push_lambda_parameters(
         &mut self,
         lambda_parameters: Vec<FieldRef>,