diff --git a/datafusion/substrait/src/logical_plan/producer/expr/field_reference.rs b/datafusion/substrait/src/logical_plan/producer/expr/field_reference.rs index aa34317a6e292..e4e50b3aaf588 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/field_reference.rs +++ b/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 { + substrait_field_ref_with_root(index, RootType::RootReference(RootReference {})) +} + +fn substrait_field_ref_with_root( + index: usize, + root_type: RootType, ) -> datafusion::common::Result { 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 { - // 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)] diff --git a/datafusion/substrait/src/logical_plan/producer/expr/mod.rs b/datafusion/substrait/src/logical_plan/producer/expr/mod.rs index c728af2f1458d..a95e46a8c06ca 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/mod.rs +++ b/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) => { diff --git a/datafusion/substrait/src/logical_plan/producer/expr/subquery.rs b/datafusion/substrait/src/logical_plan/producer/expr/subquery.rs index 97699c2132781..4affbbaf29aef 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/subquery.rs +++ b/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> { + 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 { - 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 { - 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( diff --git a/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs b/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs index 6d54d32cad3db..c3a25657019c5 100644 --- a/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs +++ b/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 { + 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 { + 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, } 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 { + // 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, diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 018e1aef80ea1..2492a295591a5 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -751,6 +751,141 @@ async fn roundtrip_not_exists_substrait() -> Result<()> { Ok(()) } +/// Assert that the Substrait plan contains a field reference with an +/// `OuterReference` root type at the given `steps_out` depth. +fn assert_contains_outer_reference(proto: &Plan, steps_out: u32) { + let proto_str = format!("{proto:?}"); + let expected = format!("OuterReference {{ steps_out: {steps_out} }}"); + assert!( + proto_str.contains(&expected), + "expected Substrait plan to contain an outer reference with steps_out={steps_out}" + ); +} + +#[tokio::test] +async fn roundtrip_correlated_exists() -> Result<()> { + let ctx = create_context().await?; + let plan = ctx + .sql("SELECT b FROM data WHERE EXISTS (SELECT 1 FROM data2 WHERE data2.a = data.a)") + .await? + .into_unoptimized_plan(); + + let proto = to_substrait_plan(&plan, &ctx.state())?; + assert_contains_outer_reference(&proto, 1); + + let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; + assert_eq!(plan.schema(), plan2.schema()); + assert_snapshot!( + plan2, + @r" + Projection: data.b + Filter: EXISTS () + Subquery: + Projection: Int64(1) + Filter: data2.a = outer_ref(data.a) + TableScan: data2 + TableScan: data + " + ); + Ok(()) +} + +#[tokio::test] +async fn roundtrip_correlated_in_subquery() -> Result<()> { + let ctx = create_context().await?; + let plan = ctx + .sql("SELECT b FROM data WHERE a IN (SELECT data2.a FROM data2 WHERE data2.d = data.d)") + .await? + .into_unoptimized_plan(); + + let proto = to_substrait_plan(&plan, &ctx.state())?; + assert_contains_outer_reference(&proto, 1); + + let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; + assert_eq!(plan.schema(), plan2.schema()); + assert_snapshot!( + plan2, + @r" + Projection: data.b + Filter: data.a IN () + Subquery: + Projection: data2.a + Filter: data2.d = outer_ref(data.d) + TableScan: data2 + TableScan: data + " + ); + Ok(()) +} + +#[tokio::test] +async fn roundtrip_correlated_scalar_subquery() -> Result<()> { + let ctx = create_context().await?; + let plan = ctx + .sql("SELECT a FROM data WHERE a < (SELECT sum(data2.a) FROM data2 WHERE data2.a = data.a)") + .await? + .into_unoptimized_plan(); + + let proto = to_substrait_plan(&plan, &ctx.state())?; + assert_contains_outer_reference(&proto, 1); + + let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; + assert_eq!(plan.schema(), plan2.schema()); + assert_snapshot!( + plan2, + @r" + Projection: data.a + Filter: data.a < () + Subquery: + Projection: sum(data2.a) + Aggregate: groupBy=[[]], aggr=[[sum(data2.a)]] + Filter: data2.a = outer_ref(data.a) + TableScan: data2 + TableScan: data + " + ); + Ok(()) +} + +#[tokio::test] +async fn roundtrip_nested_correlated_subquery() -> Result<()> { + let ctx = create_context().await?; + // The innermost subquery references `data.a` from the outermost query, + // two subquery boundaries out. + let plan = ctx + .sql( + "SELECT b FROM data WHERE EXISTS (\ + SELECT 1 FROM data2 WHERE data2.a = data.a AND EXISTS (\ + SELECT 1 FROM book WHERE book.isbn = data.a))", + ) + .await? + .into_unoptimized_plan(); + + let proto = to_substrait_plan(&plan, &ctx.state())?; + assert_contains_outer_reference(&proto, 1); + assert_contains_outer_reference(&proto, 2); + + let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; + assert_eq!(plan.schema(), plan2.schema()); + assert_snapshot!( + plan2, + @r" + Projection: data.b + Filter: EXISTS () + Subquery: + Projection: Int64(1) + Filter: data2.a = outer_ref(data.a) AND EXISTS () + Subquery: + Projection: Int64(1) + Filter: book.isbn = outer_ref(data.a) + TableScan: book + TableScan: data2 + TableScan: data + " + ); + Ok(()) +} + #[tokio::test] async fn roundtrip_not_exists_filter_left_anti_join() -> Result<()> { let plan = generate_plan_from_sql(