diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/aggregate_function.rs b/datafusion/substrait/src/logical_plan/consumer/expr/aggregate_function.rs index a82868b6abd0e..acafa2b7a2671 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/aggregate_function.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/aggregate_function.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements +use crate::logical_plan::consumer::SubstraitConsumerContext; + use crate::logical_plan::consumer::{ SubstraitConsumer, from_substrait_func_args, substrait_fun_name, }; @@ -36,6 +38,7 @@ pub async fn from_substrait_agg_func( filter: Option>, order_by: Vec, distinct: bool, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result> { validate_aggregation_phase(f.phase)?; let Some(fn_signature) = consumer @@ -59,7 +62,8 @@ pub async fn from_substrait_agg_func( ) })?; - let args = from_substrait_func_args(consumer, &f.arguments, input_schema).await?; + let args = + from_substrait_func_args(consumer, &f.arguments, input_schema, context).await?; // Datafusion does not support aggregate functions with no arguments, so // we inject a dummy argument that does not affect the query, but allows diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/cast.rs b/datafusion/substrait/src/logical_plan/consumer/expr/cast.rs index 3c8deb34c1ad8..79a6e35986828 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/cast.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/cast.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements +use crate::logical_plan::consumer::SubstraitConsumerContext; + use crate::logical_plan::consumer::{ SubstraitConsumer, field_from_substrait_type_without_names, }; @@ -14,6 +16,7 @@ pub async fn from_cast( consumer: &impl SubstraitConsumer, cast: &substrait_expression::Cast, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { match cast.r#type.as_ref() { Some(output_type) => { @@ -22,6 +25,7 @@ pub async fn from_cast( .consume_expression( cast.input.as_ref().unwrap().as_ref(), input_schema, + context, ) .await?, ); diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs b/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs index ae7959b4ae310..9fb17b3b2cc7e 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use crate::logical_plan::consumer::SubstraitConsumer; +use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; use datafusion::common::{Column, DFSchema, not_impl_err, substrait_err}; use datafusion::logical_expr::Expr; use std::sync::Arc; @@ -15,14 +15,16 @@ pub async fn from_field_reference( consumer: &impl SubstraitConsumer, field_ref: &FieldReference, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_substrait_field_reference(consumer, field_ref, input_schema) + from_substrait_field_reference(consumer, field_ref, input_schema, context) } pub(crate) fn from_substrait_field_reference( consumer: &impl SubstraitConsumer, field_ref: &FieldReference, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { match &field_ref.reference_type { Some(DirectReference(direct)) => match &direct.reference_type.as_ref() { @@ -38,14 +40,14 @@ pub(crate) fn from_substrait_field_reference( Column::from(input_schema.qualified_field(field_idx)), )), Some(RootType::OuterReference(outer_ref)) => { - resolve_outer_reference(consumer, outer_ref, field_idx) + resolve_outer_reference(consumer, outer_ref, field_idx, context) } Some(RootType::Expression(_)) => not_impl_err!( "Expression root type in field reference is not supported" ), Some(RootType::LambdaParameterReference( LambdaParameterReference { steps_out }, - )) => consumer.lambda_variable(*steps_out as usize, field_idx), + )) => context.lambda_variable(*steps_out as usize, field_idx), } } _ => not_impl_err!( @@ -57,12 +59,13 @@ pub(crate) fn from_substrait_field_reference( } fn resolve_outer_reference( - consumer: &impl SubstraitConsumer, + _consumer: &impl SubstraitConsumer, outer_ref: &substrait::proto::expression::field_reference::OuterReference, field_idx: usize, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let steps_out = outer_ref.steps_out as usize; - let Some(outer_schema) = consumer.get_outer_schema(steps_out) else { + let Some(outer_schema) = context.get_outer_schema(steps_out) else { return substrait_err!( "OuterReference with steps_out={steps_out} \ but no outer schema is available" @@ -75,6 +78,7 @@ fn resolve_outer_reference( #[cfg(test)] mod tests { + use crate::logical_plan::consumer::SubstraitConsumerContext; use datafusion::{ common::{DFSchema, assert_contains}, prelude::SessionContext, @@ -91,9 +95,7 @@ mod tests { use crate::{ extensions::Extensions, - logical_plan::consumer::{ - DefaultSubstraitConsumer, SubstraitConsumer, from_field_reference, - }, + logical_plan::consumer::{DefaultSubstraitConsumer, from_field_reference}, }; #[tokio::test] @@ -104,10 +106,14 @@ mod tests { let session_state = SessionContext::new().state(); let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); - let err = - from_field_reference(&consumer, &lambda_field_ref, DFSchema::empty_ref()) - .await - .unwrap_err(); + let err = from_field_reference( + &consumer, + &lambda_field_ref, + DFSchema::empty_ref(), + &SubstraitConsumerContext::default(), + ) + .await + .unwrap_err(); assert_contains!(err.to_string(), "No lambda at 99 steps out, got only 0"); } @@ -119,8 +125,9 @@ mod tests { let extensions = Extensions::default(); let session_state = SessionContext::new().state(); let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); - let _names = consumer - .push_lambda_parameters( + let (context, _names) = SubstraitConsumerContext::default() + .with_lambda_parameters( + &consumer, &[Type { kind: Some(Kind::I64(I64::default())), }], @@ -128,10 +135,14 @@ mod tests { ) .unwrap(); - let err = - from_field_reference(&consumer, &lambda_field_ref, DFSchema::empty_ref()) - .await - .unwrap_err(); + let err = from_field_reference( + &consumer, + &lambda_field_ref, + DFSchema::empty_ref(), + &context, + ) + .await + .unwrap_err(); assert_contains!( err.to_string(), diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/function_arguments.rs b/datafusion/substrait/src/logical_plan/consumer/expr/function_arguments.rs index d1b4222933792..3d2e1557c8eac 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/function_arguments.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/function_arguments.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use crate::logical_plan::consumer::SubstraitConsumer; +use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; use datafusion::common::{DFSchema, not_impl_err}; use datafusion::logical_expr::Expr; use substrait::proto::FunctionArgument; @@ -13,11 +13,14 @@ pub async fn from_substrait_func_args( consumer: &impl SubstraitConsumer, arguments: &Vec, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result> { let mut args: Vec = Vec::with_capacity(arguments.len()); for arg in arguments { let arg_expr = match &arg.arg_type { - Some(ArgType::Value(e)) => consumer.consume_expression(e, input_schema).await, + Some(ArgType::Value(e)) => { + consumer.consume_expression(e, input_schema, context).await + } _ => not_impl_err!("Function argument non-Value type not supported"), }; args.push(arg_expr?); diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/if_then.rs b/datafusion/substrait/src/logical_plan/consumer/expr/if_then.rs index 759ea5323eb19..ee8f9905e67df 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/if_then.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/if_then.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use crate::logical_plan::consumer::SubstraitConsumer; +use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; use datafusion::common::DFSchema; use datafusion::logical_expr::{Case, Expr}; use substrait::proto::expression::IfThen; @@ -11,6 +11,7 @@ pub async fn from_if_then( consumer: &impl SubstraitConsumer, if_then: &IfThen, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { // Parse `ifs` // If the first element does not have a `then` part, then we can assume it's a base expression @@ -23,7 +24,11 @@ pub async fn from_if_then( if if_expr.then.is_none() { expr = Some(Box::new( consumer - .consume_expression(if_expr.r#if.as_ref().unwrap(), input_schema) + .consume_expression( + if_expr.r#if.as_ref().unwrap(), + input_schema, + context, + ) .await?, )); continue; @@ -32,12 +37,20 @@ pub async fn from_if_then( when_then_expr.push(( Box::new( consumer - .consume_expression(if_expr.r#if.as_ref().unwrap(), input_schema) + .consume_expression( + if_expr.r#if.as_ref().unwrap(), + input_schema, + context, + ) .await?, ), Box::new( consumer - .consume_expression(if_expr.then.as_ref().unwrap(), input_schema) + .consume_expression( + if_expr.then.as_ref().unwrap(), + input_schema, + context, + ) .await?, ), )); @@ -45,7 +58,9 @@ pub async fn from_if_then( // Parse `else` let else_expr = match &if_then.r#else { Some(e) => Some(Box::new( - consumer.consume_expression(e, input_schema).await?, + consumer + .consume_expression(e, input_schema, context) + .await?, )), None => None, }; diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/lambda.rs b/datafusion/substrait/src/logical_plan/consumer/expr/lambda.rs index 6269b0556623d..bcf4fe7459d81 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/lambda.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/lambda.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements +use crate::logical_plan::consumer::SubstraitConsumerContext; + use datafusion::{ common::{DFSchema, substrait_err}, prelude::{Expr, lambda}, @@ -14,26 +16,29 @@ pub async fn from_lambda( consumer: &impl SubstraitConsumer, expr: &proto::expression::Lambda, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let Some(parameters) = expr.parameters.as_ref() else { return substrait_err!("Lambda expression without parameters is not allowed"); }; - let names = consumer.push_lambda_parameters(¶meters.types, input_schema)?; + let (child_context, names) = + context.with_lambda_parameters(consumer, ¶meters.types, input_schema)?; let Some(body) = expr.body.as_ref() else { return substrait_err!("Lambda expression without body is not allowed"); }; - let body = consumer.consume_expression(body, input_schema).await?; - - consumer.pop_lambda_parameters(); + let body = consumer + .consume_expression(body, input_schema, &child_context) + .await?; Ok(lambda(names, body)) } #[cfg(test)] mod tests { + use crate::logical_plan::consumer::SubstraitConsumerContext; use datafusion::{ common::{DFSchema, assert_contains}, prelude::SessionContext, @@ -56,9 +61,14 @@ mod tests { let session_state = SessionContext::new().state(); let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); - let err = from_lambda(&consumer, &lambda, DFSchema::empty_ref()) - .await - .unwrap_err(); + let err = from_lambda( + &consumer, + &lambda, + DFSchema::empty_ref(), + &SubstraitConsumerContext::default(), + ) + .await + .unwrap_err(); assert_contains!( err.to_string(), @@ -77,9 +87,14 @@ mod tests { let session_state = SessionContext::new().state(); let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); - let err = from_lambda(&consumer, &lambda, DFSchema::empty_ref()) - .await - .unwrap_err(); + let err = from_lambda( + &consumer, + &lambda, + DFSchema::empty_ref(), + &SubstraitConsumerContext::default(), + ) + .await + .unwrap_err(); assert_contains!( err.to_string(), diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs b/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs index c7ad51c017408..1da713abbd709 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs @@ -28,6 +28,8 @@ pub use singular_or_list::*; pub use subquery::*; pub use window_function::*; +use crate::logical_plan::consumer::SubstraitConsumerContext; + use crate::extensions::Extensions; use crate::logical_plan::consumer::{ DefaultSubstraitConsumer, SubstraitConsumer, from_substrait_named_struct, @@ -46,46 +48,71 @@ pub async fn from_substrait_rex( consumer: &impl SubstraitConsumer, expression: &Expression, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { match &expression.rex_type { Some(t) => match t { - RexType::Literal(expr) => consumer.consume_literal(expr).await, + RexType::Literal(expr) => consumer.consume_literal(expr, context).await, RexType::Selection(expr) => { - consumer.consume_field_reference(expr, input_schema).await + consumer + .consume_field_reference(expr, input_schema, context) + .await } RexType::ScalarFunction(expr) => { - consumer.consume_scalar_function(expr, input_schema).await + consumer + .consume_scalar_function(expr, input_schema, context) + .await } RexType::WindowFunction(expr) => { - consumer.consume_window_function(expr, input_schema).await + consumer + .consume_window_function(expr, input_schema, context) + .await + } + RexType::IfThen(expr) => { + consumer.consume_if_then(expr, input_schema, context).await } - RexType::IfThen(expr) => consumer.consume_if_then(expr, input_schema).await, RexType::SwitchExpression(expr) => { - consumer.consume_switch(expr, input_schema).await + consumer.consume_switch(expr, input_schema, context).await } RexType::SingularOrList(expr) => { - consumer.consume_singular_or_list(expr, input_schema).await + consumer + .consume_singular_or_list(expr, input_schema, context) + .await } RexType::MultiOrList(expr) => { - consumer.consume_multi_or_list(expr, input_schema).await + consumer + .consume_multi_or_list(expr, input_schema, context) + .await } RexType::Cast(expr) => { - consumer.consume_cast(expr.as_ref(), input_schema).await + consumer + .consume_cast(expr.as_ref(), input_schema, context) + .await } RexType::Subquery(expr) => { - consumer.consume_subquery(expr.as_ref(), input_schema).await + consumer + .consume_subquery(expr.as_ref(), input_schema, context) + .await + } + RexType::Nested(expr) => { + consumer.consume_nested(expr, input_schema, context).await } - RexType::Nested(expr) => consumer.consume_nested(expr, input_schema).await, #[expect(deprecated)] - RexType::Enum(expr) => consumer.consume_enum(expr, input_schema).await, + RexType::Enum(expr) => { + consumer.consume_enum(expr, input_schema, context).await + } RexType::DynamicParameter(expr) => { - consumer.consume_dynamic_parameter(expr, input_schema).await + consumer + .consume_dynamic_parameter(expr, input_schema, context) + .await } RexType::Lambda(lambda) => { - consumer.consume_lambda(lambda.as_ref(), input_schema).await + consumer + .consume_lambda(lambda.as_ref(), input_schema, context) + .await } RexType::LambdaInvocation(_) => { not_impl_err!("Lambda invocations are not supported") @@ -140,7 +167,11 @@ pub async fn from_substrait_extended_expr( } }?; let expr = consumer - .consume_expression(scalar_expr, &input_schema) + .consume_expression( + scalar_expr, + &input_schema, + &SubstraitConsumerContext::default(), + ) .await?; let output_field = expr.to_field(&input_schema)?.1; let mut names_idx = 0; @@ -179,10 +210,13 @@ pub async fn from_substrait_rex_vec( consumer: &impl SubstraitConsumer, exprs: &Vec, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result> { let mut expressions: Vec = Vec::with_capacity(exprs.len()); for expr in exprs { - let expression = consumer.consume_expression(expr, input_schema).await?; + let expression = consumer + .consume_expression(expr, input_schema, context) + .await?; expressions.push(expression); } Ok(expressions) @@ -191,6 +225,7 @@ pub async fn from_substrait_rex_vec( #[cfg(test)] mod tests { use crate::extensions::Extensions; + use crate::logical_plan::consumer::SubstraitConsumerContext; use crate::logical_plan::consumer::utils::tests::test_consumer; use crate::logical_plan::consumer::*; use datafusion::common::{DFSchema, assert_contains}; @@ -221,7 +256,14 @@ mod tests { extensions.register_function("count"); consumer.extensions = &extensions; - match from_substrait_rex(&consumer, &substrait, &DFSchema::empty()).await? { + match from_substrait_rex( + &consumer, + &substrait, + &DFSchema::empty(), + &SubstraitConsumerContext::default(), + ) + .await? + { Expr::WindowFunction(window_function) => { assert_eq!(window_function.params.order_by.len(), 1) } @@ -248,7 +290,14 @@ mod tests { extensions.register_function("count"); consumer.extensions = &extensions; - match from_substrait_rex(&consumer, &substrait, &DFSchema::empty()).await? { + match from_substrait_rex( + &consumer, + &substrait, + &DFSchema::empty(), + &SubstraitConsumerContext::default(), + ) + .await? + { Expr::WindowFunction(window_function) => { assert_eq!(window_function.params.args.len(), 1) } @@ -275,9 +324,14 @@ mod tests { extensions.register_function("count"); consumer.extensions = &extensions; - let err = from_substrait_rex(&consumer, &substrait, &DFSchema::empty()) - .await - .unwrap_err(); + let err = from_substrait_rex( + &consumer, + &substrait, + &DFSchema::empty(), + &SubstraitConsumerContext::default(), + ) + .await + .unwrap_err(); assert_contains!(err.to_string(), "Invalid window aggregation invocation 3"); } } diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/nested.rs b/datafusion/substrait/src/logical_plan/consumer/expr/nested.rs index 4017c0af264f4..6bd0daa600887 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/nested.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/nested.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use crate::logical_plan::consumer::SubstraitConsumer; +use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; use datafusion::common::{DFSchema, not_impl_err, substrait_err}; use datafusion::execution::FunctionRegistry; use datafusion::logical_expr::Expr; @@ -19,6 +19,7 @@ pub async fn from_nested( consumer: &impl SubstraitConsumer, nested: &Nested, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let Some(nested_type) = &nested.nested_type else { return substrait_err!("Nested expression requires a nested_type"); @@ -34,7 +35,11 @@ pub async fn from_nested( let mut args = Vec::with_capacity(list.values.len()); for value in &list.values { - args.push(consumer.consume_expression(value, input_schema).await?); + args.push( + consumer + .consume_expression(value, input_schema, context) + .await?, + ); } let make_array_udf = consumer.get_function_registry().udf("make_array")?; @@ -88,7 +93,13 @@ mod tests { })), }; - let expr = from_nested(&consumer, &nested, &schema).await?; + let expr = from_nested( + &consumer, + &nested, + &schema, + &SubstraitConsumerContext::default(), + ) + .await?; assert_eq!( format!("{expr}"), "make_array(Int64(1), Int64(2), Int64(3))" @@ -107,7 +118,13 @@ mod tests { nested_type: Some(NestedType::List(List { values: vec![] })), }; - let result = from_nested(&consumer, &nested, &schema).await; + let result = from_nested( + &consumer, + &nested, + &schema, + &SubstraitConsumerContext::default(), + ) + .await; assert!(result.is_err()); assert!( result @@ -129,7 +146,13 @@ mod tests { nested_type: None, }; - let result = from_nested(&consumer, &nested, &schema).await; + let result = from_nested( + &consumer, + &nested, + &schema, + &SubstraitConsumerContext::default(), + ) + .await; assert!(result.is_err()); assert!(result.unwrap_err().to_string().contains("nested_type")); diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs b/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs index 2be837015feae..059e517355bd0 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements +use crate::logical_plan::consumer::SubstraitConsumerContext; + use crate::logical_plan::consumer::{SubstraitConsumer, from_substrait_func_args}; use datafusion::common::Result; use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; @@ -17,6 +19,7 @@ pub async fn from_scalar_function( consumer: &impl SubstraitConsumer, f: &ScalarFunction, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> Result { let Some(fn_signature) = consumer .get_extensions() @@ -30,7 +33,8 @@ pub async fn from_scalar_function( }; let fn_name = substrait_fun_name(fn_signature); - let args = from_substrait_func_args(consumer, &f.arguments, input_schema).await?; + let args = + from_substrait_func_args(consumer, &f.arguments, input_schema, context).await?; let higher_order_func = consumer .get_function_registry() @@ -448,6 +452,7 @@ impl BuiltinExprBuilder { mod tests { use super::{arg_list_to_binary_op_tree, case_insensitive_option}; use crate::extensions::Extensions; + use crate::logical_plan::consumer::SubstraitConsumerContext; use crate::logical_plan::consumer::tests::TEST_SESSION_STATE; use crate::logical_plan::consumer::{DefaultSubstraitConsumer, SubstraitConsumer}; use datafusion::arrow::datatypes::{DataType, Field, Schema}; @@ -489,7 +494,13 @@ mod tests { let df_schema = DFSchema::try_from(schema).unwrap(); // Consume the expression and ensure we don't crash - let _ = consumer.consume_scalar_function(&func, &df_schema).await?; + let _ = consumer + .consume_scalar_function( + &func, + &df_schema, + &SubstraitConsumerContext::default(), + ) + .await?; Ok(()) } @@ -560,7 +571,13 @@ mod tests { let df_schema = DFSchema::try_from(schema).unwrap(); // Consume the expression and ensure we don't get an error - let _ = consumer.consume_scalar_function(&func, &df_schema).await?; + let _ = consumer + .consume_scalar_function( + &func, + &df_schema, + &SubstraitConsumerContext::default(), + ) + .await?; Ok(()) } @@ -624,7 +641,11 @@ mod tests { }; let result = consumer - .consume_scalar_function(&func_like, &df_schema) + .consume_scalar_function( + &func_like, + &df_schema, + &SubstraitConsumerContext::default(), + ) .await?; if let Expr::Like(like) = result { @@ -643,7 +664,11 @@ mod tests { }; let result = consumer - .consume_scalar_function(&func_not_like, &df_schema) + .consume_scalar_function( + &func_not_like, + &df_schema, + &SubstraitConsumerContext::default(), + ) .await?; if let Expr::Like(like) = result { diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/singular_or_list.rs b/datafusion/substrait/src/logical_plan/consumer/expr/singular_or_list.rs index 3d32408e69cc6..aed563e65499f 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/singular_or_list.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/singular_or_list.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements +use crate::logical_plan::consumer::SubstraitConsumerContext; + use crate::logical_plan::consumer::{SubstraitConsumer, from_substrait_rex_vec}; use datafusion::common::DFSchema; use datafusion::logical_expr::Expr; @@ -12,16 +14,18 @@ pub async fn from_singular_or_list( consumer: &impl SubstraitConsumer, expr: &SingularOrList, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let substrait_expr = expr.value.as_ref().unwrap(); let substrait_list = expr.options.as_ref(); Ok(Expr::InList(InList { expr: Box::new( consumer - .consume_expression(substrait_expr, input_schema) + .consume_expression(substrait_expr, input_schema, context) .await?, ), - list: from_substrait_rex_vec(consumer, substrait_list, input_schema).await?, + list: from_substrait_rex_vec(consumer, substrait_list, input_schema, context) + .await?, negated: false, })) } diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/subquery.rs b/datafusion/substrait/src/logical_plan/consumer/expr/subquery.rs index b9d10ba66d3cf..5d07faf3e9d0a 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/subquery.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/subquery.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use crate::logical_plan::consumer::SubstraitConsumer; +use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; use datafusion::common::{DFSchema, Spans, substrait_datafusion_err, substrait_err}; use datafusion::logical_expr::expr::{Exists, InSubquery, SetComparison, SetQuantifier}; use datafusion::logical_expr::{Expr, LogicalPlan, Operator, Subquery}; @@ -18,22 +18,22 @@ use substrait::proto::expression::subquery::set_predicate::PredicateOp; /// /// Substrait represents correlated references using `OuterReference` /// field references with a `steps_out` depth. To resolve these, -/// the consumer maintains a stack of outer schemas. +/// the context carries the enclosing query schemas. async fn consume_subquery_rel( consumer: &impl SubstraitConsumer, rel: &Rel, outer_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - consumer.push_outer_schema(Arc::new(outer_schema.clone())); - let result = consumer.consume_rel(rel).await; - consumer.pop_outer_schema(); - result + let child_context = context.with_outer_schema(Arc::new(outer_schema.clone())); + consumer.consume_rel(rel, &child_context).await } pub async fn from_subquery( consumer: &impl SubstraitConsumer, subquery: &substrait_expression::Subquery, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { match &subquery.subquery_type { Some(subquery_type) => match subquery_type { @@ -46,14 +46,22 @@ pub async fn from_subquery( let needle_expr = &in_predicate.needles[0]; let haystack_expr = &in_predicate.haystack; if let Some(haystack_expr) = haystack_expr { - let haystack_expr = - consume_subquery_rel(consumer, haystack_expr, input_schema) - .await?; + let haystack_expr = consume_subquery_rel( + consumer, + haystack_expr, + input_schema, + context, + ) + .await?; let outer_refs = haystack_expr.all_out_ref_exprs(); Ok(Expr::InSubquery(InSubquery { expr: Box::new( consumer - .consume_expression(needle_expr, input_schema) + .consume_expression( + needle_expr, + input_schema, + context, + ) .await?, ), subquery: Subquery { @@ -75,6 +83,7 @@ pub async fn from_subquery( consumer, &(query.input.clone()).unwrap_or_default(), input_schema, + context, ) .await?; let outer_ref_columns = plan.all_out_ref_exprs(); @@ -93,6 +102,7 @@ pub async fn from_subquery( consumer, &relation.clone().unwrap_or_default(), input_schema, + context, ) .await?; let outer_ref_columns = plan.all_out_ref_exprs(); @@ -143,8 +153,11 @@ pub async fn from_subquery( } }; - let left_expr = consumer.consume_expression(left, input_schema).await?; - let plan = consume_subquery_rel(consumer, right, input_schema).await?; + let left_expr = consumer + .consume_expression(left, input_schema, context) + .await?; + let plan = + consume_subquery_rel(consumer, right, input_schema, context).await?; let outer_ref_columns = plan.all_out_ref_exprs(); Ok(Expr::SetComparison(SetComparison::new( diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/window_function.rs b/datafusion/substrait/src/logical_plan/consumer/expr/window_function.rs index 77cf317db546c..f34e75c358931 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/window_function.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/window_function.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements +use crate::logical_plan::consumer::SubstraitConsumerContext; + use super::aggregate_function::validate_aggregation_phase; use crate::logical_plan::consumer::{ SubstraitConsumer, from_substrait_func_args, from_substrait_rex_vec, @@ -26,6 +28,7 @@ pub async fn from_window_function( consumer: &impl SubstraitConsumer, window: &WindowFunction, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { validate_aggregation_phase(window.phase)?; let Some(fn_signature) = consumer @@ -54,7 +57,7 @@ pub async fn from_window_function( }?; let mut order_by = - from_substrait_sorts(consumer, &window.sorts, input_schema).await?; + from_substrait_sorts(consumer, &window.sorts, input_schema, context).await?; let bound_units = match BoundsType::try_from(window.bounds_type).map_err(|e| { plan_datafusion_err!("Invalid bound type {}: {e}", window.bounds_type) @@ -86,7 +89,8 @@ pub async fn from_window_function( let args = if fun.name() == "count" && window.arguments.is_empty() { vec![Expr::Literal(ScalarValue::Int64(Some(1)), None)] } else { - from_substrait_func_args(consumer, &window.arguments, input_schema).await? + from_substrait_func_args(consumer, &window.arguments, input_schema, context) + .await? }; let distinct = match AggregationInvocation::try_from(window.invocation).map_err(|e| { @@ -107,6 +111,7 @@ pub async fn from_window_function( consumer, &window.partitions, input_schema, + context, ) .await?, order_by, diff --git a/datafusion/substrait/src/logical_plan/consumer/plan.rs b/datafusion/substrait/src/logical_plan/consumer/plan.rs index c7c119991c461..2cc2996bac05f 100644 --- a/datafusion/substrait/src/logical_plan/consumer/plan.rs +++ b/datafusion/substrait/src/logical_plan/consumer/plan.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements +use crate::logical_plan::consumer::SubstraitConsumerContext; + use super::utils::{make_renamed_schema, rename_expressions}; use super::{DefaultSubstraitConsumer, SubstraitConsumer}; use crate::extensions::Extensions; @@ -35,10 +37,16 @@ pub async fn from_substrait_plan_with_consumer( 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::Rel(rel) => Ok(consumer + .consume_rel(rel, &SubstraitConsumerContext::default()) + .await?), plan_rel::RelType::Root(root) => { - let plan = - consumer.consume_rel(root.input.as_ref().unwrap()).await?; + let plan = consumer + .consume_rel( + root.input.as_ref().unwrap(), + &SubstraitConsumerContext::default(), + ) + .await?; if root.names.is_empty() { // Backwards compatibility for plans missing names return Ok(plan); diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs index 14fe946c4dde7..c2d792b77182f 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements +use crate::logical_plan::consumer::SubstraitConsumerContext; + use crate::logical_plan::consumer::{NameTracker, SubstraitConsumer}; use crate::logical_plan::consumer::{from_substrait_agg_func, from_substrait_sorts}; use datafusion::common::{Column, DFSchemaRef, internal_err, not_impl_err}; @@ -16,13 +18,16 @@ use substrait::proto::aggregate_rel::Grouping; pub async fn from_aggregate_rel( consumer: &impl SubstraitConsumer, agg: &AggregateRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { if let Some(input) = agg.input.as_ref() { - let input = LogicalPlanBuilder::from(consumer.consume_rel(input).await?); + let input = LogicalPlanBuilder::from(consumer.consume_rel(input, context).await?); let mut ref_group_exprs = Vec::with_capacity(agg.grouping_expressions.len()); for e in &agg.grouping_expressions { - let x = consumer.consume_expression(e, input.schema()).await?; + let x = consumer + .consume_expression(e, input.schema(), context) + .await?; ref_group_exprs.push(x); } @@ -38,6 +43,7 @@ pub async fn from_aggregate_rel( &agg.groupings[0], &ref_group_exprs, input.schema(), + context, ) .await?, ); @@ -50,6 +56,7 @@ pub async fn from_aggregate_rel( grouping, &ref_group_exprs, input.schema(), + context, ) .await?; grouping_sets.push(grouping_set); @@ -66,7 +73,9 @@ pub async fn from_aggregate_rel( for m in &agg.measures { let filter = match &m.filter { Some(fil) => Some(Box::new( - consumer.consume_expression(fil, input.schema()).await?, + consumer + .consume_expression(fil, input.schema(), context) + .await?, )), None => None, }; @@ -80,7 +89,8 @@ pub async fn from_aggregate_rel( _ => false, }; let order_by = - from_substrait_sorts(consumer, &f.sorts, input.schema()).await?; + from_substrait_sorts(consumer, &f.sorts, input.schema(), context) + .await?; from_substrait_agg_func( consumer, @@ -89,6 +99,7 @@ pub async fn from_aggregate_rel( filter, order_by, distinct, + context, ) .await } @@ -163,11 +174,14 @@ async fn from_substrait_grouping( grouping: &Grouping, expressions: &[Expr], input_schema: &DFSchemaRef, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result> { let mut group_exprs = vec![]; if !grouping.grouping_expressions.is_empty() { for e in &grouping.grouping_expressions { - let expr = consumer.consume_expression(e, input_schema).await?; + let expr = consumer + .consume_expression(e, input_schema, context) + .await?; group_exprs.push(expr); } return Ok(group_exprs); 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 ef8d93d0a0822..3ca075152779b 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/cross_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/cross_rel.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use crate::logical_plan::consumer::SubstraitConsumer; +use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder}; use datafusion::logical_expr::requalify_sides_if_needed; @@ -12,12 +12,17 @@ use substrait::proto::CrossRel; pub async fn from_cross_rel( consumer: &impl SubstraitConsumer, cross: &CrossRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let left = LogicalPlanBuilder::from( - consumer.consume_rel(cross.left.as_ref().unwrap()).await?, + consumer + .consume_rel(cross.left.as_ref().unwrap(), context) + .await?, ); let right = LogicalPlanBuilder::from( - consumer.consume_rel(cross.right.as_ref().unwrap()).await?, + consumer + .consume_rel(cross.right.as_ref().unwrap(), context) + .await?, ); 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/exchange_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/exchange_rel.rs index fd9419bc592df..cee70e4e7b9e4 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/exchange_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/exchange_rel.rs @@ -2,8 +2,8 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use crate::logical_plan::consumer::SubstraitConsumer; use crate::logical_plan::consumer::from_substrait_field_reference; +use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; use datafusion::common::{not_impl_err, substrait_err}; use datafusion::logical_expr::{LogicalPlan, Partitioning, Repartition}; use std::sync::Arc; @@ -13,11 +13,12 @@ use substrait::proto::exchange_rel::ExchangeKind; pub async fn from_exchange_rel( consumer: &impl SubstraitConsumer, exchange: &ExchangeRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let Some(input) = exchange.input.as_ref() else { return substrait_err!("Unexpected empty input in ExchangeRel"); }; - let input = Arc::new(consumer.consume_rel(input).await?); + let input = Arc::new(consumer.consume_rel(input, context).await?); let Some(exchange_kind) = &exchange.exchange_kind else { return substrait_err!("Unexpected empty input in ExchangeRel"); @@ -29,8 +30,12 @@ pub async fn from_exchange_rel( let mut partition_columns = Vec::with_capacity(scatter_fields.fields.len()); let input_schema = input.schema(); for field_ref in &scatter_fields.fields { - let column = - from_substrait_field_reference(consumer, field_ref, input_schema)?; + let column = from_substrait_field_reference( + consumer, + field_ref, + input_schema, + context, + )?; partition_columns.push(column); } Partitioning::Hash(partition_columns, exchange.partition_count as usize) diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/fetch_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/fetch_rel.rs index ec51c8b231dfe..c5fde0ebb31d2 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/fetch_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/fetch_rel.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use crate::logical_plan::consumer::SubstraitConsumer; +use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; use async_recursion::async_recursion; use datafusion::common::{DFSchema, DFSchemaRef, not_impl_err}; use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, lit}; @@ -12,16 +12,19 @@ use substrait::proto::{FetchRel, fetch_rel}; pub async fn from_fetch_rel( consumer: &impl SubstraitConsumer, fetch: &FetchRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { if let Some(input) = fetch.input.as_ref() { - let input = LogicalPlanBuilder::from(consumer.consume_rel(input).await?); + let input = LogicalPlanBuilder::from(consumer.consume_rel(input, context).await?); let empty_schema = DFSchemaRef::new(DFSchema::empty()); let offset = match &fetch.offset_mode { #[expect(deprecated)] Some(fetch_rel::OffsetMode::Offset(offset)) => Some(lit(*offset)), - Some(fetch_rel::OffsetMode::OffsetExpr(expr)) => { - Some(consumer.consume_expression(expr, &empty_schema).await?) - } + Some(fetch_rel::OffsetMode::OffsetExpr(expr)) => Some( + consumer + .consume_expression(expr, &empty_schema, context) + .await?, + ), None => None, }; let count = match &fetch.count_mode { @@ -30,9 +33,11 @@ pub async fn from_fetch_rel( // -1 means that ALL records should be returned, equivalent to None (*count != -1).then(|| lit(*count)) } - Some(fetch_rel::CountMode::CountExpr(expr)) => { - Some(consumer.consume_expression(expr, &empty_schema).await?) - } + Some(fetch_rel::CountMode::CountExpr(expr)) => Some( + consumer + .consume_expression(expr, &empty_schema, context) + .await?, + ), None => None, }; input.limit_by_expr(offset, count)?.build() diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/filter_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/filter_rel.rs index c9c7e624c036b..84a6b12007fb3 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/filter_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/filter_rel.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use crate::logical_plan::consumer::SubstraitConsumer; +use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; use async_recursion::async_recursion; use datafusion::common::not_impl_err; use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder}; @@ -12,12 +12,13 @@ use substrait::proto::FilterRel; pub async fn from_filter_rel( consumer: &impl SubstraitConsumer, filter: &FilterRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { if let Some(input) = filter.input.as_ref() { - let input = LogicalPlanBuilder::from(consumer.consume_rel(input).await?); + let input = LogicalPlanBuilder::from(consumer.consume_rel(input, context).await?); if let Some(condition) = filter.condition.as_ref() { let expr = consumer - .consume_expression(condition, input.schema()) + .consume_expression(condition, input.schema(), context) .await?; input.filter(expr)?.build() } else { 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 7723e781e2578..c20cec0bbfbac 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/join_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/join_rel.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use crate::logical_plan::consumer::SubstraitConsumer; +use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; 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; @@ -15,16 +15,21 @@ use substrait::proto::{JoinRel, join_rel}; pub async fn from_join_rel( consumer: &impl SubstraitConsumer, join: &JoinRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { if join.post_join_filter.is_some() { 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?, + consumer + .consume_rel(join.left.as_ref().unwrap(), context) + .await?, ); let right = LogicalPlanBuilder::from( - consumer.consume_rel(join.right.as_ref().unwrap()).await?, + consumer + .consume_rel(join.right.as_ref().unwrap(), context) + .await?, ); let (left, right, _requalified) = requalify_sides_if_needed(left, right)?; @@ -37,7 +42,9 @@ pub async fn from_join_rel( // Otherwise, build join with only the filter, without join keys match &join.expression.as_ref() { Some(expr) => { - let on = consumer.consume_expression(expr, &in_join_schema).await?; + let on = consumer + .consume_expression(expr, &in_join_schema, context) + .await?; // The join expression can contain both equal and non-equal ops. // As of datafusion 31.0.0, the equal and non equal join conditions are in separate fields. // So we extract each part as follows: diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/mod.rs b/datafusion/substrait/src/logical_plan/consumer/rel/mod.rs index d109802b66d01..398a13eefd70b 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/mod.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/mod.rs @@ -24,8 +24,8 @@ pub use read_rel::*; pub use set_rel::*; pub use sort_rel::*; -use crate::logical_plan::consumer::SubstraitConsumer; use crate::logical_plan::consumer::utils::NameTracker; +use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; use async_recursion::async_recursion; use datafusion::common::{Column, not_impl_err, substrait_datafusion_err, substrait_err}; use datafusion::logical_expr::builder::project; @@ -40,25 +40,34 @@ use substrait::proto::{Rel, RelCommon, rel_common}; pub async fn from_substrait_rel( consumer: &impl SubstraitConsumer, relation: &Rel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let plan: datafusion::common::Result = match &relation.rel_type { Some(rel_type) => match rel_type { - RelType::Read(rel) => consumer.consume_read(rel).await, - RelType::Filter(rel) => consumer.consume_filter(rel).await, - RelType::Fetch(rel) => consumer.consume_fetch(rel).await, - RelType::Aggregate(rel) => consumer.consume_aggregate(rel).await, - RelType::Sort(rel) => consumer.consume_sort(rel).await, - RelType::Join(rel) => consumer.consume_join(rel).await, - RelType::Project(rel) => consumer.consume_project(rel).await, - RelType::Set(rel) => consumer.consume_set(rel).await, - RelType::ExtensionSingle(rel) => consumer.consume_extension_single(rel).await, - RelType::ExtensionMulti(rel) => consumer.consume_extension_multi(rel).await, - RelType::ExtensionLeaf(rel) => consumer.consume_extension_leaf(rel).await, - RelType::Cross(rel) => consumer.consume_cross(rel).await, + RelType::Read(rel) => consumer.consume_read(rel, context).await, + RelType::Filter(rel) => consumer.consume_filter(rel, context).await, + RelType::Fetch(rel) => consumer.consume_fetch(rel, context).await, + RelType::Aggregate(rel) => consumer.consume_aggregate(rel, context).await, + RelType::Sort(rel) => consumer.consume_sort(rel, context).await, + RelType::Join(rel) => consumer.consume_join(rel, context).await, + RelType::Project(rel) => consumer.consume_project(rel, context).await, + RelType::Set(rel) => consumer.consume_set(rel, context).await, + RelType::ExtensionSingle(rel) => { + consumer.consume_extension_single(rel, context).await + } + RelType::ExtensionMulti(rel) => { + consumer.consume_extension_multi(rel, context).await + } + RelType::ExtensionLeaf(rel) => { + consumer.consume_extension_leaf(rel, context).await + } + RelType::Cross(rel) => consumer.consume_cross(rel, context).await, RelType::Window(rel) => { - consumer.consume_consistent_partition_window(rel).await + consumer + .consume_consistent_partition_window(rel, context) + .await } - RelType::Exchange(rel) => consumer.consume_exchange(rel).await, + RelType::Exchange(rel) => consumer.consume_exchange(rel, context).await, rt => not_impl_err!("{rt:?} rel not supported yet"), }, None => return substrait_err!("rel must set rel_type"), diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs index 52083aa365a11..977684917c3a8 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs @@ -2,8 +2,8 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use crate::logical_plan::consumer::SubstraitConsumer; use crate::logical_plan::consumer::utils::NameTracker; +use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; use async_recursion::async_recursion; use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion::common::{Column, not_impl_err}; @@ -19,9 +19,10 @@ use substrait::proto::ProjectRel; pub async fn from_project_rel( consumer: &impl SubstraitConsumer, p: &ProjectRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { if let Some(input) = p.input.as_ref() { - let input = consumer.consume_rel(input).await?; + let input = consumer.consume_rel(input, context).await?; let original_schema = Arc::clone(input.schema()); // Ensure that all expressions have a unique display name, so that @@ -44,7 +45,9 @@ pub async fn from_project_rel( // Expr contains Arc with interior mutability but is intentionally used as hash key let mut window_exprs: HashSet = HashSet::new(); for expr in &p.expressions { - let e = consumer.consume_expression(expr, input.schema()).await?; + let e = consumer + .consume_expression(expr, input.schema(), context) + .await?; // The project's builder uses columnize_expr(..) to transform // nested window expressions into column references. window_exprs.extend(find_window_exprs([&e])); diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs index 04f310b72f4bf..b85176e4d1488 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs @@ -2,10 +2,10 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use crate::logical_plan::consumer::SubstraitConsumer; use crate::logical_plan::consumer::from_substrait_literal; use crate::logical_plan::consumer::from_substrait_named_struct; use crate::logical_plan::consumer::utils::ensure_schema_compatibility; +use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; use datafusion::common::{ DFSchema, DFSchemaRef, TableReference, not_impl_err, plan_err, substrait_datafusion_err, substrait_err, @@ -25,6 +25,7 @@ use url::Url; pub async fn from_read_rel( consumer: &impl SubstraitConsumer, read: &ReadRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { async fn read_with_schema( consumer: &impl SubstraitConsumer, @@ -32,11 +33,12 @@ pub async fn from_read_rel( schema: DFSchema, projection: Option<&MaskExpression>, filter: Option<&Expression>, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let schema = schema.replace_qualifier(table_ref.clone()); let filters = if let Some(f) = filter { - let filter_expr = consumer.consume_expression(f, &schema).await?; + let filter_expr = consumer.consume_expression(f, &schema, context).await?; split_conjunction_owned(filter_expr) } else { vec![] @@ -96,6 +98,7 @@ pub async fn from_read_rel( substrait_schema, read.projection.as_ref(), read.filter.as_deref(), + context, ) .await } @@ -158,7 +161,11 @@ pub async fn from_read_rel( } _ => { consumer - .consume_expression(expression, &substrait_schema) + .consume_expression( + expression, + &substrait_schema, + context, + ) .await? } }; @@ -218,6 +225,7 @@ pub async fn from_read_rel( substrait_schema, read.projection.as_ref(), read.filter.as_deref(), + context, ) .await } 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 13821295da320..c42b8c00f0141 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use crate::logical_plan::consumer::SubstraitConsumer; +use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; use datafusion::arrow::datatypes::Field; use datafusion::common::{JoinType, NullEquality, not_impl_err, substrait_err}; use datafusion::logical_expr::{ @@ -15,26 +15,33 @@ use substrait::proto::{Rel, SetRel}; pub async fn from_set_rel( consumer: &impl SubstraitConsumer, set: &SetRel, + context: &SubstraitConsumerContext, ) -> 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::UnionAll => union_rels(consumer, &set.inputs, true, context).await, + SetOp::UnionDistinct => { + union_rels(consumer, &set.inputs, false, context).await + } SetOp::IntersectionPrimary => intersect_rel( - consumer.consume_rel(&set.inputs[0]).await?, - union_rels(consumer, &set.inputs[1..], true).await?, + consumer.consume_rel(&set.inputs[0], context).await?, + union_rels(consumer, &set.inputs[1..], true, context).await?, false, ), SetOp::IntersectionMultiset => { - intersect_rels(consumer, &set.inputs, false).await + intersect_rels(consumer, &set.inputs, false, context).await } SetOp::IntersectionMultisetAll => { - intersect_rels(consumer, &set.inputs, true).await + intersect_rels(consumer, &set.inputs, true, context).await + } + SetOp::MinusPrimary => { + except_rels(consumer, &set.inputs, false, context).await + } + SetOp::MinusPrimaryAll => { + except_rels(consumer, &set.inputs, true, context).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:?}"), } } @@ -44,12 +51,13 @@ async fn union_rels( consumer: &impl SubstraitConsumer, rels: &[Rel], is_all: bool, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let mut union_builder = Ok(LogicalPlanBuilder::from( - consumer.consume_rel(&rels[0]).await?, + consumer.consume_rel(&rels[0], context).await?, )); for input in &rels[1..] { - let rel_plan = consumer.consume_rel(input).await?; + let rel_plan = consumer.consume_rel(input, context).await?; union_builder = if is_all { union_builder?.union(rel_plan) @@ -64,11 +72,12 @@ async fn intersect_rels( consumer: &impl SubstraitConsumer, rels: &[Rel], is_all: bool, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - let mut rel = consumer.consume_rel(&rels[0]).await?; + let mut rel = consumer.consume_rel(&rels[0], context).await?; for input in &rels[1..] { - rel = intersect_rel(rel, consumer.consume_rel(input).await?, is_all)?; + rel = intersect_rel(rel, consumer.consume_rel(input, context).await?, is_all)?; } Ok(rel) @@ -203,12 +212,16 @@ async fn except_rels( consumer: &impl SubstraitConsumer, rels: &[Rel], is_all: bool, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - let mut rel = consumer.consume_rel(&rels[0]).await?; + let mut rel = consumer.consume_rel(&rels[0], context).await?; for input in &rels[1..] { - rel = - LogicalPlanBuilder::except(rel, consumer.consume_rel(input).await?, is_all)?; + rel = LogicalPlanBuilder::except( + rel, + consumer.consume_rel(input, context).await?, + is_all, + )?; } Ok(rel) diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/sort_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/sort_rel.rs index afa9fe8465f67..b81addfbe5d14 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/sort_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/sort_rel.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements +use crate::logical_plan::consumer::SubstraitConsumerContext; + use crate::logical_plan::consumer::{SubstraitConsumer, from_substrait_sorts}; use datafusion::common::not_impl_err; use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder}; @@ -10,10 +12,12 @@ use substrait::proto::SortRel; pub async fn from_sort_rel( consumer: &impl SubstraitConsumer, sort: &SortRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { if let Some(input) = sort.input.as_ref() { - let input = LogicalPlanBuilder::from(consumer.consume_rel(input).await?); - let sorts = from_substrait_sorts(consumer, &sort.sorts, input.schema()).await?; + let input = LogicalPlanBuilder::from(consumer.consume_rel(input, context).await?); + let sorts = + from_substrait_sorts(consumer, &sort.sorts, input.schema(), context).await?; input.sort(sorts)?.build() } else { not_impl_err!("Sort without an input is not valid") diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index 94d104c158055..5ead4f42ea220 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -17,6 +17,7 @@ use async_trait::async_trait; use datafusion::arrow::datatypes::{DataType, FieldRef}; use datafusion::catalog::TableProvider; use datafusion::common::datatype::FieldExt; +use datafusion::common::extensions::Extensions as ContextExtensions; use datafusion::common::{ DFSchema, ScalarValue, TableReference, not_impl_err, substrait_err, }; @@ -24,7 +25,7 @@ use datafusion::execution::{FunctionRegistry, SessionState}; use datafusion::logical_expr::expr::LambdaVariable; use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; use std::collections::{HashMap, VecDeque}; -use std::sync::{Arc, RwLock}; +use std::sync::Arc; use substrait::proto::expression as substrait_expression; use substrait::proto::expression::{ Enum, FieldReference, IfThen, Literal, MultiOrList, Nested, ScalarFunction, @@ -41,6 +42,9 @@ use substrait::proto::{ /// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans. /// It can be implemented by users to allow for custom handling of relations, expressions, etc. /// +/// Relation and expression callbacks receive their lexical scope explicitly. +/// Inputs are consumed sequentially by default. +/// /// Combined with the [crate::logical_plan::producer::SubstraitProducer] this allows for fully /// customizable Substrait serde. /// @@ -60,14 +64,12 @@ use substrait::proto::{ /// # use datafusion::logical_expr::expr::ScalarFunction; /// # use datafusion_substrait::extensions::Extensions; /// # use datafusion_substrait::logical_plan::consumer::{ -/// # from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer, DefaultSubstraitLambdaConsumer +/// # from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer, SubstraitConsumerContext /// # }; /// /// struct CustomSubstraitConsumer { /// extensions: Arc, /// state: Arc, -/// // You can reuse existing consumer code related to lambdas -/// lambda_consumer: DefaultSubstraitLambdaConsumer, /// } /// /// #[async_trait] @@ -90,33 +92,9 @@ use substrait::proto::{ /// self.state.as_ref() /// } /// -/// fn push_lambda_parameters( -/// &self, -/// lambda_parameters: &[Type], -/// input_schema: &DFSchema, -/// ) -> datafusion::common::Result> { -/// self.lambda_consumer.push_lambda_parameters( -/// self, -/// lambda_parameters, -/// input_schema, -/// ) -/// } -/// -/// fn pop_lambda_parameters(&self) { -/// self.lambda_consumer.pop_lambda_parameters(); -/// } -/// -/// fn lambda_variable( -/// &self, -/// steps_out: usize, -/// field_idx: usize, -/// ) -> datafusion::common::Result { -/// self.lambda_consumer.lambda_variable(steps_out, field_idx) -/// } -/// /// // You can reuse existing consumer code to assist in handling advanced extensions -/// async fn consume_project(&self, rel: &ProjectRel) -> Result { -/// let df_plan = from_project_rel(self, rel).await?; +/// async fn consume_project(&self, rel: &ProjectRel, context: &SubstraitConsumerContext) -> Result { +/// let df_plan = from_project_rel(self, rel, context).await?; /// if let Some(advanced_extension) = rel.advanced_extension.as_ref() { /// not_impl_err!( /// "decode and handle an advanced extension: {:?}", @@ -128,10 +106,10 @@ use substrait::proto::{ /// } /// /// // You can implement a fully custom consumer method if you need special handling -/// async fn consume_filter(&self, rel: &FilterRel) -> Result { -/// let input = self.consume_rel(rel.input.as_ref().unwrap()).await?; +/// async fn consume_filter(&self, rel: &FilterRel, context: &SubstraitConsumerContext) -> Result { +/// let input = self.consume_rel(rel.input.as_ref().unwrap(), context).await?; /// let expression = -/// self.consume_expression(rel.condition.as_ref().unwrap(), input.schema()) +/// self.consume_expression(rel.condition.as_ref().unwrap(), input.schema(), context) /// .await?; /// // though this one is quite boring /// LogicalPlanBuilder::from(input).filter(expression)?.build() @@ -141,6 +119,7 @@ use substrait::proto::{ /// async fn consume_extension_leaf( /// &self, /// rel: &ExtensionLeafRel, +/// _context: &SubstraitConsumerContext, /// ) -> Result { /// not_impl_err!( /// "handle protobuf Any {} as you need", @@ -198,73 +177,90 @@ pub trait SubstraitConsumer: Send + Sync + Sized { /// All [Rel]s to be converted pass through this method. /// You can provide your own implementation if you wish to customize the conversion behaviour. - async fn consume_rel(&self, rel: &Rel) -> datafusion::common::Result { - from_substrait_rel(self, rel).await + async fn consume_rel( + &self, + rel: &Rel, + context: &SubstraitConsumerContext, + ) -> datafusion::common::Result { + from_substrait_rel(self, rel, context).await } async fn consume_read( &self, rel: &ReadRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_read_rel(self, rel).await + from_read_rel(self, rel, context).await } async fn consume_filter( &self, rel: &FilterRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_filter_rel(self, rel).await + from_filter_rel(self, rel, context).await } async fn consume_fetch( &self, rel: &FetchRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_fetch_rel(self, rel).await + from_fetch_rel(self, rel, context).await } async fn consume_aggregate( &self, rel: &AggregateRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_aggregate_rel(self, rel).await + from_aggregate_rel(self, rel, context).await } async fn consume_sort( &self, rel: &SortRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_sort_rel(self, rel).await + from_sort_rel(self, rel, context).await } async fn consume_join( &self, rel: &JoinRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_join_rel(self, rel).await + from_join_rel(self, rel, context).await } async fn consume_project( &self, rel: &ProjectRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_project_rel(self, rel).await + from_project_rel(self, rel, context).await } - async fn consume_set(&self, rel: &SetRel) -> datafusion::common::Result { - from_set_rel(self, rel).await + async fn consume_set( + &self, + rel: &SetRel, + context: &SubstraitConsumerContext, + ) -> datafusion::common::Result { + from_set_rel(self, rel, context).await } async fn consume_cross( &self, rel: &CrossRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_cross_rel(self, rel).await + from_cross_rel(self, rel, context).await } async fn consume_consistent_partition_window( &self, _rel: &ConsistentPartitionWindowRel, + _context: &SubstraitConsumerContext, ) -> datafusion::common::Result { not_impl_err!("Consistent Partition Window Rel not supported") } @@ -272,8 +268,9 @@ pub trait SubstraitConsumer: Send + Sync + Sized { async fn consume_exchange( &self, rel: &ExchangeRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_exchange_rel(self, rel).await + from_exchange_rel(self, rel, context).await } // Expression Methods @@ -287,11 +284,16 @@ pub trait SubstraitConsumer: Send + Sync + Sized { &self, expr: &Expression, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_substrait_rex(self, expr, input_schema).await + from_substrait_rex(self, expr, input_schema, context).await } - async fn consume_literal(&self, expr: &Literal) -> datafusion::common::Result { + async fn consume_literal( + &self, + expr: &Literal, + _context: &SubstraitConsumerContext, + ) -> datafusion::common::Result { from_literal(self, expr).await } @@ -299,38 +301,43 @@ pub trait SubstraitConsumer: Send + Sync + Sized { &self, expr: &FieldReference, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_field_reference(self, expr, input_schema).await + from_field_reference(self, expr, input_schema, context).await } async fn consume_scalar_function( &self, expr: &ScalarFunction, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_scalar_function(self, expr, input_schema).await + from_scalar_function(self, expr, input_schema, context).await } async fn consume_window_function( &self, expr: &WindowFunction, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_window_function(self, expr, input_schema).await + from_window_function(self, expr, input_schema, context).await } async fn consume_if_then( &self, expr: &IfThen, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_if_then(self, expr, input_schema).await + from_if_then(self, expr, input_schema, context).await } async fn consume_switch( &self, _expr: &SwitchExpression, _input_schema: &DFSchema, + _context: &SubstraitConsumerContext, ) -> datafusion::common::Result { not_impl_err!("Switch expression not supported") } @@ -339,14 +346,16 @@ pub trait SubstraitConsumer: Send + Sync + Sized { &self, expr: &SingularOrList, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_singular_or_list(self, expr, input_schema).await + from_singular_or_list(self, expr, input_schema, context).await } async fn consume_multi_or_list( &self, _expr: &MultiOrList, _input_schema: &DFSchema, + _context: &SubstraitConsumerContext, ) -> datafusion::common::Result { not_impl_err!("Multi Or List expression not supported") } @@ -355,30 +364,34 @@ pub trait SubstraitConsumer: Send + Sync + Sized { &self, expr: &substrait_expression::Cast, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_cast(self, expr, input_schema).await + from_cast(self, expr, input_schema, context).await } async fn consume_subquery( &self, expr: &substrait_expression::Subquery, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_subquery(self, expr, input_schema).await + from_subquery(self, expr, input_schema, context).await } async fn consume_nested( &self, expr: &Nested, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_nested(self, expr, input_schema).await + from_nested(self, expr, input_schema, context).await } async fn consume_enum( &self, _expr: &Enum, _input_schema: &DFSchema, + _context: &SubstraitConsumerContext, ) -> datafusion::common::Result { not_impl_err!("Enum expression not supported") } @@ -387,6 +400,7 @@ pub trait SubstraitConsumer: Send + Sync + Sized { &self, expr: &DynamicParameter, _input_schema: &DFSchema, + _context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let id = format!("${}", expr.parameter_reference + 1); let field = expr @@ -407,28 +421,9 @@ pub trait SubstraitConsumer: Send + Sync + Sized { &self, expr: &proto::expression::Lambda, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - from_lambda(self, expr, input_schema).await - } - - // Outer Schema Stack - // These methods manage a stack of outer schemas for correlated subquery support. - // When entering a subquery, the enclosing query's schema is pushed onto the stack. - // Field references with OuterReference root_type use these to resolve columns. - - /// Push an outer schema onto the stack when entering a subquery. - fn push_outer_schema(&self, _schema: Arc) {} - - /// Pop an outer schema from the stack when leaving a subquery. - fn pop_outer_schema(&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 (the caller should treat this as - /// an error in the Substrait plan). - fn get_outer_schema(&self, _steps_out: usize) -> Option> { - None + from_lambda(self, expr, input_schema, context).await } // User-Defined Functionality @@ -439,6 +434,7 @@ pub trait SubstraitConsumer: Send + Sync + Sized { async fn consume_extension_leaf( &self, rel: &ExtensionLeafRel, + _context: &SubstraitConsumerContext, ) -> datafusion::common::Result { if let Some(detail) = rel.detail.as_ref() { return substrait_err!( @@ -452,6 +448,7 @@ pub trait SubstraitConsumer: Send + Sync + Sized { async fn consume_extension_single( &self, rel: &ExtensionSingleRel, + _context: &SubstraitConsumerContext, ) -> datafusion::common::Result { if let Some(detail) = rel.detail.as_ref() { return substrait_err!( @@ -465,6 +462,7 @@ pub trait SubstraitConsumer: Send + Sync + Sized { async fn consume_extension_multi( &self, rel: &ExtensionMultiRel, + _context: &SubstraitConsumerContext, ) -> datafusion::common::Result { if let Some(detail) = rel.detail.as_ref() { return substrait_err!( @@ -530,35 +528,6 @@ pub trait SubstraitConsumer: Send + Sync + Sized { }; substrait_err!("Missing handler for user-defined literals {}", type_ref) } - - // Lambda related methods - - /// Push the given lambda parameters onto the stack when entering a lambda and - /// returns the names they got assigned - /// - /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaConsumer] and forward this method to it - fn push_lambda_parameters( - &self, - _lambda_parameters: &[Type], - _input_schema: &DFSchema, - ) -> datafusion::common::Result> { - not_impl_err!("SubstraitConsumer::push_lambda_parameters") - } - - /// Pop lambda parameters from the stack when leaving a lambda. - fn pop_lambda_parameters(&self) {} - - /// Returns an expression corresponding to the lambda variable with the given field_idx within the lambda it originates from, - /// at the lambda `step_outs` of the current scope - /// - /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaConsumer] and forward this method to it - fn lambda_variable( - &self, - _steps_out: usize, - _field_idx: usize, - ) -> datafusion::common::Result { - not_impl_err!("SubstraitConsumer::lambda_variable") - } } /// Default SubstraitConsumer for converting standard Substrait without user-defined extensions. @@ -567,18 +536,11 @@ pub trait SubstraitConsumer: Send + Sync + Sized { pub struct DefaultSubstraitConsumer<'a> { pub(super) extensions: &'a Extensions, pub(super) state: &'a SessionState, - outer_schemas: RwLock>>, - lambda_consumer: DefaultSubstraitLambdaConsumer, } impl<'a> DefaultSubstraitConsumer<'a> { pub fn new(extensions: &'a Extensions, state: &'a SessionState) -> Self { - DefaultSubstraitConsumer { - extensions, - state, - outer_schemas: RwLock::new(Vec::new()), - lambda_consumer: DefaultSubstraitLambdaConsumer::new(), - } + DefaultSubstraitConsumer { extensions, state } } } @@ -602,27 +564,10 @@ impl SubstraitConsumer for DefaultSubstraitConsumer<'_> { self.state } - fn push_outer_schema(&self, schema: Arc) { - self.outer_schemas.write().unwrap().push(schema); - } - - fn pop_outer_schema(&self) { - self.outer_schemas.write().unwrap().pop(); - } - - fn get_outer_schema(&self, steps_out: usize) -> Option> { - let schemas = self.outer_schemas.read().unwrap(); - // steps_out=1 → last element, steps_out=2 → second-to-last, etc. - // Returns None for steps_out=0 or steps_out > stack depth. - schemas - .len() - .checked_sub(steps_out) - .and_then(|idx| schemas.get(idx).cloned()) - } - async fn consume_extension_leaf( &self, rel: &ExtensionLeafRel, + _context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let Some(ext_detail) = &rel.detail else { return substrait_err!("Unexpected empty detail in ExtensionLeafRel"); @@ -637,6 +582,7 @@ impl SubstraitConsumer for DefaultSubstraitConsumer<'_> { async fn consume_extension_single( &self, rel: &ExtensionSingleRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let Some(ext_detail) = &rel.detail else { return substrait_err!("Unexpected empty detail in ExtensionSingleRel"); @@ -650,7 +596,7 @@ impl SubstraitConsumer for DefaultSubstraitConsumer<'_> { "ExtensionSingleRel missing input rel, try using ExtensionLeafRel instead" ); }; - let input_plan = self.consume_rel(input_rel).await?; + let input_plan = self.consume_rel(input_rel, context).await?; let plan = plan.with_exprs_and_inputs(plan.expressions(), vec![input_plan])?; Ok(LogicalPlan::Extension(Extension { node: plan })) } @@ -658,6 +604,7 @@ impl SubstraitConsumer for DefaultSubstraitConsumer<'_> { async fn consume_extension_multi( &self, rel: &ExtensionMultiRel, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let Some(ext_detail) = &rel.detail else { return substrait_err!("Unexpected empty detail in ExtensionMultiRel"); @@ -668,82 +615,81 @@ impl SubstraitConsumer for DefaultSubstraitConsumer<'_> { .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?; + let input_plan = self.consume_rel(input, context).await?; inputs.push(input_plan); } let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?; Ok(LogicalPlan::Extension(Extension { node: plan })) } - - fn push_lambda_parameters( - &self, - lambda_parameters: &[Type], - input_schema: &DFSchema, - ) -> datafusion::common::Result> { - self.lambda_consumer - .push_lambda_parameters(self, lambda_parameters, input_schema) - } - - fn pop_lambda_parameters(&self) { - self.lambda_consumer.pop_lambda_parameters() - } - - fn lambda_variable( - &self, - steps_out: usize, - field_idx: usize, - ) -> datafusion::common::Result { - self.lambda_consumer.lambda_variable(steps_out, field_idx) - } } -/// Default implementation of lambda related methods of the [SubstraitConsumer] trait +/// Immutable lexical scope for one path through a Substrait plan. /// -/// Can be embedded into a custom [SubstraitConsumer] to implement them -pub struct DefaultSubstraitLambdaConsumer { - inner: RwLock, -} - -struct DefaultSubstraitLambdaConsumerInner { - /// Parameters of the lambdas currently in scope, ordered from innermost - /// to outermost. Index 0 is the lambda being consumed; higher indices - /// are enclosing lambdas, matching the `steps_out` value used by - /// [`DefaultSubstraitLambdaConsumer::lambda_variable`] and `LambdaParameterReference`. +/// Pass the current context to child conversions. Enter a subquery or lambda by +/// deriving a child context; the parent stays unchanged, including on errors or +/// cancellation. Independent inputs may borrow the same context concurrently. +#[derive(Debug, Clone, Default)] +pub struct SubstraitConsumerContext { + outer_schemas: Vec>, lambda_parameters: VecDeque>, next_lambda_parameter: usize, + extensions: ContextExtensions, } -impl Default for DefaultSubstraitLambdaConsumer { - fn default() -> Self { - Self::new() - } -} - -impl DefaultSubstraitLambdaConsumer { - pub fn new() -> Self { - Self { - inner: RwLock::new(DefaultSubstraitLambdaConsumerInner { - lambda_parameters: VecDeque::new(), - next_lambda_parameter: 0, - }), +impl SubstraitConsumerContext { + /// Derives a child scope with an application-specific value, keyed by its type. + /// + /// Children inherit these values, including through lambda and subquery scopes. + /// Replacing a value leaves the parent and sibling scopes unchanged. Values + /// are shared through `Arc`; interior mutations remain visible to all scopes + /// sharing that value. Use a replacement value for scope-local changes. + #[must_use] + pub fn with_extension(&self, value: Arc) -> Self { + let mut child = self.clone(); + child.extensions.insert_arc(value); + child + } + + /// Returns the application-specific value of type `T` in this scope, if any. + pub fn get_extension(&self) -> Option> { + self.extensions.get_arc::() + } + + /// Derives a child scope whose immediately enclosing query has `schema`. + pub fn with_outer_schema(&self, schema: Arc) -> Self { + let mut child = self.clone(); + child.outer_schemas.push(schema); + child + } + + /// Returns the enclosing query schema; one means the immediate parent. + /// Zero or a depth beyond the enclosing queries returns `None`. + pub fn get_outer_schema(&self, steps_out: usize) -> Option> { + if steps_out == 0 { + return None; } + self.outer_schemas + .len() + .checked_sub(steps_out) + .map(|index| Arc::clone(&self.outer_schemas[index])) } - pub fn push_lambda_parameters( + /// Derives a child lambda scope and returns its assigned parameter names. + pub fn with_lambda_parameters( &self, consumer: &impl SubstraitConsumer, lambda_parameters: &[Type], input_schema: &DFSchema, - ) -> datafusion::common::Result> { - let mut inner = self.inner.write().unwrap(); + ) -> datafusion::common::Result<(Self, Vec)> { + let mut child = self.clone(); let lambda_parameters = lambda_parameters .iter() .map(|ty| { let (assigned_number, default_name) = - next_lambda_parameter_name(inner.next_lambda_parameter, input_schema); + next_lambda_parameter_name(child.next_lambda_parameter, input_schema); - inner.next_lambda_parameter = assigned_number + 1; + child.next_lambda_parameter = assigned_number + 1; Ok(field_from_substrait_type_without_names(consumer, ty)? .renamed(&default_name)) @@ -752,21 +698,18 @@ impl DefaultSubstraitLambdaConsumer { let names = lambda_parameters.iter().map(|f| f.name().clone()).collect(); - inner.lambda_parameters.push_front(lambda_parameters); + child.lambda_parameters.push_front(lambda_parameters); - Ok(names) - } - - pub fn pop_lambda_parameters(&self) { - self.inner.write().unwrap().lambda_parameters.pop_front(); + Ok((child, names)) } + /// Resolves a parameter in the current lambda (zero) or an enclosing lambda. pub fn lambda_variable( &self, steps_out: usize, field_idx: usize, ) -> datafusion::common::Result { - let lambda_parameters = &self.inner.read().unwrap().lambda_parameters; + let lambda_parameters = &self.lambda_parameters; let Some(lambda_parameters) = lambda_parameters.get(steps_out) else { return substrait_err!( @@ -815,73 +758,105 @@ mod tests { use super::*; use crate::logical_plan::consumer::utils::tests::test_consumer; use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::logical_expr::ExprSchemable; - fn make_schema(fields: &[(&str, DataType)]) -> Arc { - let arrow_fields: Vec = fields - .iter() - .map(|(name, dt)| Field::new(*name, dt.clone(), true)) - .collect(); + fn schema(name: &str, data_type: DataType) -> Arc { Arc::new( - DFSchema::try_from(Schema::new(arrow_fields)) - .expect("failed to create schema"), + DFSchema::try_from(Schema::new(vec![Field::new(name, data_type, true)])) + .unwrap(), ) } #[test] - fn test_get_outer_schema_empty_stack() { - let consumer = test_consumer(); - - // No schemas pushed — any steps_out should return None - assert!(consumer.get_outer_schema(0).is_none()); - assert!(consumer.get_outer_schema(1).is_none()); - assert!(consumer.get_outer_schema(2).is_none()); + fn extensions_inherit_and_override_without_changing_other_scopes() { + struct Scope(&'static str); + + let root = SubstraitConsumerContext::default(); + let value = Arc::new(Scope("parent")); + let parent = root.with_extension(Arc::clone(&value)); + let sibling = parent.with_outer_schema(schema("a", DataType::Int64)); + let (child, _) = parent + .with_lambda_parameters(&test_consumer(), &[], &DFSchema::empty()) + .unwrap(); + assert!(Arc::ptr_eq( + &child.get_extension::().unwrap(), + &value + )); + let child = child.with_extension(Arc::new(Scope("child"))); + + assert!(root.get_extension::().is_none()); + assert_eq!(parent.get_extension::().unwrap().0, "parent"); + assert!(Arc::ptr_eq( + &sibling.get_extension::().unwrap(), + &value + )); + assert_eq!(child.get_extension::().unwrap().0, "child"); } #[test] - fn test_get_outer_schema_single_level() { - let consumer = test_consumer(); - - let schema_a = make_schema(&[("a", DataType::Int64)]); - consumer.push_outer_schema(Arc::clone(&schema_a)); - - // steps_out=1 returns the one pushed schema - let result = consumer.get_outer_schema(1).unwrap(); - assert_eq!(result.fields().len(), 1); - assert_eq!(result.fields()[0].name(), "a"); - - // steps_out=0 and steps_out=2 are out of range - assert!(consumer.get_outer_schema(0).is_none()); - assert!(consumer.get_outer_schema(2).is_none()); - - consumer.pop_outer_schema(); - assert!(consumer.get_outer_schema(1).is_none()); + fn outer_scopes_inherit_without_changing_the_parent() { + let root = SubstraitConsumerContext::default(); + let a = schema("a", DataType::Int64); + let b = schema("b", DataType::Utf8); + let parent = root.with_outer_schema(Arc::clone(&a)); + let child = parent.with_outer_schema(Arc::clone(&b)); + + assert!(root.get_outer_schema(1).is_none()); + assert_eq!(parent.get_outer_schema(1), Some(Arc::clone(&a))); + assert!(parent.get_outer_schema(2).is_none()); + assert_eq!(child.get_outer_schema(1), Some(b)); + assert_eq!(child.get_outer_schema(2), Some(a)); + assert!(child.get_outer_schema(0).is_none()); + assert!(child.get_outer_schema(3).is_none()); } #[test] - fn test_get_outer_schema_nested() { + fn lambda_scopes_inherit_without_changing_parents_or_siblings() { let consumer = test_consumer(); - - let schema_a = make_schema(&[("a", DataType::Int64)]); - let schema_b = make_schema(&[("b", DataType::Utf8)]); - - consumer.push_outer_schema(Arc::clone(&schema_a)); - consumer.push_outer_schema(Arc::clone(&schema_b)); - - // steps_out=1 returns the most recent (schema_b) - let result = consumer.get_outer_schema(1).unwrap(); - assert_eq!(result.fields()[0].name(), "b"); - - // steps_out=2 returns the grandparent (schema_a) - let result = consumer.get_outer_schema(2).unwrap(); - assert_eq!(result.fields()[0].name(), "a"); - - // steps_out=3 exceeds depth - assert!(consumer.get_outer_schema(3).is_none()); - - // Pop one level — now steps_out=1 returns schema_a - consumer.pop_outer_schema(); - let result = consumer.get_outer_schema(1).unwrap(); - assert_eq!(result.fields()[0].name(), "a"); - assert!(consumer.get_outer_schema(2).is_none()); + let root = SubstraitConsumerContext::default(); + let schema = DFSchema::empty(); + let integer = [Type { + kind: Some(r#type::Kind::I64(r#type::I64::default())), + }]; + let string = [Type { + kind: Some(r#type::Kind::String(r#type::String::default())), + }]; + let (a, a_names) = root + .with_lambda_parameters(&consumer, &integer, &schema) + .unwrap(); + let (b, b_names) = root + .with_lambda_parameters(&consumer, &string, &schema) + .unwrap(); + let (nested, nested_names) = a + .with_lambda_parameters(&consumer, &string, &schema) + .unwrap(); + + assert_eq!(a_names, vec!["p0"]); + assert_eq!(b_names, vec!["p0"]); + assert_eq!(nested_names, vec!["p1"]); + assert_eq!( + a.lambda_variable(0, 0).unwrap().get_type(&schema).unwrap(), + DataType::Int64 + ); + assert_eq!( + b.lambda_variable(0, 0).unwrap().get_type(&schema).unwrap(), + DataType::Utf8 + ); + assert_eq!( + nested + .lambda_variable(0, 0) + .unwrap() + .get_type(&schema) + .unwrap(), + DataType::Utf8 + ); + assert_eq!( + nested.lambda_variable(1, 0).unwrap(), + a.lambda_variable(0, 0).unwrap() + ); + assert!(root.lambda_variable(0, 0).is_err()); + assert!(a.lambda_variable(1, 0).is_err()); + assert!(b.lambda_variable(1, 0).is_err()); + assert!(nested.lambda_variable(0, 1).is_err()); } } diff --git a/datafusion/substrait/src/logical_plan/consumer/utils.rs b/datafusion/substrait/src/logical_plan/consumer/utils.rs index c792e6897acd3..ec2b190a0fdab 100644 --- a/datafusion/substrait/src/logical_plan/consumer/utils.rs +++ b/datafusion/substrait/src/logical_plan/consumer/utils.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -use crate::logical_plan::consumer::SubstraitConsumer; +use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UnionFields}; use datafusion::common::{ Column, DFSchema, DFSchemaRef, TableReference, exec_err, not_impl_err, @@ -497,11 +497,12 @@ pub async fn from_substrait_sorts( consumer: &impl SubstraitConsumer, substrait_sorts: &Vec, input_schema: &DFSchema, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result> { let mut sorts: Vec = Vec::with_capacity(substrait_sorts.len()); for s in substrait_sorts { let expr = consumer - .consume_expression(s.expr.as_ref().unwrap(), input_schema) + .consume_expression(s.expr.as_ref().unwrap(), input_schema, context) .await?; let asc_nullfirst = match &s.sort_kind { Some(k) => match k { diff --git a/datafusion/substrait/tests/cases/logical_plans.rs b/datafusion/substrait/tests/cases/logical_plans.rs index c6a8275338770..1ccc57ba7cfb9 100644 --- a/datafusion/substrait/tests/cases/logical_plans.rs +++ b/datafusion/substrait/tests/cases/logical_plans.rs @@ -545,7 +545,7 @@ mod tests { assert_snapshot!( plan, @" - Projection: array_transform2(make_array(make_array(data3.p1)), (p0, p2) -> array_concat(array_transform2(p0, (p3, p4) -> p3 * p2 * p4), array_transform2(p0, (p5, p6) -> p5 * p2 * p6))) AS array_transform2(make_array(make_array(data3.p1)),(v, i) -> array_concat(array_transform2(v,(v, j) -> v * i * j),array_transform2(v,(v, j) -> v * i * j))) + Projection: array_transform2(make_array(make_array(data3.p1)), (p0, p2) -> array_concat(array_transform2(p0, (p3, p4) -> p3 * p2 * p4), array_transform2(p0, (p3, p4) -> p3 * p2 * p4))) AS array_transform2(make_array(make_array(data3.p1)),(v, i) -> array_concat(array_transform2(v,(v, j) -> v * i * j),array_transform2(v,(v, j) -> v * i * j))) TableScan: data3 " ); @@ -641,7 +641,7 @@ mod tests { assert_snapshot!( df.into_optimized_plan()?, @r#" - Projection: array_any_match(array_transform(t.tags, (p0) -> p0), (p1) -> p1 = Dictionary(UInt32, Utf8("c"))) AS matched + Projection: array_any_match(array_transform(t.tags, (p0) -> p0), (p0) -> p0 = Dictionary(UInt32, Utf8("c"))) AS matched TableScan: t projection=[tags] "# ); diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 5ebf2364ab13b..4886c51d93e53 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -2645,7 +2645,7 @@ async fn roundtrip_array_transform_higher_order_function() -> Result<()> { assert_snapshot!( plan, @" - Projection: array_transform2(make_array(make_array(data3.p1)), (p0) -> array_concat(array_transform2(p0, (p2) -> p2 * Int64(2)), array_transform2(p0, (p3) -> p3 * Int64(2)))) AS array_transform2(make_array(make_array(data3.p1)),(v) -> array_concat(array_transform2(v,(v) -> v * Int64(2)),array_transform2(v,(v) -> v * Int64(2)))) + Projection: array_transform2(make_array(make_array(data3.p1)), (p0) -> array_concat(array_transform2(p0, (p2) -> p2 * Int64(2)), array_transform2(p0, (p2) -> p2 * Int64(2)))) AS array_transform2(make_array(make_array(data3.p1)),(v) -> array_concat(array_transform2(v,(v) -> v * Int64(2)),array_transform2(v,(v) -> v * Int64(2)))) TableScan: data3 projection=[p1] " ); @@ -2698,6 +2698,29 @@ async fn roundtrip_lambda_parameter_struct_field_names() -> Result<()> { Ok(()) } +// Sibling lambda parameter names may repeat, but nested captures must still +// resolve against their own enclosing scopes during optimization and execution. +#[tokio::test] +async fn roundtrip_executes_sibling_and_nested_lambda_scopes() -> Result<()> { + let ctx = create_context_with_dialect(Some(Dialect::Databricks)).await?; + let df = ctx.sql( + "SELECT array_transform([1, 2], x -> x + 10) AS ints, + array_transform(['a', 'b'], s -> upper(s)) AS strings, + array_transform([1, 2], x -> array_transform([10, 20], y -> x + y)) AS nested", + ).await?; + let plan = df.logical_plan().clone(); + let expected = df.collect().await?; + assert_eq!( + expected.iter().map(|batch| batch.num_rows()).sum::(), + 1 + ); + let proto = to_substrait_plan(&plan, &ctx.state())?; + let consumed = from_substrait_plan(&ctx.state(), &proto).await?; + let actual = DataFrame::new(ctx.state(), consumed).collect().await?; + assert_eq!(actual, expected); + Ok(()) +} + pub(crate) async fn higher_order_function_ctx() -> Result { let ctx = create_context_with_dialect(Some(Dialect::Databricks)).await?; diff --git a/docs/source/library-user-guide/upgrading/57.0.0.md b/docs/source/library-user-guide/upgrading/57.0.0.md new file mode 100644 index 0000000000000..2f63a9d39a392 --- /dev/null +++ b/docs/source/library-user-guide/upgrading/57.0.0.md @@ -0,0 +1,40 @@ + + +# Upgrade Guides + +## DataFusion 57.0.0 + +**Note:** DataFusion `57.0.0` has not been released yet. The information provided +in this section pertains to features and changes that have already been merged +to the main branch and are awaiting release in this version. + +### Substrait consumers use explicit scope context + +Add a final `&SubstraitConsumerContext` argument to custom relation and expression +callbacks on `SubstraitConsumer`. Forward it when calling conversion helpers or +consuming child relations and expressions. + +Remove implementations of `push_outer_schema`, `pop_outer_schema`, +`get_outer_schema`, `push_lambda_parameters`, `pop_lambda_parameters`, and +`lambda_variable`, and remove embedded `DefaultSubstraitLambdaConsumer` values. +Use the context's `get_outer_schema` and `lambda_variable` methods to resolve +outer references and lambda parameters. To enter a scope, derive a child with +`context.with_outer_schema(...)` or `context.with_lambda_parameters(...)` and pass +it to nested conversions. No pop operation is needed; the parent stays unchanged. + +`from_substrait_plan` and `from_substrait_plan_with_consumer` create the root +context. For a new conversion started directly through a helper, pass +`&SubstraitConsumerContext::default()`. Nested conversions must receive the current +context. + +Custom scoped values can be stored with `context.with_extension(Arc::new(value))` +and retrieved with `context.get_extension::()`. Pass the returned child +context to nested conversions. Children inherit values; replacing a value leaves +parents and siblings unchanged. The values themselves are shared through `Arc`. + +Update logical-plan snapshots if needed: generated lambda parameter names are +local to lexical scopes, so sibling lambdas may reuse names. diff --git a/docs/source/library-user-guide/upgrading/index.rst b/docs/source/library-user-guide/upgrading/index.rst index f64662f5865d2..e3cdacc104f39 100644 --- a/docs/source/library-user-guide/upgrading/index.rst +++ b/docs/source/library-user-guide/upgrading/index.rst @@ -8,6 +8,7 @@ Upgrade Guides .. toctree:: :maxdepth: 1 + DataFusion 57.0.0 <57.0.0> DataFusion 56.0.0 <56.0.0> DataFusion 55.0.0 <55.0.0> DataFusion 54.0.0 <54.0.0>