From 7e9180a51031032125d7f7ca5c9f590f9b0c8380 Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Fri, 23 Oct 2026 15:22:59 -0400 Subject: [PATCH 1/9] feat(substrait): add a hook for consuming relation inputs --- .../src/logical_plan/consumer/rel/join_rel.rs | 25 +- .../src/logical_plan/consumer/rel/set_rel.rs | 25 +- .../consumer/substrait_consumer.rs | 225 ++++++++++++++++++ 3 files changed, 259 insertions(+), 16 deletions(-) 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..932da6951903d 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/join_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/join_rel.rs @@ -3,7 +3,10 @@ // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements use crate::logical_plan::consumer::SubstraitConsumer; -use datafusion::common::{Column, JoinType, NullEquality, not_impl_err, plan_err}; +use datafusion::common::{ + Column, JoinType, NullEquality, internal_datafusion_err, not_impl_err, plan_err, + substrait_err, +}; use datafusion::logical_expr::requalify_sides_if_needed; use datafusion::logical_expr::utils::split_conjunction_owned; use datafusion::logical_expr::{ @@ -20,12 +23,20 @@ pub async fn from_join_rel( 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?, - ); - let right = LogicalPlanBuilder::from( - consumer.consume_rel(join.right.as_ref().unwrap()).await?, - ); + let (Some(left), Some(right)) = (join.left.as_deref(), join.right.as_deref()) else { + return substrait_err!("JoinRel requires left and right inputs"); + }; + let [left, right] = consumer + .consume_rel_inputs(&[left, right]) + .await? + .try_into() + .map_err(|_| { + internal_datafusion_err!( + "JoinRel consumer returned the wrong number of inputs" + ) + })?; + let left = LogicalPlanBuilder::from(left); + let right = LogicalPlanBuilder::from(right); let (left, right, _requalified) = requalify_sides_if_needed(left, right)?; let join_type = from_substrait_jointype(join.r#type)?; 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..0fa296d7b1e86 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs @@ -4,7 +4,9 @@ use crate::logical_plan::consumer::SubstraitConsumer; use datafusion::arrow::datatypes::Field; -use datafusion::common::{JoinType, NullEquality, not_impl_err, substrait_err}; +use datafusion::common::{ + JoinType, NullEquality, internal_err, not_impl_err, substrait_err, +}; use datafusion::logical_expr::{ Cast, Expr, LogicalPlan, LogicalPlanBuilder, requalify_sides_if_needed, }; @@ -45,16 +47,21 @@ async fn union_rels( rels: &[Rel], is_all: bool, ) -> datafusion::common::Result { - let mut union_builder = Ok(LogicalPlanBuilder::from( - consumer.consume_rel(&rels[0]).await?, - )); - for input in &rels[1..] { - let rel_plan = consumer.consume_rel(input).await?; - + let inputs = rels.iter().collect::>(); + let plans = consumer.consume_rel_inputs(&inputs).await?; + if plans.len() != rels.len() { + return internal_err!("SetRel consumer returned the wrong number of inputs"); + } + let mut plans = plans.into_iter(); + let Some(first) = plans.next() else { + return substrait_err!("Set operation requires at least two inputs"); + }; + let mut union_builder = Ok(LogicalPlanBuilder::from(first)); + for plan in plans { union_builder = if is_all { - union_builder?.union(rel_plan) + union_builder?.union(plan) } else { - union_builder?.union_distinct(rel_plan) + union_builder?.union_distinct(plan) }; } union_builder?.build() diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index 94d104c158055..46816e13144db 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -202,6 +202,24 @@ pub trait SubstraitConsumer: Send + Sync + Sized { from_substrait_rel(self, rel).await } + /// Consumes independent inputs of a JOIN or UNION, returning plans in input order. + /// + /// Implementations must return exactly one plan per input, in the same order. + /// On error, implementations must drop any unfinished input futures. + /// + /// The default preserves sequential consumption. Implementations may override it + /// when they can isolate mutable traversal state between inputs. + async fn consume_rel_inputs( + &self, + rels: &[&Rel], + ) -> datafusion::common::Result> { + let mut plans = Vec::with_capacity(rels.len()); + for rel in rels { + plans.push(self.consume_rel(rel).await?); + } + Ok(plans) + } + async fn consume_read( &self, rel: &ReadRel, @@ -884,4 +902,211 @@ mod tests { assert_eq!(result.fields()[0].name(), "a"); assert!(consumer.get_outer_schema(2).is_none()); } + + // A custom consumer exercises the trait default independently of the default + // consumer's scheduling policy. + struct InputConsumer { + extensions: Extensions, + state: SessionState, + reads: std::sync::Mutex>, + ready: tokio::sync::Semaphore, + } + + #[async_trait] + impl SubstraitConsumer for InputConsumer { + fn get_extensions(&self) -> &Extensions { + &self.extensions + } + + fn get_function_registry(&self) -> &impl FunctionRegistry { + &self.state + } + + async fn resolve_table_ref( + &self, + _: &TableReference, + ) -> datafusion::common::Result>> { + unreachable!("test consumes reads directly") + } + + async fn consume_rel( + &self, + rel: &Rel, + ) -> datafusion::common::Result { + let Some(proto::rel::RelType::Read(read)) = &rel.rel_type else { + panic!("expected a read") + }; + let Some(proto::read_rel::ReadType::NamedTable(table)) = &read.read_type + else { + panic!("expected a named table") + }; + self.reads.lock().unwrap().push(table.names[0].clone()); + self.ready.acquire().await.unwrap().forget(); + if table.names[0] == "fail" { + return substrait_err!("input failed"); + } + let provider = + datafusion::catalog::empty::EmptyTable::new(Arc::new(Schema::new(vec![ + Field::new("x", DataType::Int64, false), + ]))); + datafusion::logical_expr::LogicalPlanBuilder::scan( + table.names[0].clone(), + datafusion::datasource::provider_as_source(Arc::new(provider)), + None, + )? + .build() + } + } + + fn input_consumer() -> InputConsumer { + InputConsumer { + extensions: Extensions::default(), + state: datafusion::execution::SessionStateBuilder::new_with_default_features( + ) + .build(), + reads: std::sync::Mutex::default(), + ready: tokio::sync::Semaphore::new(0), + } + } + + fn input_read(name: &str) -> Rel { + Rel { + rel_type: Some(proto::rel::RelType::Read(Box::new(ReadRel { + read_type: Some(proto::read_rel::ReadType::NamedTable( + proto::read_rel::NamedTable { + names: vec![name.into()], + ..Default::default() + }, + )), + ..Default::default() + }))), + } + } + + #[tokio::test] + async fn input_hook_default_is_sequential_and_stops_on_error() { + use std::future::{Future, poll_fn}; + use std::task::Poll; + let consumer = input_consumer(); + let rels = [input_read("a"), input_read("fail"), input_read("unreached")]; + let inputs = rels.iter().collect::>(); + let mut pending = Box::pin(consumer.consume_rel_inputs(&inputs)); + poll_fn(|cx| { + assert!(pending.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + assert_eq!(*consumer.reads.lock().unwrap(), ["a"]); + consumer.ready.add_permits(3); + assert!( + pending + .await + .unwrap_err() + .to_string() + .contains("input failed") + ); + assert_eq!(*consumer.reads.lock().unwrap(), ["a", "fail"]); + } + + struct HookConsumer { + inner: InputConsumer, + calls: std::sync::Mutex>, + truncate: bool, + } + + #[async_trait] + impl SubstraitConsumer for HookConsumer { + fn get_extensions(&self) -> &Extensions { + &self.inner.extensions + } + fn get_function_registry(&self) -> &impl FunctionRegistry { + &self.inner.state + } + async fn resolve_table_ref( + &self, + _: &TableReference, + ) -> datafusion::common::Result>> { + unreachable!("test consumes reads directly") + } + async fn consume_rel_inputs( + &self, + rels: &[&Rel], + ) -> datafusion::common::Result> { + self.calls.lock().unwrap().push(rels.len()); + let mut plans = self.inner.consume_rel_inputs(rels).await?; + if self.truncate { + plans.pop(); + } + Ok(plans) + } + } + + #[tokio::test] + async fn join_and_union_use_the_input_hook() { + for truncate in [false, true] { + let consumer = HookConsumer { + inner: input_consumer(), + calls: std::sync::Mutex::default(), + truncate, + }; + consumer.inner.ready.add_permits(6); + let join = JoinRel { + left: Some(Box::new(input_read("a"))), + right: Some(Box::new(input_read("b"))), + r#type: proto::join_rel::JoinType::Inner as i32, + ..Default::default() + }; + let result = consumer.consume_join(&join).await; + if truncate { + assert!( + result + .unwrap_err() + .to_string() + .contains("wrong number of inputs") + ); + } else { + let LogicalPlan::Join(join) = result.unwrap() else { + panic!("expected join") + }; + assert!( + join.left + .display_indent() + .to_string() + .contains("TableScan: a") + ); + assert!( + join.right + .display_indent() + .to_string() + .contains("TableScan: b") + ); + } + for op in [ + proto::set_rel::SetOp::UnionAll, + proto::set_rel::SetOp::UnionDistinct, + ] { + let set = SetRel { + inputs: vec![input_read("a"), input_read("b")], + op: op as i32, + ..Default::default() + }; + let result = consumer.consume_set(&set).await; + if truncate { + assert!( + result + .unwrap_err() + .to_string() + .contains("wrong number of inputs") + ); + } else { + let plan = result.unwrap().display_indent().to_string(); + assert!( + plan.find("TableScan: a").unwrap() + < plan.find("TableScan: b").unwrap() + ); + } + } + assert_eq!(*consumer.calls.lock().unwrap(), [2, 2, 2]); + } + } } From 41a1f3654c6f0fd6e755b1265309ba1951e146b1 Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Thu, 29 Oct 2026 12:00:27 -0400 Subject: [PATCH 2/9] refactor(substrait): pass immutable scope context through consumers --- Cargo.lock | 1 + datafusion/substrait/Cargo.toml | 1 + datafusion/substrait/README.md | 27 + .../consumer/expr/aggregate_function.rs | 6 +- .../src/logical_plan/consumer/expr/cast.rs | 4 + .../consumer/expr/field_reference.rs | 49 +- .../consumer/expr/function_arguments.rs | 7 +- .../src/logical_plan/consumer/expr/if_then.rs | 25 +- .../src/logical_plan/consumer/expr/lambda.rs | 35 +- .../src/logical_plan/consumer/expr/mod.rs | 96 ++- .../src/logical_plan/consumer/expr/nested.rs | 33 +- .../consumer/expr/scalar_function.rs | 35 +- .../consumer/expr/singular_or_list.rs | 8 +- .../logical_plan/consumer/expr/subquery.rs | 37 +- .../consumer/expr/window_function.rs | 9 +- .../src/logical_plan/consumer/plan.rs | 14 +- .../consumer/rel/aggregate_rel.rs | 24 +- .../logical_plan/consumer/rel/cross_rel.rs | 11 +- .../logical_plan/consumer/rel/exchange_rel.rs | 13 +- .../logical_plan/consumer/rel/fetch_rel.rs | 21 +- .../logical_plan/consumer/rel/filter_rel.rs | 7 +- .../src/logical_plan/consumer/rel/join_rel.rs | 9 +- .../src/logical_plan/consumer/rel/mod.rs | 39 +- .../logical_plan/consumer/rel/project_rel.rs | 9 +- .../src/logical_plan/consumer/rel/read_rel.rs | 14 +- .../src/logical_plan/consumer/rel/set_rel.rs | 43 +- .../src/logical_plan/consumer/rel/sort_rel.rs | 8 +- .../consumer/substrait_consumer.rs | 614 +++++++++++------- .../src/logical_plan/consumer/utils.rs | 5 +- .../substrait/tests/cases/logical_plans.rs | 4 +- .../tests/cases/roundtrip_logical_plan.rs | 25 +- 31 files changed, 825 insertions(+), 408 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e3bfa7f5959b1..4ed040839f7a4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2792,6 +2792,7 @@ dependencies = [ "chrono", "datafusion", "datafusion-functions-aggregate", + "futures", "half", "insta", "itertools 0.15.0", diff --git a/datafusion/substrait/Cargo.toml b/datafusion/substrait/Cargo.toml index d6fc6daa9e6d6..4aa8868c17371 100644 --- a/datafusion/substrait/Cargo.toml +++ b/datafusion/substrait/Cargo.toml @@ -21,6 +21,7 @@ rust-version = { workspace = true } workspace = true [dependencies] +futures = { workspace = true } async-recursion = "1.2" async-trait = { workspace = true } chrono = { workspace = true } diff --git a/datafusion/substrait/README.md b/datafusion/substrait/README.md index 7feeafbeeacb1..5f6cdaefe63fd 100644 --- a/datafusion/substrait/README.md +++ b/datafusion/substrait/README.md @@ -15,3 +15,30 @@ plans. See [API Docs] for details and examples. [apache datafusion]: https://datafusion.apache.org/ [substrait]: https://substrait.io [api docs]: https://docs.rs/datafusion-substrait/latest + +## Migrating custom consumers to explicit scope context + +Relation and expression callbacks on `SubstraitConsumer` now take a final +`&SubstraitConsumerContext` argument. Forward that context when delegating to the +standard conversion helpers or consuming child relations and expressions. +`from_substrait_plan` and `from_substrait_plan_with_consumer` still create the root +context for you. Direct calls to conversion helpers should start with +`SubstraitConsumerContext::default()`. + +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. +The standard handlers resolve scopes through the context. Custom handlers that +introduce scopes should use `context.with_outer_schema(...)` or +`context.with_lambda_parameters(...)` and pass the resulting child context to +nested conversions. There is no matching pop operation; the parent context stays +unchanged on success, error, and cancellation. + +`consume_rel_inputs` now consumes JOIN and UNION inputs concurrently on the same +consumer. Keep traversal scopes in the context. Shared services and caches must +support concurrent calls. Consumers requiring sequential input processing can +override `consume_rel_inputs` and await each `consume_rel(rel, context)` in order. + +Generated lambda parameter names are now local to lexical scopes, so sibling +lambdas may reuse names. Logical plan displays can change without changing query +results. 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 932da6951903d..c4b3354eeb2e7 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, internal_datafusion_err, not_impl_err, plan_err, substrait_err, @@ -18,6 +18,7 @@ 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"); @@ -27,7 +28,7 @@ pub async fn from_join_rel( return substrait_err!("JoinRel requires left and right inputs"); }; let [left, right] = consumer - .consume_rel_inputs(&[left, right]) + .consume_rel_inputs(&[left, right], context) .await? .try_into() .map_err(|_| { @@ -48,7 +49,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 0fa296d7b1e86..0311d8a0fc796 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, internal_err, not_impl_err, substrait_err, @@ -17,26 +17,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:?}"), } } @@ -46,9 +53,10 @@ async fn union_rels( consumer: &impl SubstraitConsumer, rels: &[Rel], is_all: bool, + context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let inputs = rels.iter().collect::>(); - let plans = consumer.consume_rel_inputs(&inputs).await?; + let plans = consumer.consume_rel_inputs(&inputs, context).await?; if plans.len() != rels.len() { return internal_err!("SetRel consumer returned the wrong number of inputs"); } @@ -71,11 +79,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) @@ -210,12 +219,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 46816e13144db..ad67f731a83fc 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -23,8 +23,10 @@ use datafusion::common::{ use datafusion::execution::{FunctionRegistry, SessionState}; use datafusion::logical_expr::expr::LambdaVariable; use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; +use futures::TryStreamExt; +use futures::stream::FuturesUnordered; 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 +43,11 @@ 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. +/// Independent inputs may invoke callbacks concurrently on the same consumer. +/// Keep traversal scopes in the context and make shared services concurrency-safe, +/// or override [`Self::consume_rel_inputs`] to consume inputs sequentially. +/// /// Combined with the [crate::logical_plan::producer::SubstraitProducer] this allows for fully /// customizable Substrait serde. /// @@ -60,14 +67,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 +95,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 +109,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 +122,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,8 +180,12 @@ 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 } /// Consumes independent inputs of a JOIN or UNION, returning plans in input order. @@ -207,82 +193,102 @@ pub trait SubstraitConsumer: Send + Sync + Sized { /// Implementations must return exactly one plan per input, in the same order. /// On error, implementations must drop any unfinished input futures. /// - /// The default preserves sequential consumption. Implementations may override it - /// when they can isolate mutable traversal state between inputs. + /// The default consumes inputs concurrently using the same immutable context. + /// Overrides can provide a different scheduling policy. async fn consume_rel_inputs( &self, rels: &[&Rel], + context: &SubstraitConsumerContext, ) -> datafusion::common::Result> { - let mut plans = Vec::with_capacity(rels.len()); - for rel in rels { - plans.push(self.consume_rel(rel).await?); + let pending = FuturesUnordered::new(); + for (index, rel) in rels.iter().enumerate() { + pending.push(async move { + self.consume_rel(rel, context) + .await + .map(|plan| (index, plan)) + }); } - Ok(plans) + let mut completed = pending.try_collect::>().await?; + completed.sort_unstable_by_key(|(index, _)| *index); + Ok(completed.into_iter().map(|(_, plan)| plan).collect()) } 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") } @@ -290,8 +296,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 @@ -305,11 +312,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 } @@ -317,38 +329,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") } @@ -357,14 +374,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") } @@ -373,30 +392,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") } @@ -405,6 +428,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 @@ -425,28 +449,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 @@ -457,6 +462,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!( @@ -470,6 +476,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!( @@ -483,6 +490,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!( @@ -548,35 +556,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. @@ -585,18 +564,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 } } } @@ -620,27 +592,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"); @@ -655,6 +610,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"); @@ -668,7 +624,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 })) } @@ -676,6 +632,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"); @@ -686,82 +643,62 @@ 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, } -impl Default for DefaultSubstraitLambdaConsumer { - fn default() -> Self { - Self::new() +impl SubstraitConsumerContext { + /// 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 } -} -impl DefaultSubstraitLambdaConsumer { - pub fn new() -> Self { - Self { - inner: RwLock::new(DefaultSubstraitLambdaConsumerInner { - lambda_parameters: VecDeque::new(), - next_lambda_parameter: 0, - }), + /// 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)) @@ -770,21 +707,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!( @@ -831,9 +765,22 @@ fn next_lambda_parameter_name( #[cfg(test)] mod tests { use super::*; - use crate::logical_plan::consumer::utils::tests::test_consumer; use datafusion::arrow::datatypes::{Field, Schema}; + use futures::poll; + use proto::expression::{ReferenceSegment, field_reference, reference_segment}; + use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; + use tokio::sync::Semaphore; + + struct ActiveInput<'a>(&'a AtomicUsize); + + impl Drop for ActiveInput<'_> { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::SeqCst); + } + } + fn make_schema(fields: &[(&str, DataType)]) -> Arc { let arrow_fields: Vec = fields .iter() @@ -847,69 +794,73 @@ mod tests { #[test] fn test_get_outer_schema_empty_stack() { - let consumer = test_consumer(); + let context = SubstraitConsumerContext::default(); // 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()); + assert!(context.get_outer_schema(0).is_none()); + assert!(context.get_outer_schema(1).is_none()); + assert!(context.get_outer_schema(2).is_none()); } #[test] fn test_get_outer_schema_single_level() { - let consumer = test_consumer(); + let context = SubstraitConsumerContext::default(); let schema_a = make_schema(&[("a", DataType::Int64)]); - consumer.push_outer_schema(Arc::clone(&schema_a)); + let parent = context.clone(); + let context = context.with_outer_schema(Arc::clone(&schema_a)); // steps_out=1 returns the one pushed schema - let result = consumer.get_outer_schema(1).unwrap(); + let result = context.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()); + assert!(context.get_outer_schema(0).is_none()); + assert!(context.get_outer_schema(2).is_none()); - consumer.pop_outer_schema(); - assert!(consumer.get_outer_schema(1).is_none()); + let context = parent; + assert!(context.get_outer_schema(1).is_none()); } #[test] fn test_get_outer_schema_nested() { - let consumer = test_consumer(); + let context = SubstraitConsumerContext::default(); 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)); + let context = context.with_outer_schema(Arc::clone(&schema_a)); + let parent = context.clone(); + let context = context.with_outer_schema(Arc::clone(&schema_b)); // steps_out=1 returns the most recent (schema_b) - let result = consumer.get_outer_schema(1).unwrap(); + let result = context.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(); + let result = context.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()); + assert!(context.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(); + // The parent is unchanged — steps_out=1 still returns schema_a + let context = parent; + let result = context.get_outer_schema(1).unwrap(); assert_eq!(result.fields()[0].name(), "a"); - assert!(consumer.get_outer_schema(2).is_none()); + assert!(context.get_outer_schema(2).is_none()); } // A custom consumer exercises the trait default independently of the default // consumer's scheduling policy. struct InputConsumer { - extensions: Extensions, - state: SessionState, - reads: std::sync::Mutex>, - ready: tokio::sync::Semaphore, + extensions: Arc, + state: Arc, + reads: Arc>>, + ready: Arc, + gates: Arc>, + active: Arc, } #[async_trait] @@ -919,7 +870,7 @@ mod tests { } fn get_function_registry(&self) -> &impl FunctionRegistry { - &self.state + self.state.as_ref() } async fn resolve_table_ref( @@ -929,9 +880,22 @@ mod tests { unreachable!("test consumes reads directly") } + async fn consume_expression( + &self, + expr: &Expression, + input_schema: &DFSchema, + context: &SubstraitConsumerContext, + ) -> datafusion::common::Result { + self.active.fetch_add(1, Ordering::SeqCst); + let _active = ActiveInput(&self.active); + self.ready.acquire().await.unwrap().forget(); + from_substrait_rex(self, expr, input_schema, context).await + } + async fn consume_rel( &self, rel: &Rel, + _context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let Some(proto::rel::RelType::Read(read)) = &rel.rel_type else { panic!("expected a read") @@ -941,7 +905,10 @@ mod tests { panic!("expected a named table") }; self.reads.lock().unwrap().push(table.names[0].clone()); - self.ready.acquire().await.unwrap().forget(); + self.active.fetch_add(1, Ordering::SeqCst); + let _active = ActiveInput(&self.active); + let gate = self.gates.get(&table.names[0]).unwrap_or(&self.ready); + gate.acquire().await.unwrap().forget(); if table.names[0] == "fail" { return substrait_err!("input failed"); } @@ -960,12 +927,15 @@ mod tests { fn input_consumer() -> InputConsumer { InputConsumer { - extensions: Extensions::default(), - state: datafusion::execution::SessionStateBuilder::new_with_default_features( - ) - .build(), - reads: std::sync::Mutex::default(), - ready: tokio::sync::Semaphore::new(0), + extensions: Arc::new(Extensions::default()), + state: Arc::new( + datafusion::execution::SessionStateBuilder::new_with_default_features() + .build(), + ), + reads: Arc::new(Mutex::default()), + ready: Arc::new(Semaphore::new(0)), + gates: Arc::new(HashMap::new()), + active: Arc::new(AtomicUsize::new(0)), } } @@ -983,34 +953,169 @@ mod tests { } } + fn identity_lambda(kind: r#type::Kind, steps_out: u32) -> proto::expression::Lambda { + proto::expression::Lambda { + parameters: Some(r#type::Struct { + types: vec![Type { kind: Some(kind) }], + ..Default::default() + }), + body: Some(Box::new(Expression { + rex_type: Some(proto::expression::RexType::Selection(Box::new( + FieldReference { + reference_type: Some( + field_reference::ReferenceType::DirectReference( + ReferenceSegment { + reference_type: Some( + reference_segment::ReferenceType::StructField( + Box::new(reference_segment::StructField { + field: 0, + child: None, + }), + ), + ), + }, + ), + ), + root_type: Some( + field_reference::RootType::LambdaParameterReference( + field_reference::LambdaParameterReference { steps_out }, + ), + ), + }, + ))), + })), + } + } + + fn assert_lambda_body(expr: Expr, name: &str, data_type: DataType) { + let Expr::Lambda(lambda) = expr else { + panic!("expected lambda") + }; + let Expr::LambdaVariable(variable) = *lambda.body else { + panic!("expected lambda variable") + }; + assert_eq!(variable.name, name); + assert_eq!(variable.field.unwrap().data_type(), &data_type); + } + #[tokio::test] - async fn input_hook_default_is_sequential_and_stops_on_error() { - use std::future::{Future, poll_fn}; - use std::task::Poll; + async fn shared_consumer_keeps_concurrent_lambda_scopes_independent() { let consumer = input_consumer(); - let rels = [input_read("a"), input_read("fail"), input_read("unreached")]; - let inputs = rels.iter().collect::>(); - let mut pending = Box::pin(consumer.consume_rel_inputs(&inputs)); - poll_fn(|cx| { - assert!(pending.as_mut().poll(cx).is_pending()); - Poll::Ready(()) - }) - .await; - assert_eq!(*consumer.reads.lock().unwrap(), ["a"]); - consumer.ready.add_permits(3); + let context = SubstraitConsumerContext::default(); + let schema = DFSchema::empty(); + let integer = identity_lambda(r#type::Kind::I64(r#type::I64::default()), 0); + let string = identity_lambda(r#type::Kind::String(r#type::String::default()), 0); + let mut first = Box::pin(consumer.consume_lambda(&integer, &schema, &context)); + let mut second = Box::pin(consumer.consume_lambda(&string, &schema, &context)); + assert!(poll!(first.as_mut()).is_pending()); + assert!(poll!(second.as_mut()).is_pending()); + assert_eq!(consumer.active.load(Ordering::SeqCst), 2); + consumer.ready.add_permits(2); + assert_lambda_body(first.await.unwrap(), "p0", DataType::Int64); + assert_lambda_body(second.await.unwrap(), "p0", DataType::Utf8); + assert!(context.lambda_variable(0, 0).is_err()); + } + + #[tokio::test] + async fn lambda_error_and_cancellation_preserve_parent_scope() { + let consumer = input_consumer(); + let schema = DFSchema::empty(); + let root = SubstraitConsumerContext::default(); + let (context, names) = root + .with_lambda_parameters( + &consumer, + &[Type { + kind: Some(r#type::Kind::I64(r#type::I64::default())), + }], + &schema, + ) + .unwrap(); + assert_eq!(names, vec!["p0"]); + let nested = identity_lambda(r#type::Kind::String(r#type::String::default()), 1); + let mut cancelled = Box::pin(consumer.consume_lambda(&nested, &schema, &context)); + assert!(poll!(cancelled.as_mut()).is_pending()); + drop(cancelled); + assert_eq!(consumer.active.load(Ordering::SeqCst), 0); + let mut invalid = nested.clone(); + invalid.body = None; assert!( - pending + consumer + .consume_lambda(&invalid, &schema, &context) .await - .unwrap_err() + .is_err() + ); + consumer.ready.add_permits(1); + let result = consumer + .consume_lambda(&nested, &schema, &context) + .await + .unwrap(); + let Expr::Lambda(ref lambda) = result else { + panic!("expected lambda") + }; + assert_eq!(lambda.params, vec!["p1"]); + assert_lambda_body(result, "p0", DataType::Int64); + assert!(context.lambda_variable(1, 0).is_err()); + assert!(root.lambda_variable(0, 0).is_err()); + } + + #[tokio::test] + async fn concurrent_inputs_overlap_and_preserve_input_order() { + let mut consumer = input_consumer(); + consumer.gates = Arc::new(HashMap::from([ + ("a".into(), Semaphore::new(0)), + ("b".into(), Semaphore::new(0)), + ])); + let rels = [input_read("a"), input_read("b")]; + let inputs = rels.iter().collect::>(); + let context = SubstraitConsumerContext::default(); + let mut pending = Box::pin(consumer.consume_rel_inputs(&inputs, &context)); + assert!(poll!(&mut pending).is_pending()); + assert_eq!(*consumer.reads.lock().unwrap(), ["a", "b"]); + assert_eq!(consumer.active.load(Ordering::SeqCst), 2); + consumer.gates["b"].add_permits(1); + assert!(poll!(&mut pending).is_pending()); + assert_eq!(consumer.active.load(Ordering::SeqCst), 1); + consumer.gates["a"].add_permits(1); + let plans = pending.await.unwrap(); + assert_eq!(plans.len(), 2); + assert!( + plans[0] + .display_indent() + .to_string() + .contains("TableScan: a") + ); + assert!( + plans[1] + .display_indent() .to_string() - .contains("input failed") + .contains("TableScan: b") ); - assert_eq!(*consumer.reads.lock().unwrap(), ["a", "fail"]); + assert_eq!(consumer.active.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn concurrent_inputs_cancel_pending_siblings_on_error() { + let mut consumer = input_consumer(); + consumer.gates = Arc::new(HashMap::from([("fail".into(), Semaphore::new(0))])); + // Exercise more inputs than the small-future optimization in try_join_all. + let mut rels = (0..40) + .map(|i| input_read(&format!("a{i}"))) + .collect::>(); + rels.push(input_read("fail")); + let inputs = rels.iter().collect::>(); + let context = SubstraitConsumerContext::default(); + let mut pending = Box::pin(consumer.consume_rel_inputs(&inputs, &context)); + assert!(poll!(&mut pending).is_pending()); + assert_eq!(consumer.active.load(Ordering::SeqCst), 41); + consumer.gates["fail"].add_permits(1); + let error = pending.await.unwrap_err(); + assert!(error.to_string().contains("input failed")); + assert_eq!(consumer.active.load(Ordering::SeqCst), 0); } struct HookConsumer { inner: InputConsumer, - calls: std::sync::Mutex>, + calls: Mutex>, truncate: bool, } @@ -1020,7 +1125,7 @@ mod tests { &self.inner.extensions } fn get_function_registry(&self) -> &impl FunctionRegistry { - &self.inner.state + self.inner.state.as_ref() } async fn resolve_table_ref( &self, @@ -1031,9 +1136,10 @@ mod tests { async fn consume_rel_inputs( &self, rels: &[&Rel], + context: &SubstraitConsumerContext, ) -> datafusion::common::Result> { self.calls.lock().unwrap().push(rels.len()); - let mut plans = self.inner.consume_rel_inputs(rels).await?; + let mut plans = self.inner.consume_rel_inputs(rels, context).await?; if self.truncate { plans.pop(); } @@ -1046,7 +1152,7 @@ mod tests { for truncate in [false, true] { let consumer = HookConsumer { inner: input_consumer(), - calls: std::sync::Mutex::default(), + calls: Mutex::default(), truncate, }; consumer.inner.ready.add_permits(6); @@ -1056,7 +1162,9 @@ mod tests { r#type: proto::join_rel::JoinType::Inner as i32, ..Default::default() }; - let result = consumer.consume_join(&join).await; + let result = consumer + .consume_join(&join, &SubstraitConsumerContext::default()) + .await; if truncate { assert!( result @@ -1090,7 +1198,9 @@ mod tests { op: op as i32, ..Default::default() }; - let result = consumer.consume_set(&set).await; + let result = consumer + .consume_set(&set, &SubstraitConsumerContext::default()) + .await; if truncate { assert!( result 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?; From 0e977f002db7b23d65c98fc4ef3077889093e57e Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Thu, 29 Oct 2026 13:34:38 -0400 Subject: [PATCH 3/9] test(substrait): use table_scan in consumer test fixture --- .../src/logical_plan/consumer/substrait_consumer.rs | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index ad67f731a83fc..ef8ca55f71578 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -766,6 +766,7 @@ fn next_lambda_parameter_name( mod tests { use super::*; use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::logical_expr::logical_plan::builder::table_scan; use futures::poll; use proto::expression::{ReferenceSegment, field_reference, reference_segment}; @@ -912,13 +913,9 @@ mod tests { if table.names[0] == "fail" { return substrait_err!("input failed"); } - let provider = - datafusion::catalog::empty::EmptyTable::new(Arc::new(Schema::new(vec![ - Field::new("x", DataType::Int64, false), - ]))); - datafusion::logical_expr::LogicalPlanBuilder::scan( - table.names[0].clone(), - datafusion::datasource::provider_as_source(Arc::new(provider)), + table_scan( + Some(table.names[0].as_str()), + &Schema::new(vec![Field::new("x", DataType::Int64, false)]), None, )? .build() From 99ae0a9a26cd86d2d76c33cfd6c4ebb16cf6f879 Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Thu, 29 Oct 2026 14:11:34 -0400 Subject: [PATCH 4/9] test(substrait): simplify scope and concurrent input tests --- .../consumer/substrait_consumer.rs | 530 +++++------------- 1 file changed, 149 insertions(+), 381 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index ef8ca55f71578..4ae3092022067 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -765,156 +765,148 @@ fn next_lambda_parameter_name( #[cfg(test)] mod tests { use super::*; + use crate::logical_plan::consumer::utils::tests::test_consumer; use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::common::assert_contains; + use datafusion::logical_expr::ExprSchemable; use datafusion::logical_expr::logical_plan::builder::table_scan; - use futures::poll; - use proto::expression::{ReferenceSegment, field_reference, reference_segment}; - use std::sync::Mutex; - use std::sync::atomic::{AtomicUsize, Ordering}; use tokio::sync::Semaphore; - struct ActiveInput<'a>(&'a AtomicUsize); - - impl Drop for ActiveInput<'_> { - fn drop(&mut self) { - self.0.fetch_sub(1, Ordering::SeqCst); - } - } - - 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 context = SubstraitConsumerContext::default(); - - // No schemas pushed — any steps_out should return None - assert!(context.get_outer_schema(0).is_none()); - assert!(context.get_outer_schema(1).is_none()); - assert!(context.get_outer_schema(2).is_none()); - } - - #[test] - fn test_get_outer_schema_single_level() { - let context = SubstraitConsumerContext::default(); - - let schema_a = make_schema(&[("a", DataType::Int64)]); - let parent = context.clone(); - let context = context.with_outer_schema(Arc::clone(&schema_a)); - - // steps_out=1 returns the one pushed schema - let result = context.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!(context.get_outer_schema(0).is_none()); - assert!(context.get_outer_schema(2).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)); - let context = parent; - assert!(context.get_outer_schema(1).is_none()); + 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() { - let context = SubstraitConsumerContext::default(); - - let schema_a = make_schema(&[("a", DataType::Int64)]); - let schema_b = make_schema(&[("b", DataType::Utf8)]); - - let context = context.with_outer_schema(Arc::clone(&schema_a)); - let parent = context.clone(); - let context = context.with_outer_schema(Arc::clone(&schema_b)); - - // steps_out=1 returns the most recent (schema_b) - let result = context.get_outer_schema(1).unwrap(); - assert_eq!(result.fields()[0].name(), "b"); - - // steps_out=2 returns the grandparent (schema_a) - let result = context.get_outer_schema(2).unwrap(); - assert_eq!(result.fields()[0].name(), "a"); - - // steps_out=3 exceeds depth - assert!(context.get_outer_schema(3).is_none()); - - // The parent is unchanged — steps_out=1 still returns schema_a - let context = parent; - let result = context.get_outer_schema(1).unwrap(); - assert_eq!(result.fields()[0].name(), "a"); - assert!(context.get_outer_schema(2).is_none()); - } + fn lambda_scopes_inherit_without_changing_parents_or_siblings() { + let consumer = test_consumer(); + 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(); - // A custom consumer exercises the trait default independently of the default - // consumer's scheduling policy. - struct InputConsumer { - extensions: Arc, - state: Arc, - reads: Arc>>, - ready: Arc, - gates: Arc>, - active: Arc, + 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()); + } + + // Only input completion is controlled; JOIN/UNION conversion and batch + // scheduling use the standard consumer implementations. + struct GatedConsumer { + consumer: DefaultSubstraitConsumer<'static>, + a_ready: Semaphore, + b_ready: Semaphore, + // Held permits count pending inputs and are released even on cancellation. + pending: Semaphore, + } + + impl GatedConsumer { + fn new(inputs: usize) -> Self { + Self { + consumer: test_consumer(), + a_ready: Semaphore::new(0), + b_ready: Semaphore::new(0), + pending: Semaphore::new(inputs), + } + } } #[async_trait] - impl SubstraitConsumer for InputConsumer { + impl SubstraitConsumer for GatedConsumer { fn get_extensions(&self) -> &Extensions { - &self.extensions + self.consumer.get_extensions() } - fn get_function_registry(&self) -> &impl FunctionRegistry { - self.state.as_ref() + self.consumer.get_function_registry() } - async fn resolve_table_ref( &self, _: &TableReference, ) -> datafusion::common::Result>> { unreachable!("test consumes reads directly") } - - async fn consume_expression( - &self, - expr: &Expression, - input_schema: &DFSchema, - context: &SubstraitConsumerContext, - ) -> datafusion::common::Result { - self.active.fetch_add(1, Ordering::SeqCst); - let _active = ActiveInput(&self.active); - self.ready.acquire().await.unwrap().forget(); - from_substrait_rex(self, expr, input_schema, context).await - } - async fn consume_rel( &self, rel: &Rel, _context: &SubstraitConsumerContext, ) -> datafusion::common::Result { let Some(proto::rel::RelType::Read(read)) = &rel.rel_type else { - panic!("expected a read") + panic!("expected read") }; let Some(proto::read_rel::ReadType::NamedTable(table)) = &read.read_type else { - panic!("expected a named table") + panic!("expected named table") + }; + let name = table.names[0].as_str(); + let _pending = self.pending.acquire().await.unwrap(); + let gate = if name == "a" { + &self.a_ready + } else { + &self.b_ready }; - self.reads.lock().unwrap().push(table.names[0].clone()); - self.active.fetch_add(1, Ordering::SeqCst); - let _active = ActiveInput(&self.active); - let gate = self.gates.get(&table.names[0]).unwrap_or(&self.ready); gate.acquire().await.unwrap().forget(); - if table.names[0] == "fail" { + if name == "fail" { return substrait_err!("input failed"); } table_scan( - Some(table.names[0].as_str()), + Some(name), &Schema::new(vec![Field::new("x", DataType::Int64, false)]), None, )? @@ -922,21 +914,7 @@ mod tests { } } - fn input_consumer() -> InputConsumer { - InputConsumer { - extensions: Arc::new(Extensions::default()), - state: Arc::new( - datafusion::execution::SessionStateBuilder::new_with_default_features() - .build(), - ), - reads: Arc::new(Mutex::default()), - ready: Arc::new(Semaphore::new(0)), - gates: Arc::new(HashMap::new()), - active: Arc::new(AtomicUsize::new(0)), - } - } - - fn input_read(name: &str) -> Rel { + fn read(name: &str) -> Rel { Rel { rel_type: Some(proto::rel::RelType::Read(Box::new(ReadRel { read_type: Some(proto::read_rel::ReadType::NamedTable( @@ -950,270 +928,60 @@ mod tests { } } - fn identity_lambda(kind: r#type::Kind, steps_out: u32) -> proto::expression::Lambda { - proto::expression::Lambda { - parameters: Some(r#type::Struct { - types: vec![Type { kind: Some(kind) }], + #[tokio::test] + async fn join_and_unions_consume_concurrently_in_input_order() { + let relations = [ + proto::rel::RelType::Join(Box::new(JoinRel { + left: Some(Box::new(read("a"))), + right: Some(Box::new(read("b"))), + r#type: proto::join_rel::JoinType::Inner as i32, ..Default::default() - }), - body: Some(Box::new(Expression { - rex_type: Some(proto::expression::RexType::Selection(Box::new( - FieldReference { - reference_type: Some( - field_reference::ReferenceType::DirectReference( - ReferenceSegment { - reference_type: Some( - reference_segment::ReferenceType::StructField( - Box::new(reference_segment::StructField { - field: 0, - child: None, - }), - ), - ), - }, - ), - ), - root_type: Some( - field_reference::RootType::LambdaParameterReference( - field_reference::LambdaParameterReference { steps_out }, - ), - ), - }, - ))), })), + proto::rel::RelType::Set(SetRel { + inputs: vec![read("a"), read("b")], + op: proto::set_rel::SetOp::UnionAll as i32, + ..Default::default() + }), + proto::rel::RelType::Set(SetRel { + inputs: vec![read("a"), read("b")], + op: proto::set_rel::SetOp::UnionDistinct as i32, + ..Default::default() + }), + ]; + for rel_type in relations { + let consumer = GatedConsumer::new(2); + let context = SubstraitConsumerContext::default(); + let rel = Rel { + rel_type: Some(rel_type), + }; + let mut result = Box::pin(from_substrait_rel(&consumer, &rel, &context)); + assert!(poll!(result.as_mut()).is_pending()); + assert_eq!(consumer.pending.available_permits(), 0); + consumer.b_ready.add_permits(1); // B finishes first. + assert!(poll!(result.as_mut()).is_pending()); + assert_eq!(consumer.pending.available_permits(), 1); + consumer.a_ready.add_permits(1); + let plan = result.await.unwrap().display_indent().to_string(); + assert!( + plan.find("TableScan: a").unwrap() < plan.find("TableScan: b").unwrap() + ); + assert_eq!(consumer.pending.available_permits(), 2); } } - fn assert_lambda_body(expr: Expr, name: &str, data_type: DataType) { - let Expr::Lambda(lambda) = expr else { - panic!("expected lambda") - }; - let Expr::LambdaVariable(variable) = *lambda.body else { - panic!("expected lambda variable") - }; - assert_eq!(variable.name, name); - assert_eq!(variable.field.unwrap().data_type(), &data_type); - } - - #[tokio::test] - async fn shared_consumer_keeps_concurrent_lambda_scopes_independent() { - let consumer = input_consumer(); - let context = SubstraitConsumerContext::default(); - let schema = DFSchema::empty(); - let integer = identity_lambda(r#type::Kind::I64(r#type::I64::default()), 0); - let string = identity_lambda(r#type::Kind::String(r#type::String::default()), 0); - let mut first = Box::pin(consumer.consume_lambda(&integer, &schema, &context)); - let mut second = Box::pin(consumer.consume_lambda(&string, &schema, &context)); - assert!(poll!(first.as_mut()).is_pending()); - assert!(poll!(second.as_mut()).is_pending()); - assert_eq!(consumer.active.load(Ordering::SeqCst), 2); - consumer.ready.add_permits(2); - assert_lambda_body(first.await.unwrap(), "p0", DataType::Int64); - assert_lambda_body(second.await.unwrap(), "p0", DataType::Utf8); - assert!(context.lambda_variable(0, 0).is_err()); - } - - #[tokio::test] - async fn lambda_error_and_cancellation_preserve_parent_scope() { - let consumer = input_consumer(); - let schema = DFSchema::empty(); - let root = SubstraitConsumerContext::default(); - let (context, names) = root - .with_lambda_parameters( - &consumer, - &[Type { - kind: Some(r#type::Kind::I64(r#type::I64::default())), - }], - &schema, - ) - .unwrap(); - assert_eq!(names, vec!["p0"]); - let nested = identity_lambda(r#type::Kind::String(r#type::String::default()), 1); - let mut cancelled = Box::pin(consumer.consume_lambda(&nested, &schema, &context)); - assert!(poll!(cancelled.as_mut()).is_pending()); - drop(cancelled); - assert_eq!(consumer.active.load(Ordering::SeqCst), 0); - let mut invalid = nested.clone(); - invalid.body = None; - assert!( - consumer - .consume_lambda(&invalid, &schema, &context) - .await - .is_err() - ); - consumer.ready.add_permits(1); - let result = consumer - .consume_lambda(&nested, &schema, &context) - .await - .unwrap(); - let Expr::Lambda(ref lambda) = result else { - panic!("expected lambda") - }; - assert_eq!(lambda.params, vec!["p1"]); - assert_lambda_body(result, "p0", DataType::Int64); - assert!(context.lambda_variable(1, 0).is_err()); - assert!(root.lambda_variable(0, 0).is_err()); - } - - #[tokio::test] - async fn concurrent_inputs_overlap_and_preserve_input_order() { - let mut consumer = input_consumer(); - consumer.gates = Arc::new(HashMap::from([ - ("a".into(), Semaphore::new(0)), - ("b".into(), Semaphore::new(0)), - ])); - let rels = [input_read("a"), input_read("b")]; - let inputs = rels.iter().collect::>(); - let context = SubstraitConsumerContext::default(); - let mut pending = Box::pin(consumer.consume_rel_inputs(&inputs, &context)); - assert!(poll!(&mut pending).is_pending()); - assert_eq!(*consumer.reads.lock().unwrap(), ["a", "b"]); - assert_eq!(consumer.active.load(Ordering::SeqCst), 2); - consumer.gates["b"].add_permits(1); - assert!(poll!(&mut pending).is_pending()); - assert_eq!(consumer.active.load(Ordering::SeqCst), 1); - consumer.gates["a"].add_permits(1); - let plans = pending.await.unwrap(); - assert_eq!(plans.len(), 2); - assert!( - plans[0] - .display_indent() - .to_string() - .contains("TableScan: a") - ); - assert!( - plans[1] - .display_indent() - .to_string() - .contains("TableScan: b") - ); - assert_eq!(consumer.active.load(Ordering::SeqCst), 0); - } - #[tokio::test] - async fn concurrent_inputs_cancel_pending_siblings_on_error() { - let mut consumer = input_consumer(); - consumer.gates = Arc::new(HashMap::from([("fail".into(), Semaphore::new(0))])); - // Exercise more inputs than the small-future optimization in try_join_all. - let mut rels = (0..40) - .map(|i| input_read(&format!("a{i}"))) - .collect::>(); - rels.push(input_read("fail")); - let inputs = rels.iter().collect::>(); + async fn input_failure_cancels_pending_siblings() { + let consumer = GatedConsumer::new(41); + // A later failure must not wait for earlier pending inputs, even in a wide UNION. + let mut inputs = vec![read("a"); 40]; + inputs.push(read("fail")); + let inputs = inputs.iter().collect::>(); let context = SubstraitConsumerContext::default(); - let mut pending = Box::pin(consumer.consume_rel_inputs(&inputs, &context)); - assert!(poll!(&mut pending).is_pending()); - assert_eq!(consumer.active.load(Ordering::SeqCst), 41); - consumer.gates["fail"].add_permits(1); - let error = pending.await.unwrap_err(); - assert!(error.to_string().contains("input failed")); - assert_eq!(consumer.active.load(Ordering::SeqCst), 0); - } - - struct HookConsumer { - inner: InputConsumer, - calls: Mutex>, - truncate: bool, - } - - #[async_trait] - impl SubstraitConsumer for HookConsumer { - fn get_extensions(&self) -> &Extensions { - &self.inner.extensions - } - fn get_function_registry(&self) -> &impl FunctionRegistry { - self.inner.state.as_ref() - } - async fn resolve_table_ref( - &self, - _: &TableReference, - ) -> datafusion::common::Result>> { - unreachable!("test consumes reads directly") - } - async fn consume_rel_inputs( - &self, - rels: &[&Rel], - context: &SubstraitConsumerContext, - ) -> datafusion::common::Result> { - self.calls.lock().unwrap().push(rels.len()); - let mut plans = self.inner.consume_rel_inputs(rels, context).await?; - if self.truncate { - plans.pop(); - } - Ok(plans) - } - } - - #[tokio::test] - async fn join_and_union_use_the_input_hook() { - for truncate in [false, true] { - let consumer = HookConsumer { - inner: input_consumer(), - calls: Mutex::default(), - truncate, - }; - consumer.inner.ready.add_permits(6); - let join = JoinRel { - left: Some(Box::new(input_read("a"))), - right: Some(Box::new(input_read("b"))), - r#type: proto::join_rel::JoinType::Inner as i32, - ..Default::default() - }; - let result = consumer - .consume_join(&join, &SubstraitConsumerContext::default()) - .await; - if truncate { - assert!( - result - .unwrap_err() - .to_string() - .contains("wrong number of inputs") - ); - } else { - let LogicalPlan::Join(join) = result.unwrap() else { - panic!("expected join") - }; - assert!( - join.left - .display_indent() - .to_string() - .contains("TableScan: a") - ); - assert!( - join.right - .display_indent() - .to_string() - .contains("TableScan: b") - ); - } - for op in [ - proto::set_rel::SetOp::UnionAll, - proto::set_rel::SetOp::UnionDistinct, - ] { - let set = SetRel { - inputs: vec![input_read("a"), input_read("b")], - op: op as i32, - ..Default::default() - }; - let result = consumer - .consume_set(&set, &SubstraitConsumerContext::default()) - .await; - if truncate { - assert!( - result - .unwrap_err() - .to_string() - .contains("wrong number of inputs") - ); - } else { - let plan = result.unwrap().display_indent().to_string(); - assert!( - plan.find("TableScan: a").unwrap() - < plan.find("TableScan: b").unwrap() - ); - } - } - assert_eq!(*consumer.calls.lock().unwrap(), [2, 2, 2]); - } + let mut result = Box::pin(consumer.consume_rel_inputs(&inputs, &context)); + assert!(poll!(result.as_mut()).is_pending()); + assert_eq!(consumer.pending.available_permits(), 0); + consumer.b_ready.add_permits(1); + assert_contains!(result.await.unwrap_err().to_string(), "input failed"); + assert_eq!(consumer.pending.available_permits(), 41); } } From 58ae265704ad25317283ab03deca2abbdf353404 Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Fri, 6 Nov 2026 12:28:38 -0500 Subject: [PATCH 5/9] refactor(substrait): preserve sequential consumption with explicit scopes --- Cargo.lock | 1 - datafusion/substrait/Cargo.toml | 1 - datafusion/substrait/README.md | 28 +-- .../consumer/substrait_consumer.rs | 168 +----------------- .../library-user-guide/upgrading/56.0.0.md | 28 +++ 5 files changed, 39 insertions(+), 187 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4ed040839f7a4..e3bfa7f5959b1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2792,7 +2792,6 @@ dependencies = [ "chrono", "datafusion", "datafusion-functions-aggregate", - "futures", "half", "insta", "itertools 0.15.0", diff --git a/datafusion/substrait/Cargo.toml b/datafusion/substrait/Cargo.toml index 4aa8868c17371..d6fc6daa9e6d6 100644 --- a/datafusion/substrait/Cargo.toml +++ b/datafusion/substrait/Cargo.toml @@ -21,7 +21,6 @@ rust-version = { workspace = true } workspace = true [dependencies] -futures = { workspace = true } async-recursion = "1.2" async-trait = { workspace = true } chrono = { workspace = true } diff --git a/datafusion/substrait/README.md b/datafusion/substrait/README.md index 5f6cdaefe63fd..a9652901d2b01 100644 --- a/datafusion/substrait/README.md +++ b/datafusion/substrait/README.md @@ -16,29 +16,5 @@ plans. See [API Docs] for details and examples. [substrait]: https://substrait.io [api docs]: https://docs.rs/datafusion-substrait/latest -## Migrating custom consumers to explicit scope context - -Relation and expression callbacks on `SubstraitConsumer` now take a final -`&SubstraitConsumerContext` argument. Forward that context when delegating to the -standard conversion helpers or consuming child relations and expressions. -`from_substrait_plan` and `from_substrait_plan_with_consumer` still create the root -context for you. Direct calls to conversion helpers should start with -`SubstraitConsumerContext::default()`. - -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. -The standard handlers resolve scopes through the context. Custom handlers that -introduce scopes should use `context.with_outer_schema(...)` or -`context.with_lambda_parameters(...)` and pass the resulting child context to -nested conversions. There is no matching pop operation; the parent context stays -unchanged on success, error, and cancellation. - -`consume_rel_inputs` now consumes JOIN and UNION inputs concurrently on the same -consumer. Keep traversal scopes in the context. Shared services and caches must -support concurrent calls. Consumers requiring sequential input processing can -override `consume_rel_inputs` and await each `consume_rel(rel, context)` in order. - -Generated lambda parameter names are now local to lexical scopes, so sibling -lambdas may reuse names. Logical plan displays can change without changing query -results. +For custom consumer migration instructions, see the +[DataFusion 56.0.0 upgrade guide](../../docs/source/library-user-guide/upgrading/56.0.0.md#substrait-consumers-use-explicit-scope-context). diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index 4ae3092022067..a4013b333ea85 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -23,8 +23,6 @@ use datafusion::common::{ use datafusion::execution::{FunctionRegistry, SessionState}; use datafusion::logical_expr::expr::LambdaVariable; use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; -use futures::TryStreamExt; -use futures::stream::FuturesUnordered; use std::collections::{HashMap, VecDeque}; use std::sync::Arc; use substrait::proto::expression as substrait_expression; @@ -44,9 +42,8 @@ use substrait::proto::{ /// It can be implemented by users to allow for custom handling of relations, expressions, etc. /// /// Relation and expression callbacks receive their lexical scope explicitly. -/// Independent inputs may invoke callbacks concurrently on the same consumer. -/// Keep traversal scopes in the context and make shared services concurrency-safe, -/// or override [`Self::consume_rel_inputs`] to consume inputs sequentially. +/// Inputs are consumed sequentially by default. Custom consumers may override +/// [`Self::consume_rel_inputs`] to choose a different scheduling policy. /// /// Combined with the [crate::logical_plan::producer::SubstraitProducer] this allows for fully /// customizable Substrait serde. @@ -193,24 +190,19 @@ pub trait SubstraitConsumer: Send + Sync + Sized { /// Implementations must return exactly one plan per input, in the same order. /// On error, implementations must drop any unfinished input futures. /// - /// The default consumes inputs concurrently using the same immutable context. - /// Overrides can provide a different scheduling policy. + /// The default consumes inputs sequentially, stopping at the first error. + /// Overrides can provide a different scheduling policy, but must account for + /// any additional mutable traversal state in the custom consumer. async fn consume_rel_inputs( &self, rels: &[&Rel], context: &SubstraitConsumerContext, ) -> datafusion::common::Result> { - let pending = FuturesUnordered::new(); - for (index, rel) in rels.iter().enumerate() { - pending.push(async move { - self.consume_rel(rel, context) - .await - .map(|plan| (index, plan)) - }); + let mut plans = Vec::with_capacity(rels.len()); + for rel in rels { + plans.push(self.consume_rel(rel, context).await?); } - let mut completed = pending.try_collect::>().await?; - completed.sort_unstable_by_key(|(index, _)| *index); - Ok(completed.into_iter().map(|(_, plan)| plan).collect()) + Ok(plans) } async fn consume_read( @@ -767,11 +759,7 @@ mod tests { use super::*; use crate::logical_plan::consumer::utils::tests::test_consumer; use datafusion::arrow::datatypes::{Field, Schema}; - use datafusion::common::assert_contains; use datafusion::logical_expr::ExprSchemable; - use datafusion::logical_expr::logical_plan::builder::table_scan; - use futures::poll; - use tokio::sync::Semaphore; fn schema(name: &str, data_type: DataType) -> Arc { Arc::new( @@ -846,142 +834,4 @@ mod tests { assert!(b.lambda_variable(1, 0).is_err()); assert!(nested.lambda_variable(0, 1).is_err()); } - - // Only input completion is controlled; JOIN/UNION conversion and batch - // scheduling use the standard consumer implementations. - struct GatedConsumer { - consumer: DefaultSubstraitConsumer<'static>, - a_ready: Semaphore, - b_ready: Semaphore, - // Held permits count pending inputs and are released even on cancellation. - pending: Semaphore, - } - - impl GatedConsumer { - fn new(inputs: usize) -> Self { - Self { - consumer: test_consumer(), - a_ready: Semaphore::new(0), - b_ready: Semaphore::new(0), - pending: Semaphore::new(inputs), - } - } - } - - #[async_trait] - impl SubstraitConsumer for GatedConsumer { - fn get_extensions(&self) -> &Extensions { - self.consumer.get_extensions() - } - fn get_function_registry(&self) -> &impl FunctionRegistry { - self.consumer.get_function_registry() - } - async fn resolve_table_ref( - &self, - _: &TableReference, - ) -> datafusion::common::Result>> { - unreachable!("test consumes reads directly") - } - async fn consume_rel( - &self, - rel: &Rel, - _context: &SubstraitConsumerContext, - ) -> datafusion::common::Result { - let Some(proto::rel::RelType::Read(read)) = &rel.rel_type else { - panic!("expected read") - }; - let Some(proto::read_rel::ReadType::NamedTable(table)) = &read.read_type - else { - panic!("expected named table") - }; - let name = table.names[0].as_str(); - let _pending = self.pending.acquire().await.unwrap(); - let gate = if name == "a" { - &self.a_ready - } else { - &self.b_ready - }; - gate.acquire().await.unwrap().forget(); - if name == "fail" { - return substrait_err!("input failed"); - } - table_scan( - Some(name), - &Schema::new(vec![Field::new("x", DataType::Int64, false)]), - None, - )? - .build() - } - } - - fn read(name: &str) -> Rel { - Rel { - rel_type: Some(proto::rel::RelType::Read(Box::new(ReadRel { - read_type: Some(proto::read_rel::ReadType::NamedTable( - proto::read_rel::NamedTable { - names: vec![name.into()], - ..Default::default() - }, - )), - ..Default::default() - }))), - } - } - - #[tokio::test] - async fn join_and_unions_consume_concurrently_in_input_order() { - let relations = [ - proto::rel::RelType::Join(Box::new(JoinRel { - left: Some(Box::new(read("a"))), - right: Some(Box::new(read("b"))), - r#type: proto::join_rel::JoinType::Inner as i32, - ..Default::default() - })), - proto::rel::RelType::Set(SetRel { - inputs: vec![read("a"), read("b")], - op: proto::set_rel::SetOp::UnionAll as i32, - ..Default::default() - }), - proto::rel::RelType::Set(SetRel { - inputs: vec![read("a"), read("b")], - op: proto::set_rel::SetOp::UnionDistinct as i32, - ..Default::default() - }), - ]; - for rel_type in relations { - let consumer = GatedConsumer::new(2); - let context = SubstraitConsumerContext::default(); - let rel = Rel { - rel_type: Some(rel_type), - }; - let mut result = Box::pin(from_substrait_rel(&consumer, &rel, &context)); - assert!(poll!(result.as_mut()).is_pending()); - assert_eq!(consumer.pending.available_permits(), 0); - consumer.b_ready.add_permits(1); // B finishes first. - assert!(poll!(result.as_mut()).is_pending()); - assert_eq!(consumer.pending.available_permits(), 1); - consumer.a_ready.add_permits(1); - let plan = result.await.unwrap().display_indent().to_string(); - assert!( - plan.find("TableScan: a").unwrap() < plan.find("TableScan: b").unwrap() - ); - assert_eq!(consumer.pending.available_permits(), 2); - } - } - - #[tokio::test] - async fn input_failure_cancels_pending_siblings() { - let consumer = GatedConsumer::new(41); - // A later failure must not wait for earlier pending inputs, even in a wide UNION. - let mut inputs = vec![read("a"); 40]; - inputs.push(read("fail")); - let inputs = inputs.iter().collect::>(); - let context = SubstraitConsumerContext::default(); - let mut result = Box::pin(consumer.consume_rel_inputs(&inputs, &context)); - assert!(poll!(result.as_mut()).is_pending()); - assert_eq!(consumer.pending.available_permits(), 0); - consumer.b_ready.add_permits(1); - assert_contains!(result.await.unwrap_err().to_string(), "input failed"); - assert_eq!(consumer.pending.available_permits(), 41); - } } diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 07714dcd5589c..c479388f39d06 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -1621,3 +1621,31 @@ let return_type: Option = definition `CreateFunction::try_new` takes `Option` for the return type. See [issue #18235](https://github.com/apache/datafusion/issues/18235) for details. + +### Substrait consumers use explicit scope context + +Relation and expression callbacks on `SubstraitConsumer` now take a final +`&SubstraitConsumerContext` argument. Forward that context when delegating to the +standard conversion helpers or consuming child relations and expressions. +`from_substrait_plan` and `from_substrait_plan_with_consumer` still create the root +context for you. When starting a new conversion directly through a helper, use +`SubstraitConsumerContext::default()`; nested conversions must receive the current +context. + +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. +The standard handlers resolve scopes through the context. Custom handlers that +introduce scopes should use `context.with_outer_schema(...)` or +`context.with_lambda_parameters(...)` and pass the resulting child context to +nested conversions. There is no matching pop operation; the parent context stays +unchanged on success, error, and cancellation. + +`consume_rel_inputs` consumes inputs sequentially by default, stopping at the first +error. This change does not enable parallel input consumption. Custom consumers +can continue to keep their own mutable traversal state; consumers that opt into +concurrent consumption must isolate that state themselves. + +Generated lambda parameter names are now local to lexical scopes, so sibling +lambdas may reuse names. Logical plan displays can change without changing query +results. From 87ae6252def7961e42820784da646a2bf6dcadd6 Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Fri, 6 Nov 2026 13:10:05 -0500 Subject: [PATCH 6/9] docs(substrait): keep migration notes in the upgrade guide --- datafusion/substrait/README.md | 3 --- 1 file changed, 3 deletions(-) diff --git a/datafusion/substrait/README.md b/datafusion/substrait/README.md index a9652901d2b01..7feeafbeeacb1 100644 --- a/datafusion/substrait/README.md +++ b/datafusion/substrait/README.md @@ -15,6 +15,3 @@ plans. See [API Docs] for details and examples. [apache datafusion]: https://datafusion.apache.org/ [substrait]: https://substrait.io [api docs]: https://docs.rs/datafusion-substrait/latest - -For custom consumer migration instructions, see the -[DataFusion 56.0.0 upgrade guide](../../docs/source/library-user-guide/upgrading/56.0.0.md#substrait-consumers-use-explicit-scope-context). From 4d8d227911be5e4f1d91d8f97db17d22fb0b6e1e Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Fri, 6 Nov 2026 18:03:57 -0500 Subject: [PATCH 7/9] feat(substrait): support scoped consumer context extensions --- .../consumer/substrait_consumer.rs | 46 +++++++++++++++++++ .../library-user-guide/upgrading/56.0.0.md | 7 ++- 2 files changed, 51 insertions(+), 2 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index a4013b333ea85..70dbdef9e9688 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, }; @@ -653,9 +654,28 @@ pub struct SubstraitConsumerContext { outer_schemas: Vec>, lambda_parameters: VecDeque>, next_lambda_parameter: usize, + extensions: ContextExtensions, } 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(); @@ -768,6 +788,32 @@ mod tests { ) } + #[test] + 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 outer_scopes_inherit_without_changing_the_parent() { let root = SubstraitConsumerContext::default(); diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index c479388f39d06..cfbe438053665 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -1643,8 +1643,11 @@ unchanged on success, error, and cancellation. `consume_rel_inputs` consumes inputs sequentially by default, stopping at the first error. This change does not enable parallel input consumption. Custom consumers -can continue to keep their own mutable traversal state; consumers that opt into -concurrent consumption must isolate that state themselves. +can move application-specific scoped values into the context using +`context.with_extension(Arc::new(value))` and retrieve them with +`context.get_extension::()`. Child contexts inherit values; replacing a +value affects only that child and its descendants. Values are shared through +`Arc`, so interior mutability still requires care when opting into concurrency. Generated lambda parameter names are now local to lexical scopes, so sibling lambdas may reuse names. Logical plan displays can change without changing query From 4e24776b80e068d6f59d86a7bde8b84b2a1b6487 Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Tue, 10 Nov 2026 08:17:31 -0500 Subject: [PATCH 8/9] refactor(substrait): defer input consumption hook to a follow-up --- .../src/logical_plan/consumer/rel/join_rel.rs | 29 +++++++------------ .../src/logical_plan/consumer/rel/set_rel.rs | 25 ++++++---------- .../consumer/substrait_consumer.rs | 23 +-------------- .../library-user-guide/upgrading/56.0.0.md | 3 +- 4 files changed, 22 insertions(+), 58 deletions(-) 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 c4b3354eeb2e7..c20cec0bbfbac 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/join_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/join_rel.rs @@ -3,10 +3,7 @@ // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; -use datafusion::common::{ - Column, JoinType, NullEquality, internal_datafusion_err, not_impl_err, plan_err, - substrait_err, -}; +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; use datafusion::logical_expr::{ @@ -24,20 +21,16 @@ pub async fn from_join_rel( return not_impl_err!("JoinRel with post_join_filter is not yet supported"); } - let (Some(left), Some(right)) = (join.left.as_deref(), join.right.as_deref()) else { - return substrait_err!("JoinRel requires left and right inputs"); - }; - let [left, right] = consumer - .consume_rel_inputs(&[left, right], context) - .await? - .try_into() - .map_err(|_| { - internal_datafusion_err!( - "JoinRel consumer returned the wrong number of inputs" - ) - })?; - let left = LogicalPlanBuilder::from(left); - let right = LogicalPlanBuilder::from(right); + let left: LogicalPlanBuilder = LogicalPlanBuilder::from( + consumer + .consume_rel(join.left.as_ref().unwrap(), context) + .await?, + ); + let right = LogicalPlanBuilder::from( + consumer + .consume_rel(join.right.as_ref().unwrap(), context) + .await?, + ); let (left, right, _requalified) = requalify_sides_if_needed(left, right)?; let join_type = from_substrait_jointype(join.r#type)?; 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 0311d8a0fc796..c42b8c00f0141 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs @@ -4,9 +4,7 @@ use crate::logical_plan::consumer::{SubstraitConsumer, SubstraitConsumerContext}; use datafusion::arrow::datatypes::Field; -use datafusion::common::{ - JoinType, NullEquality, internal_err, not_impl_err, substrait_err, -}; +use datafusion::common::{JoinType, NullEquality, not_impl_err, substrait_err}; use datafusion::logical_expr::{ Cast, Expr, LogicalPlan, LogicalPlanBuilder, requalify_sides_if_needed, }; @@ -55,21 +53,16 @@ async fn union_rels( is_all: bool, context: &SubstraitConsumerContext, ) -> datafusion::common::Result { - let inputs = rels.iter().collect::>(); - let plans = consumer.consume_rel_inputs(&inputs, context).await?; - if plans.len() != rels.len() { - return internal_err!("SetRel consumer returned the wrong number of inputs"); - } - let mut plans = plans.into_iter(); - let Some(first) = plans.next() else { - return substrait_err!("Set operation requires at least two inputs"); - }; - let mut union_builder = Ok(LogicalPlanBuilder::from(first)); - for plan in plans { + let mut union_builder = Ok(LogicalPlanBuilder::from( + consumer.consume_rel(&rels[0], context).await?, + )); + for input in &rels[1..] { + let rel_plan = consumer.consume_rel(input, context).await?; + union_builder = if is_all { - union_builder?.union(plan) + union_builder?.union(rel_plan) } else { - union_builder?.union_distinct(plan) + union_builder?.union_distinct(rel_plan) }; } union_builder?.build() diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index 70dbdef9e9688..5ead4f42ea220 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -43,8 +43,7 @@ use substrait::proto::{ /// 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. Custom consumers may override -/// [`Self::consume_rel_inputs`] to choose a different scheduling policy. +/// Inputs are consumed sequentially by default. /// /// Combined with the [crate::logical_plan::producer::SubstraitProducer] this allows for fully /// customizable Substrait serde. @@ -186,26 +185,6 @@ pub trait SubstraitConsumer: Send + Sync + Sized { from_substrait_rel(self, rel, context).await } - /// Consumes independent inputs of a JOIN or UNION, returning plans in input order. - /// - /// Implementations must return exactly one plan per input, in the same order. - /// On error, implementations must drop any unfinished input futures. - /// - /// The default consumes inputs sequentially, stopping at the first error. - /// Overrides can provide a different scheduling policy, but must account for - /// any additional mutable traversal state in the custom consumer. - async fn consume_rel_inputs( - &self, - rels: &[&Rel], - context: &SubstraitConsumerContext, - ) -> datafusion::common::Result> { - let mut plans = Vec::with_capacity(rels.len()); - for rel in rels { - plans.push(self.consume_rel(rel, context).await?); - } - Ok(plans) - } - async fn consume_read( &self, rel: &ReadRel, diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index cfbe438053665..0d64162d16cc0 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -1641,8 +1641,7 @@ introduce scopes should use `context.with_outer_schema(...)` or nested conversions. There is no matching pop operation; the parent context stays unchanged on success, error, and cancellation. -`consume_rel_inputs` consumes inputs sequentially by default, stopping at the first -error. This change does not enable parallel input consumption. Custom consumers +Input consumption remains sequential. Custom consumers can move application-specific scoped values into the context using `context.with_extension(Arc::new(value))` and retrieve them with `context.get_extension::()`. Child contexts inherit values; replacing a From 861b01489c1781c8b54e3895bc1e305b7e0d9284 Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Tue, 10 Nov 2026 09:51:50 -0500 Subject: [PATCH 9/9] docs(substrait): document consumer migration for DataFusion 57 --- .../library-user-guide/upgrading/56.0.0.md | 30 -------------- .../library-user-guide/upgrading/57.0.0.md | 40 +++++++++++++++++++ .../library-user-guide/upgrading/index.rst | 1 + 3 files changed, 41 insertions(+), 30 deletions(-) create mode 100644 docs/source/library-user-guide/upgrading/57.0.0.md diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 0d64162d16cc0..07714dcd5589c 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -1621,33 +1621,3 @@ let return_type: Option = definition `CreateFunction::try_new` takes `Option` for the return type. See [issue #18235](https://github.com/apache/datafusion/issues/18235) for details. - -### Substrait consumers use explicit scope context - -Relation and expression callbacks on `SubstraitConsumer` now take a final -`&SubstraitConsumerContext` argument. Forward that context when delegating to the -standard conversion helpers or consuming child relations and expressions. -`from_substrait_plan` and `from_substrait_plan_with_consumer` still create the root -context for you. When starting a new conversion directly through a helper, use -`SubstraitConsumerContext::default()`; nested conversions must receive the current -context. - -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. -The standard handlers resolve scopes through the context. Custom handlers that -introduce scopes should use `context.with_outer_schema(...)` or -`context.with_lambda_parameters(...)` and pass the resulting child context to -nested conversions. There is no matching pop operation; the parent context stays -unchanged on success, error, and cancellation. - -Input consumption remains sequential. Custom consumers -can move application-specific scoped values into the context using -`context.with_extension(Arc::new(value))` and retrieve them with -`context.get_extension::()`. Child contexts inherit values; replacing a -value affects only that child and its descendants. Values are shared through -`Arc`, so interior mutability still requires care when opting into concurrency. - -Generated lambda parameter names are now local to lexical scopes, so sibling -lambdas may reuse names. Logical plan displays can change without changing query -results. 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>