refactor(substrait): pass immutable scope context through consumers

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<Box<Expr>>,
     order_by: Vec<SortExpr>,
     distinct: bool,
+    context: &SubstraitConsumerContext,
 ) -> datafusion::common::Result<Arc<Expr>> {
     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

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<Expr> {
     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?,
             );

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<Expr> {
-    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<Expr> {
     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<Expr> {
     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,19 +125,24 @@ 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())),
                 }],
                 DFSchema::empty_ref(),
             )
             .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(),

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<FunctionArgument>,
     input_schema: &DFSchema,
+    context: &SubstraitConsumerContext,
 ) -> datafusion::common::Result<Vec<Expr>> {
     let mut args: Vec<Expr> = 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?);

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<Expr> {
     // 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,20 +37,30 @@ 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?,
             ),
         ));
     }
     // 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,
     };

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<Expr> {
     let Some(parameters) = expr.parameters.as_ref() else {
         return substrait_err!("Lambda expression without parameters is not allowed");
     };
 
-    let names = consumer.push_lambda_parameters(&parameters.types, input_schema)?;
+    let (child_context, names) =
+        context.with_lambda_parameters(consumer, &parameters.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(),