diff --git a/Cargo.lock b/Cargo.lock index c3c9bd6a76e9c..bbd4d77e1953f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2771,7 +2771,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -2915,7 +2915,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4174,7 +4174,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4928,7 +4928,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] @@ -5301,7 +5301,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -5556,6 +5556,7 @@ version = "1.0.149" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" dependencies = [ + "indexmap 2.13.0", "itoa", "memchr", "serde", @@ -5913,11 +5914,12 @@ dependencies = [ [[package]] name = "substrait" -version = "0.62.2" +version = "0.63.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62fc4b483a129b9772ccb9c3f7945a472112fdd9140da87f8a4e7f1d44e045d0" +checksum = "e620ff4d5c02fd6f7752931aa74b16a26af66a63022cc1ad412c77edbe0bab47" dependencies = [ "heck", + "indexmap 2.13.0", "pbjson 0.8.0", "pbjson-build 0.8.0", "pbjson-types", @@ -6009,7 +6011,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -6959,7 +6961,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/datafusion-testing b/datafusion-testing index eccb0e4a42634..13bbae38776c2 160000 --- a/datafusion-testing +++ b/datafusion-testing @@ -1 +1 @@ -Subproject commit eccb0e4a426344ef3faf534cd60e02e9c3afd3ac +Subproject commit 13bbae38776c2bfbc1fab1be7e7220222d4284bf diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 99688a52a75c1..c8de25df316cc 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -4532,49 +4532,49 @@ mod tests { [ { "Plan": { + "Node Type": "Projection", "Expressions": [ "employee_csv.id" ], - "Node Type": "Projection", - "Output": [ - "id" - ], "Plans": [ { - "Condition": "employee_csv.state IN ()", "Node Type": "Filter", - "Output": [ - "id", - "state" - ], + "Condition": "employee_csv.state IN ()", "Plans": [ { "Node Type": "Subquery", - "Output": [ - "state" - ], "Plans": [ { "Node Type": "TableScan", + "Relation Name": "employee_csv", + "Plans": [], "Output": [ "state" - ], - "Plans": [], - "Relation Name": "employee_csv" + ] } + ], + "Output": [ + "state" ] }, { "Node Type": "TableScan", + "Relation Name": "employee_csv", + "Plans": [], "Output": [ "id", "state" - ], - "Plans": [], - "Relation Name": "employee_csv" + ] } + ], + "Output": [ + "id", + "state" ] } + ], + "Output": [ + "id" ] } } diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index 6b50a1d1cbcd4..bbd319f00bbd7 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -642,11 +642,11 @@ logical_plan 02)--{ 03)----"Plan": { 04)------"Node Type": "Values", -05)------"Output": [ -06)--------"column1" -07)------], -08)------"Plans": [], -09)------"Values": "(Int64(1))" +05)------"Values": "(Int64(1))", +06)------"Plans": [], +07)------"Output": [ +08)--------"column1" +09)------] 10)----} 11)--} 12)] diff --git a/datafusion/substrait/Cargo.toml b/datafusion/substrait/Cargo.toml index 85479c344860d..a0f203cec8db6 100644 --- a/datafusion/substrait/Cargo.toml +++ b/datafusion/substrait/Cargo.toml @@ -44,7 +44,7 @@ object_store = { workspace = true } # We need to match the version in substrait, so we don't use the workspace version here pbjson-types = { version = "0.8.0" } prost = { workspace = true } -substrait = { version = "=0.62.2", features = ["serde"] } +substrait = { version = "0.63.0", features = ["serde"] } url = { workspace = true } tokio = { workspace = true, features = ["fs"] } diff --git a/datafusion/substrait/src/extensions.rs b/datafusion/substrait/src/extensions.rs index 0f848270babb9..78c357f3b8886 100644 --- a/datafusion/substrait/src/extensions.rs +++ b/datafusion/substrait/src/extensions.rs @@ -113,14 +113,10 @@ impl TryFrom<&Vec> for Extensions { } impl From for Vec { - // Silence deprecation warnings for `extension_uri_reference` during the uri -> urn migration - // See: https://github.com/substrait-io/substrait/issues/856 - #[expect(deprecated)] fn from(val: Extensions) -> Vec { let mut extensions = vec![]; for (f_anchor, f_name) in val.functions { let function_extension = ExtensionFunction { - extension_uri_reference: u32::MAX, extension_urn_reference: u32::MAX, function_anchor: f_anchor, name: f_name, @@ -133,7 +129,6 @@ impl From for Vec { for (t_anchor, t_name) in val.types { let type_extension = ExtensionType { - extension_uri_reference: u32::MAX, // https://github.com/apache/datafusion/issues/11545 extension_urn_reference: u32::MAX, // https://github.com/apache/datafusion/issues/11545 type_anchor: t_anchor, name: t_name, @@ -146,8 +141,7 @@ impl From for Vec { for (tv_anchor, tv_name) in val.type_variations { let type_variation_extension = ExtensionTypeVariation { - extension_uri_reference: u32::MAX, // We don't register proper extension URIs yet - extension_urn_reference: u32::MAX, // We don't register proper extension URIs yet + extension_urn_reference: u32::MAX, // We don't register proper extension URNs yet type_variation_anchor: tv_anchor, name: tv_name, }; 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 50d93a4600a00..aab9bde949203 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs @@ -21,7 +21,7 @@ use datafusion::logical_expr::Expr; use std::sync::Arc; use substrait::proto::expression::FieldReference; use substrait::proto::expression::field_reference::ReferenceType::DirectReference; -use substrait::proto::expression::field_reference::RootType; +use substrait::proto::expression::field_reference::{LambdaParameterReference, RootType}; use substrait::proto::expression::reference_segment::ReferenceType::StructField; pub async fn from_field_reference( @@ -56,6 +56,9 @@ pub(crate) fn from_substrait_field_reference( 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), } } _ => not_impl_err!( @@ -82,3 +85,83 @@ fn resolve_outer_reference( let col = Column::from((qualifier, field)); Ok(Expr::OuterReferenceColumn(Arc::clone(field), col)) } + +#[cfg(test)] +mod tests { + use datafusion::{ + common::{DFSchema, assert_contains}, + prelude::SessionContext, + }; + use substrait::proto::{ + Type, + expression::{ + FieldReference, ReferenceSegment, + field_reference::{self, LambdaParameterReference, RootType}, + reference_segment::{ReferenceType, StructField}, + }, + r#type::{I64, Kind}, + }; + + use crate::{ + extensions::Extensions, + logical_plan::consumer::{ + DefaultSubstraitConsumer, SubstraitConsumer, from_field_reference, + }, + }; + + #[tokio::test] + async fn test_lambda_variable_invalid_steps_out() { + let lambda_field_ref = lambda_field_ref(0, 99); + + let extensions = Extensions::default(); + let session_state = SessionContext::new().state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); + + let err = from_field_reference(&consumer, &lambda_field_ref, &DFSchema::empty()) + .await + .unwrap_err(); + + assert_contains!(err.to_string(), "No lambda at 99 steps out, got only 0"); + } + + #[tokio::test] + async fn test_lambda_variable_invalid_field_idx() { + let lambda_field_ref = lambda_field_ref(1, 0); + + let extensions = Extensions::default(); + let session_state = SessionContext::new().state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); + let _names = consumer + .push_lambda_parameters( + &[Type { + kind: Some(Kind::I64(I64::default())), + }], + &DFSchema::empty(), + ) + .unwrap(); + + let err = from_field_reference(&consumer, &lambda_field_ref, &DFSchema::empty()) + .await + .unwrap_err(); + + assert_contains!( + err.to_string(), + "At lambda 0 steps out, no field at index 1, got only 1" + ); + } + + fn lambda_field_ref(field: i32, steps_out: u32) -> FieldReference { + FieldReference { + reference_type: Some(field_reference::ReferenceType::DirectReference( + ReferenceSegment { + reference_type: Some(ReferenceType::StructField(Box::new( + StructField { field, child: None }, + ))), + }, + )), + root_type: Some(RootType::LambdaParameterReference( + LambdaParameterReference { steps_out }, + )), + } + } +} diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/lambda.rs b/datafusion/substrait/src/logical_plan/consumer/expr/lambda.rs new file mode 100644 index 0000000000000..2fb71024b054f --- /dev/null +++ b/datafusion/substrait/src/logical_plan/consumer/expr/lambda.rs @@ -0,0 +1,102 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion::{ + common::{DFSchema, substrait_err}, + prelude::{Expr, lambda}, +}; +use substrait::proto; + +use crate::logical_plan::consumer::SubstraitConsumer; + +pub async fn from_lambda( + consumer: &impl SubstraitConsumer, + expr: &proto::expression::Lambda, + input_schema: &DFSchema, +) -> 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 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(); + + Ok(lambda(names, body)) +} + +#[cfg(test)] +mod tests { + use datafusion::{ + common::{DFSchema, assert_contains}, + prelude::SessionContext, + }; + use substrait::proto::{self, Expression, r#type::Struct}; + + use crate::{ + extensions::Extensions, + logical_plan::consumer::{DefaultSubstraitConsumer, from_lambda}, + }; + + #[tokio::test] + async fn test_lambda_without_body() { + let lambda = proto::expression::Lambda { + parameters: Some(Struct::default()), + body: None, + }; + + let extensions = Extensions::default(); + let session_state = SessionContext::new().state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); + + let err = from_lambda(&consumer, &lambda, &DFSchema::empty()) + .await + .unwrap_err(); + + assert_contains!( + err.to_string(), + "Lambda expression without body is not allowed" + ); + } + + #[tokio::test] + async fn test_lambda_without_parameters() { + let lambda = proto::expression::Lambda { + parameters: None, + body: Some(Box::new(Expression::default())), + }; + + let extensions = Extensions::default(); + let session_state = SessionContext::new().state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); + + let err = from_lambda(&consumer, &lambda, &DFSchema::empty()) + .await + .unwrap_err(); + + assert_contains!( + err.to_string(), + "Lambda expression without parameters is not allowed" + ); + } +} diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs b/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs index ad38b6addee0b..d7d7a69581f05 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs @@ -43,7 +43,7 @@ use prost::Message; use std::sync::Arc; use substrait::proto; use substrait::proto::expression::Literal; -use substrait::proto::expression::literal::user_defined::Val; +use substrait::proto::expression::literal::user_defined::{TypeAnchorType, Val}; use substrait::proto::expression::literal::{ IntervalCompound, IntervalDayToSecond, IntervalYearToMonth, LiteralType, interval_day_to_second, @@ -474,11 +474,17 @@ pub(crate) fn from_substrait_literal( ))) }; - if let Some(name) = consumer - .get_extensions() - .types - .get(&user_defined.type_reference) - { + let type_ref = match user_defined.type_anchor_type { + Some(TypeAnchorType::TypeReference(ref_val)) => ref_val, + Some(TypeAnchorType::TypeAliasReference(_)) => { + return not_impl_err!( + "Type alias references in user-defined literals are not yet supported" + ); + } + None => 0, + }; + + if let Some(name) = consumer.get_extensions().types.get(&type_ref) { match name.as_ref() { FLOAT_16_TYPE_NAME => { // Rules for encoding fp16 Substrait literals are defined as part of Arrow here: @@ -518,14 +524,14 @@ pub(crate) fn from_substrait_literal( _ => { return not_impl_err!( "Unsupported Substrait user defined type with ref {} and name {}", - user_defined.type_reference, + type_ref, name ); } } } else { #[expect(deprecated)] - match user_defined.type_reference { + match type_ref { // Kept for backwards compatibility, producers should useIntervalYearToMonth instead INTERVAL_YEAR_MONTH_TYPE_REF => { let Some(Val::Value(raw_val)) = user_defined.val.as_ref() else { @@ -568,7 +574,7 @@ pub(crate) fn from_substrait_literal( _ => { return not_impl_err!( "Unsupported Substrait user defined type literal with ref {}", - user_defined.type_reference + type_ref ); } } diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs b/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs index 037d4ff688c25..2fcc11f4e417d 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs @@ -20,6 +20,7 @@ mod cast; mod field_reference; mod function_arguments; mod if_then; +mod lambda; mod literal; mod nested; mod scalar_function; @@ -32,6 +33,7 @@ pub use cast::*; pub use field_reference::*; pub use function_arguments::*; pub use if_then::*; +pub use lambda::*; pub use literal::*; pub use nested::*; pub use scalar_function::*; @@ -95,6 +97,12 @@ pub async fn from_substrait_rex( RexType::DynamicParameter(expr) => { consumer.consume_dynamic_parameter(expr, input_schema).await } + RexType::Lambda(lambda) => { + consumer.consume_lambda(lambda.as_ref(), input_schema).await + } + RexType::LambdaInvocation(_) => { + not_impl_err!("Lambda invocations are not supported") + } }, None => substrait_err!("Expression must set rex_type: {expression:?}"), } 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 1a0fb3f55f609..4cd856fc562e8 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs @@ -30,7 +30,6 @@ pub async fn from_scalar_function( f: &ScalarFunction, input_schema: &DFSchema, ) -> Result { - //TODO: handle higher order functions, as they are also encoded as scalar functions let Some(fn_signature) = consumer .get_extensions() .functions @@ -45,6 +44,20 @@ 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 higher_order_func = consumer + .get_function_registry() + .higher_order_function(fn_name) + .or_else(|e| { + if let Some(alt_name) = substrait_to_df_name(fn_name) { + consumer + .get_function_registry() + .higher_order_function(alt_name) + .or(Err(e)) + } else { + Err(e) + } + }); + let udf_func = consumer.get_function_registry().udf(fn_name).or_else(|e| { if let Some(alt_name) = substrait_to_df_name(fn_name) { consumer.get_function_registry().udf(alt_name).or(Err(e)) @@ -53,9 +66,14 @@ pub async fn from_scalar_function( } }); - // try to first match the requested function into registered udfs, then built-in ops + // try to first match the requested function into registered higher-order functions, then udfs, built-in ops // and finally built-in expressions - if let Ok(func) = udf_func { + if let Ok(func) = higher_order_func { + Ok(Expr::HigherOrderFunction(expr::HigherOrderFunction::new( + func.to_owned(), + args, + ))) + } else if let Ok(func) = udf_func { Ok(Expr::ScalarFunction(expr::ScalarFunction::new_udf( func.to_owned(), args, diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index c2bc16ea89be6..3bf9e9cce9dec 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -23,21 +23,24 @@ use super::{ from_substrait_rex, from_window_function, }; use crate::extensions::Extensions; +use crate::logical_plan::consumer::{from_lambda, from_substrait_type_without_names}; use async_trait::async_trait; -use datafusion::arrow::datatypes::DataType; +use datafusion::arrow::datatypes::{DataType, Field, FieldRef}; use datafusion::catalog::TableProvider; use datafusion::common::{ DFSchema, ScalarValue, TableReference, not_impl_err, substrait_err, }; use datafusion::execution::{FunctionRegistry, SessionState}; +use datafusion::logical_expr::expr::LambdaVariable; use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; +use std::collections::VecDeque; use std::sync::{Arc, RwLock}; -use substrait::proto; use substrait::proto::expression as substrait_expression; use substrait::proto::expression::{ Enum, FieldReference, IfThen, Literal, MultiOrList, Nested, ScalarFunction, SingularOrList, SwitchExpression, WindowFunction, }; +use substrait::proto::{self, Type}; use substrait::proto::{ AggregateRel, ConsistentPartitionWindowRel, CrossRel, DynamicParameter, ExchangeRel, Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel, @@ -62,17 +65,19 @@ use substrait::proto::{ /// # use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder}; /// # use std::sync::Arc; /// # use substrait::proto; -/// # use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel}; +/// # use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel, Type}; /// # use datafusion::arrow::datatypes::DataType; /// # 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 +/// # from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer, DefaultSubstraitLambdaConsumer /// # }; /// /// struct CustomSubstraitConsumer { /// extensions: Arc, /// state: Arc, +/// // You can reuse existing consumer code related to lambdas +/// lambda_consumer: DefaultSubstraitLambdaConsumer, /// } /// /// #[async_trait] @@ -95,6 +100,30 @@ 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?; @@ -141,7 +170,15 @@ use substrait::proto::{ /// /// // and user-defined literals /// fn consume_user_defined_literal(&self, literal: &proto::expression::literal::UserDefined) -> Result { -/// let type_string = self.extensions.types.get(&literal.type_reference).unwrap(); +/// // extract type_reference from the new TypeAnchorType oneof +/// let type_ref = match literal.type_anchor_type { +/// Some(proto::expression::literal::user_defined::TypeAnchorType::TypeReference(r)) => r, +/// Some(proto::expression::literal::user_defined::TypeAnchorType::TypeAliasReference(_)) => { +/// return not_impl_err!("Type alias references are not yet supported") +/// } +/// None => 0, +/// }; +/// let type_string = self.extensions.types.get(&type_ref).unwrap(); /// match type_string.as_str() { /// "u!foo" => not_impl_err!("handle foo conversion"), /// "u!bar" => not_impl_err!("handle bar conversion"), @@ -364,6 +401,14 @@ pub trait SubstraitConsumer: Send + Sync + Sized { not_impl_err!("Dynamic Parameter expression not supported") } + async fn consume_lambda( + &self, + expr: &proto::expression::Lambda, + input_schema: &DFSchema, + ) -> 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. @@ -444,10 +489,51 @@ pub trait SubstraitConsumer: Send + Sync + Sized { &self, user_defined_literal: &proto::expression::literal::UserDefined, ) -> datafusion::common::Result { - substrait_err!( - "Missing handler for user-defined literals {}", - user_defined_literal.type_reference - ) + let type_ref = match user_defined_literal.type_anchor_type { + Some( + proto::expression::literal::user_defined::TypeAnchorType::TypeReference( + ref_val, + ), + ) => ref_val, + Some( + proto::expression::literal::user_defined::TypeAnchorType::TypeAliasReference(_), + ) => { + return not_impl_err!( + "Type alias references in user-defined literals are not yet supported" + ) + } + None => 0, + }; + 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") } } @@ -458,6 +544,7 @@ pub struct DefaultSubstraitConsumer<'a> { pub(super) extensions: &'a Extensions, pub(super) state: &'a SessionState, outer_schemas: RwLock>>, + lambda_consumer: DefaultSubstraitLambdaConsumer, } impl<'a> DefaultSubstraitConsumer<'a> { @@ -466,6 +553,7 @@ impl<'a> DefaultSubstraitConsumer<'a> { extensions, state, outer_schemas: RwLock::new(Vec::new()), + lambda_consumer: DefaultSubstraitLambdaConsumer::new(), } } } @@ -562,6 +650,140 @@ impl SubstraitConsumer for DefaultSubstraitConsumer<'_> { 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 +/// +/// 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`. + lambda_parameters: VecDeque>, + next_lambda_parameter: usize, +} + +impl Default for DefaultSubstraitLambdaConsumer { + fn default() -> Self { + Self::new() + } +} + +impl DefaultSubstraitLambdaConsumer { + pub fn new() -> Self { + Self { + inner: RwLock::new(DefaultSubstraitLambdaConsumerInner { + lambda_parameters: VecDeque::new(), + next_lambda_parameter: 0, + }), + } + } + + pub fn push_lambda_parameters( + &self, + consumer: &impl SubstraitConsumer, + lambda_parameters: &[Type], + input_schema: &DFSchema, + ) -> datafusion::common::Result> { + let mut inner = self.inner.write().unwrap(); + + let lambda_parameters: Vec = lambda_parameters + .iter() + .map(|ty| { + let (assigned_number, default_name) = + next_lambda_parameter_name(inner.next_lambda_parameter, input_schema); + + inner.next_lambda_parameter = assigned_number + 1; + + let data_type = from_substrait_type_without_names(consumer, ty)?; + Ok(Arc::new(Field::new(default_name, data_type, true))) + }) + .collect::>>()?; + + let names = lambda_parameters.iter().map(|f| f.name().clone()).collect(); + + inner.lambda_parameters.push_front(lambda_parameters); + + Ok(names) + } + + pub fn pop_lambda_parameters(&self) { + self.inner.write().unwrap().lambda_parameters.pop_front(); + } + + pub fn lambda_variable( + &self, + steps_out: usize, + field_idx: usize, + ) -> datafusion::common::Result { + let lambda_parameters = &self.inner.read().unwrap().lambda_parameters; + + let Some(lambda_parameters) = lambda_parameters.get(steps_out) else { + return substrait_err!( + "No lambda at {steps_out} steps out, got only {}", + lambda_parameters.len() + ); + }; + + let Some(var) = lambda_parameters.get(field_idx) else { + return substrait_err!( + "At lambda {steps_out} steps out, no field at index {field_idx}, got only {}", + lambda_parameters.len() + ); + }; + + Ok(Expr::LambdaVariable(LambdaVariable::new( + var.name().clone(), + Some(Arc::clone(var)), + ))) + } +} + +/// Returns the next available lambda parameter name and the index it was assigned. +/// +/// Names follow the pattern `pN` where `N` starts at `next_lambda_parameter`. If `pN` +/// conflicts with an existing column name in `input_schema`, `N` is incremented until +/// a free name is found. +fn next_lambda_parameter_name( + mut next_lambda_parameter: usize, + input_schema: &DFSchema, +) -> (usize, String) { + loop { + let name = format!("p{next_lambda_parameter}"); + + // avoid conflicts with column names + if !input_schema.has_column_with_unqualified_name(&name) { + return (next_lambda_parameter, name); + } + + next_lambda_parameter += 1; + } } #[cfg(test)] diff --git a/datafusion/substrait/src/logical_plan/producer/expr/lambda.rs b/datafusion/substrait/src/logical_plan/producer/expr/lambda.rs new file mode 100644 index 0000000000000..0d32dab2ccadd --- /dev/null +++ b/datafusion/substrait/src/logical_plan/producer/expr/lambda.rs @@ -0,0 +1,48 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion::{common::DFSchemaRef, logical_expr::expr::Lambda}; +use substrait::proto::{ + Expression, + expression::RexType, + r#type::{Nullability, Struct}, +}; + +use crate::logical_plan::producer::SubstraitProducer; + +pub fn from_lambda( + producer: &mut impl SubstraitProducer, + lambda: &Lambda, + schema: &DFSchemaRef, +) -> Result { + Ok(Expression { + rex_type: Some(RexType::Lambda(Box::new( + substrait::proto::expression::Lambda { + parameters: Some(Struct { + nullability: Nullability::Required as i32, + type_variation_reference: 0, + types: lambda + .params + .iter() + .map(|p| producer.lambda_parameter_type(p)) + .collect::>()?, + }), + body: Some(Box::new(producer.handle_expr(&lambda.body, schema)?)), + }, + ))), + }) +} diff --git a/datafusion/substrait/src/logical_plan/producer/expr/lambda_variable.rs b/datafusion/substrait/src/logical_plan/producer/expr/lambda_variable.rs new file mode 100644 index 0000000000000..3d7f06e2332a1 --- /dev/null +++ b/datafusion/substrait/src/logical_plan/producer/expr/lambda_variable.rs @@ -0,0 +1,49 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion::logical_expr::expr::LambdaVariable; +use substrait::proto::{ + Expression, + expression::{ + FieldReference, ReferenceSegment, RexType, + field_reference::{LambdaParameterReference, ReferenceType, RootType}, + reference_segment::{self, StructField}, + }, +}; + +use crate::logical_plan::producer::SubstraitProducer; + +pub fn from_lambda_variable( + producer: &mut impl SubstraitProducer, + lambda_variable: &LambdaVariable, + _schema: &datafusion::common::DFSchema, +) -> Result { + let (steps_out, field) = producer.lambda_variable(&lambda_variable.name)?; + + Ok(Expression { + rex_type: Some(RexType::Selection(Box::new(FieldReference { + reference_type: Some(ReferenceType::DirectReference(ReferenceSegment { + reference_type: Some(reference_segment::ReferenceType::StructField( + Box::new(StructField { field, child: None }), + )), + })), + root_type: Some(RootType::LambdaParameterReference( + LambdaParameterReference { steps_out }, + )), + }))), + }) +} diff --git a/datafusion/substrait/src/logical_plan/producer/expr/literal.rs b/datafusion/substrait/src/logical_plan/producer/expr/literal.rs index 8882c992dca1c..bbed7ee9be417 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/literal.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/literal.rs @@ -117,7 +117,7 @@ pub(crate) fn to_substrait_literal( ( LiteralType::UserDefined( substrait::proto::expression::literal::UserDefined { - type_reference: type_anchor, + type_anchor_type: Some(substrait::proto::expression::literal::user_defined::TypeAnchorType::TypeReference(type_anchor)), type_parameters: vec![], val: Some(substrait::proto::expression::literal::user_defined::Val::Value( pbjson_types::Any { diff --git a/datafusion/substrait/src/logical_plan/producer/expr/mod.rs b/datafusion/substrait/src/logical_plan/producer/expr/mod.rs index 63d97610b8346..9f46c3d34cf67 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/mod.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/mod.rs @@ -19,6 +19,8 @@ mod aggregate_function; mod cast; mod field_reference; mod if_then; +mod lambda; +mod lambda_variable; mod literal; mod scalar_function; mod singular_or_list; @@ -29,6 +31,8 @@ pub use aggregate_function::*; pub use cast::*; pub use field_reference::*; pub use if_then::*; +pub use lambda::*; +pub use lambda_variable::*; pub use literal::*; pub use scalar_function::*; pub use singular_or_list::*; @@ -60,9 +64,6 @@ use substrait::version; /// /// Substrait also requires the input schema of the expressions to be included in the /// message. The field names of the input schema will be serialized. -// Silence deprecation warnings for `extension_uris` during the uri -> urn migration -// See: https://github.com/substrait-io/substrait/issues/856 -#[expect(deprecated)] pub fn to_substrait_extended_expr( exprs: &[(&Expr, &Field)], schema: &DFSchemaRef, @@ -87,7 +88,6 @@ pub fn to_substrait_extended_expr( Ok(Box::new(ExtendedExpression { advanced_extensions: None, expected_type_urls: vec![], - extension_uris: vec![], extension_urns: vec![], extensions: extensions.into(), version: Some(version::version_with_producer("datafusion")), @@ -156,10 +156,8 @@ pub fn to_substrait_rex( Expr::HigherOrderFunction(expr) => { producer.handle_higher_order_function(expr, schema) } - Expr::Lambda(expr) => not_impl_err!("Cannot convert {expr:?} to Substrait"), - Expr::LambdaVariable(expr) => { - not_impl_err!("Cannot convert {expr:?} to Substrait") - } + Expr::Lambda(expr) => producer.handle_lambda(expr, schema), + Expr::LambdaVariable(expr) => producer.handle_lambda_variable(expr, schema), } } diff --git a/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs b/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs index e36d5128cd293..e7dd2af13f9ca 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs @@ -16,7 +16,10 @@ // under the License. use crate::logical_plan::producer::{SubstraitProducer, to_substrait_literal_expr}; -use datafusion::common::{DFSchemaRef, ScalarValue, not_impl_err}; +use datafusion::common::datatype::FieldExt; +use datafusion::common::{ + DFSchemaRef, ScalarValue, internal_datafusion_err, not_impl_err, substrait_err, +}; use datafusion::logical_expr::{Between, BinaryExpr, Expr, Like, Operator, expr}; use substrait::proto::expression::{RexType, ScalarFunction}; use substrait::proto::function_argument::ArgType; @@ -35,7 +38,78 @@ pub fn from_higher_order_function( fun: &expr::HigherOrderFunction, schema: &DFSchemaRef, ) -> datafusion::common::Result { - from_function(producer, fun.name(), &fun.args, schema) + let mut lambda_parameters = fun.lambda_parameters(schema)?.into_iter(); + + let num_lambdas = fun + .args + .iter() + .filter(|arg| matches!(arg, Expr::Lambda(_))) + .count(); + + if lambda_parameters.len() != num_lambdas { + return substrait_err!( + "{} returned {} lambdas but {num_lambdas} expected", + fun.name(), + lambda_parameters.len() + ); + } + + let arguments = fun + .args + .iter() + .map(|arg| { + let arg = match arg { + Expr::Lambda(l) => { + let lambda_parameters = + lambda_parameters.next().ok_or_else(|| { + internal_datafusion_err!( + "lambda_parameters len should have been checked above" + ) + })?; + + if l.params.len() > lambda_parameters.len() { + return substrait_err!( + "Lambda defined {} parameters ({}) but function {} supports only {}", + l.params.len(), + l.params.join(","), + fun.name(), + lambda_parameters.len() + ) + } + + let named_lambda_parameters = + std::iter::zip(&l.params, lambda_parameters) + .map(|(name, parameter)| parameter.renamed(name)) + .collect(); + + producer.push_lambda_parameters(named_lambda_parameters)?; + + let arg = producer.handle_lambda(l, schema); + + producer.pop_lambda_parameters()?; + + arg + } + _ => producer.handle_expr(arg, schema), + }?; + + Ok(FunctionArgument { + arg_type: Some(ArgType::Value(arg)), + }) + }) + .collect::>()?; + + let function_anchor = producer.register_function(fun.name().to_string()); + #[expect(deprecated)] + Ok(Expression { + rex_type: Some(RexType::ScalarFunction(ScalarFunction { + function_reference: function_anchor, + arguments, + output_type: None, + options: vec![], + args: vec![], + })), + }) } fn from_function( diff --git a/datafusion/substrait/src/logical_plan/producer/plan.rs b/datafusion/substrait/src/logical_plan/producer/plan.rs index 9396329f8d3e7..3b58720dba832 100644 --- a/datafusion/substrait/src/logical_plan/producer/plan.rs +++ b/datafusion/substrait/src/logical_plan/producer/plan.rs @@ -24,9 +24,6 @@ use substrait::proto::{Plan, PlanRel, Rel, RelRoot, plan_rel}; use substrait::version; /// Convert DataFusion LogicalPlan to Substrait Plan -// Silence deprecation warnings for `extension_uris` during the uri -> urn migration -// See: https://github.com/substrait-io/substrait/issues/856 -#[expect(deprecated)] pub fn to_substrait_plan( plan: &LogicalPlan, state: &SessionState, @@ -47,7 +44,6 @@ pub fn to_substrait_plan( let extensions = producer.get_extensions(); Ok(Box::new(Plan { version: Some(version::version_with_producer("datafusion")), - extension_uris: vec![], extension_urns: vec![], extensions: extensions.into(), relations: plan_rels, diff --git a/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs b/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs index a38c1abea0de6..d10d1af091265 100644 --- a/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs +++ b/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs @@ -20,17 +20,22 @@ use crate::logical_plan::producer::{ from_aggregate, from_aggregate_function, from_alias, from_between, from_binary_expr, from_case, from_cast, from_column, from_distinct, from_empty_relation, from_exists, from_filter, from_higher_order_function, from_in_list, from_in_subquery, from_join, - from_like, from_limit, from_literal, from_projection, from_repartition, - from_scalar_function, from_scalar_subquery, from_set_comparison, from_sort, - from_subquery_alias, from_table_scan, from_try_cast, from_unary_expr, from_union, - from_values, from_window, from_window_function, to_substrait_rel, to_substrait_rex, + from_lambda, from_lambda_variable, from_like, from_limit, from_literal, + from_projection, from_repartition, from_scalar_function, from_scalar_subquery, + from_set_comparison, from_sort, from_subquery_alias, from_table_scan, from_try_cast, + from_unary_expr, from_union, from_values, from_window, from_window_function, + to_substrait_rel, to_substrait_rex, to_substrait_type, +}; +use datafusion::arrow::datatypes::FieldRef; +use datafusion::common::{ + Column, DFSchemaRef, HashMap, ScalarValue, not_impl_err, substrait_err, }; -use datafusion::common::{Column, DFSchemaRef, ScalarValue, substrait_err}; use datafusion::execution::SessionState; use datafusion::execution::registry::SerializerRegistry; use datafusion::logical_expr::Subquery; use datafusion::logical_expr::expr::{ - Alias, Exists, InList, InSubquery, SetComparison, WindowFunction, + Alias, Exists, InList, InSubquery, Lambda, LambdaVariable, Placeholder, + SetComparison, WindowFunction, }; use datafusion::logical_expr::{ Aggregate, Between, BinaryExpr, Case, Cast, Distinct, EmptyRelation, Expr, Extension, @@ -56,16 +61,19 @@ use substrait::proto::{ /// # use std::sync::Arc; /// # use substrait::proto::{Expression, Rel}; /// # use substrait::proto::rel::RelType; +/// # use datafusion::arrow::datatypes::FieldRef; /// # use datafusion::common::DFSchemaRef; /// # use datafusion::error::Result; /// # use datafusion::execution::SessionState; /// # use datafusion::logical_expr::{Between, Extension, Projection}; /// # use datafusion_substrait::extensions::Extensions; -/// # use datafusion_substrait::logical_plan::producer::{from_projection, SubstraitProducer}; +/// # use datafusion_substrait::logical_plan::producer::{from_projection, SubstraitProducer, DefaultSubstraitLambdaProducer, lambda_parameters_map}; /// /// struct CustomSubstraitProducer { /// extensions: Extensions, /// state: Arc, +/// // You can reuse existing producer code related to lambdas +/// lambda_producer: DefaultSubstraitLambdaProducer, /// } /// /// impl SubstraitProducer for CustomSubstraitProducer { @@ -82,6 +90,33 @@ use substrait::proto::{ /// self.extensions /// } /// +/// fn push_lambda_parameters( +/// &mut self, +/// lambda_parameters: Vec, +/// ) -> datafusion::common::Result<()> { +/// let lambda_parameters_map = lambda_parameters_map(self, lambda_parameters)?; +/// +/// self.lambda_producer +/// .push_lambda_parameters(lambda_parameters_map); +/// +/// Ok(()) +/// } +/// +/// fn pop_lambda_parameters(&mut self) -> datafusion::common::Result<()> { +/// self.lambda_producer.pop_lambda_parameters() +/// } +/// +/// fn lambda_variable(&self, name: &str) -> datafusion::common::Result<(u32, i32)> { +/// self.lambda_producer.lambda_variable(name) +/// } +/// +/// fn lambda_parameter_type( +/// &self, +/// name: &str, +/// ) -> datafusion::common::Result { +/// self.lambda_producer.lambda_parameter_type(name) +/// } +/// /// // You can set additional metadata on the Rels you produce /// fn handle_projection(&mut self, plan: &Projection) -> Result> { /// let mut rel = from_projection(self, plan)?; @@ -396,11 +431,73 @@ pub trait SubstraitProducer: Send + Sync + Sized { ) -> datafusion::common::Result { from_exists(self, exists, schema) } + + fn handle_placeholder( + &mut self, + placeholder: &Placeholder, + _schema: &DFSchemaRef, + ) -> datafusion::common::Result { + not_impl_err!("Cannot convert {placeholder:?} to Substrait") + } + + fn handle_lambda( + &mut self, + lambda: &Lambda, + schema: &DFSchemaRef, + ) -> datafusion::common::Result { + from_lambda(self, lambda, schema) + } + + fn handle_lambda_variable( + &mut self, + lambda_variable: &LambdaVariable, + schema: &DFSchemaRef, + ) -> datafusion::common::Result { + from_lambda_variable(self, lambda_variable, schema) + } + + // Lambda related methods + + /// Push the given `lambda_parameters` into this producer so they can be referenced by lambda variables + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaProducer] and forward this method to it + fn push_lambda_parameters( + &mut self, + _lambda_parameters: Vec, + ) -> datafusion::common::Result<()> { + not_impl_err!("SubstraitProducer::push_lambda_parameters") + } + + /// Pop the last pushed `lambda_parameters` so that it unshadow any previously shadowed lambda parameter + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaProducer] and forward this method to it + fn pop_lambda_parameters(&mut self) -> datafusion::common::Result<()> { + not_impl_err!("SubstraitProducer::pop_lambda_parameters") + } + + /// Get the (`steps_out`, `field_idx`) of the lambda variable with the given `name`. `steps_out` refers to the number + /// of lambda boundaries to traverse (0 = current lambda), and `field_idx` refers to the index within the lambda parameters + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaProducer] and forward this method to it + fn lambda_variable(&self, _name: &str) -> datafusion::common::Result<(u32, i32)> { + not_impl_err!("SubstraitProducer::lambda_variable") + } + + /// Get the type of the lambda parameter with the given `name` + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaProducer] and forward this method to it + fn lambda_parameter_type( + &self, + _name: &str, + ) -> datafusion::common::Result { + not_impl_err!("SubstraitProducer::lambda_parameter_type") + } } pub struct DefaultSubstraitProducer<'a> { extensions: Extensions, serializer_registry: &'a dyn SerializerRegistry, + lambda_producer: DefaultSubstraitLambdaProducer, } impl<'a> DefaultSubstraitProducer<'a> { @@ -408,6 +505,7 @@ impl<'a> DefaultSubstraitProducer<'a> { DefaultSubstraitProducer { extensions: Extensions::default(), serializer_registry: state.serializer_registry().as_ref(), + lambda_producer: DefaultSubstraitLambdaProducer::new(), } } } @@ -462,4 +560,112 @@ impl SubstraitProducer for DefaultSubstraitProducer<'_> { rel_type: Some(rel_type), })) } + + fn push_lambda_parameters( + &mut self, + lambda_parameters: Vec, + ) -> datafusion::common::Result<()> { + let lambda_parameters_map = lambda_parameters_map(self, lambda_parameters)?; + + self.lambda_producer + .push_lambda_parameters(lambda_parameters_map); + + Ok(()) + } + + fn pop_lambda_parameters(&mut self) -> datafusion::common::Result<()> { + self.lambda_producer.pop_lambda_parameters() + } + + fn lambda_variable(&self, name: &str) -> datafusion::common::Result<(u32, i32)> { + self.lambda_producer.lambda_variable(name) + } + + fn lambda_parameter_type( + &self, + name: &str, + ) -> datafusion::common::Result { + self.lambda_producer.lambda_parameter_type(name) + } +} + +/// Default implementation of lambda related methods of the [SubstraitProducer] trait +/// +/// Can be embedded into a custom [SubstraitProducer] to implement them +pub struct DefaultSubstraitLambdaProducer { + lambdas_variables: Vec>, +} + +impl Default for DefaultSubstraitLambdaProducer { + fn default() -> Self { + Self::new() + } +} + +impl DefaultSubstraitLambdaProducer { + pub fn new() -> Self { + Self { + lambdas_variables: Vec::new(), + } + } + + /// Note you can construct the `lambda_parameters` argument using [lambda_parameters_map] + pub fn push_lambda_parameters( + &mut self, + lambda_parameters: HashMap, + ) { + self.lambdas_variables.push(lambda_parameters); + } + + pub fn pop_lambda_parameters(&mut self) -> datafusion::common::Result<()> { + match self.lambdas_variables.pop() { + Some(_) => Ok(()), + None => substrait_err!("no lambda_parameters to pop"), + } + } + + pub fn lambda_variable(&self, name: &str) -> datafusion::common::Result<(u32, i32)> { + for (steps_out, lambda_parameters) in + self.lambdas_variables.iter().rev().enumerate() + { + if let Some((field_idx, _type)) = lambda_parameters.get(name) { + return Ok((steps_out as u32, *field_idx as i32)); + } + } + + substrait_err!("unknown lambda variable {name}") + } + + pub fn lambda_parameter_type( + &self, + name: &str, + ) -> datafusion::common::Result { + for lambda_parameters in self.lambdas_variables.iter().rev() { + if let Some((_field_idx, type_)) = lambda_parameters.get(name) { + return Ok(type_.clone()); + } + } + + substrait_err!("unknown lambda variable {name}") + } +} + +/// Produces a map of lambda parameters as expected by [DefaultSubstraitLambdaProducer::push_lambda_parameters] +pub fn lambda_parameters_map( + producer: &mut impl SubstraitProducer, + lambda_parameters: Vec, +) -> datafusion::common::Result> { + lambda_parameters + .into_iter() + .enumerate() + .map(|(field_idx, field)| { + Ok(( + field.name().clone(), + ( + field_idx, + to_substrait_type(producer, field.data_type(), field.is_nullable())?, + ), + )) + }) + .collect::>() } diff --git a/datafusion/substrait/tests/cases/logical_plans.rs b/datafusion/substrait/tests/cases/logical_plans.rs index 663a372fe2e4f..d4ac01462c879 100644 --- a/datafusion/substrait/tests/cases/logical_plans.rs +++ b/datafusion/substrait/tests/cases/logical_plans.rs @@ -19,6 +19,7 @@ #[cfg(test)] mod tests { + use crate::cases::roundtrip_logical_plan::higher_order_function_ctx; use crate::utils::test::{add_plan_schemas_to_ctx, read_json}; use datafusion::common::test_util::format_batches; use std::collections::HashSet; @@ -293,4 +294,25 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn higher_order_function() -> Result<()> { + let proto_plan = + read_json("tests/testdata/test_plans/higher_order_function.json"); + // ctx already contains the queried table + let ctx = higher_order_function_ctx().await?; + let plan = from_substrait_plan(&ctx.state(), &proto_plan).await?; + + 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))) + TableScan: data3 + " + ); + + // Trigger execution to ensure plan validity + DataFrame::new(ctx.state(), plan).show().await?; + Ok(()) + } } diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index f562be66b8e31..3a1befe8e6710 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -17,8 +17,14 @@ use crate::utils::test::read_json; use datafusion::arrow::array::ArrayRef; +use datafusion::config::Dialect; +use datafusion::execution::FunctionRegistry; use datafusion::functions_nested::map::map; -use datafusion::logical_expr::LogicalPlanBuilder; +use datafusion::logical_expr::{ + ColumnarValue, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, + HigherOrderSignature, HigherOrderUDF, LambdaParametersProgress, LogicalPlanBuilder, + ValueOrLambda, +}; use datafusion::physical_plan::Accumulator; use datafusion::scalar::ScalarValue; use datafusion_substrait::logical_plan::{ @@ -27,7 +33,9 @@ use datafusion_substrait::logical_plan::{ use std::cmp::Ordering; use std::mem::size_of_val; -use datafusion::arrow::datatypes::{DataType, Field, IntervalUnit, Schema, TimeUnit}; +use datafusion::arrow::datatypes::{ + DataType, Field, FieldRef, IntervalUnit, Schema, TimeUnit, +}; use datafusion::common::tree_node::Transformed; use datafusion::common::{DFSchema, DFSchemaRef, Spans, not_impl_err, plan_err}; use datafusion::error::Result; @@ -1775,6 +1783,215 @@ async fn roundtrip_read_filter() -> Result<()> { roundtrip_verify_read_filter_count("SELECT a FROM data where a < 5", 1).await } +#[tokio::test] +async fn roundtrip_array_transform_higher_order_function() -> Result<()> { + let ctx = higher_order_function_ctx().await?; + + // simple + roundtrip_with_ctx( + "SELECT array_transform2([data3.p1], p0 -> p0 * 2) from data3", + ctx.clone(), + ) + .await?; + + // dont use the parameter + roundtrip_with_ctx( + "SELECT array_transform2([data3.p1], p0 -> 3) from data3", + ctx.clone(), + ) + .await?; + + // multiple parameters using both + roundtrip_with_ctx( + "SELECT array_transform2([data3.p1], (p0, p2) -> p0 * p2) from data3", + ctx.clone(), + ) + .await?; + + // multiple parameters only last + roundtrip_with_ctx( + "SELECT array_transform2([data3.p1], (p0, p2) -> 2 * p2) from data3", + ctx.clone(), + ) + .await?; + + // multiple parameters use none + roundtrip_with_ctx( + "SELECT array_transform2([data3.p1], (p0, p2) -> 3) from data3", + ctx.clone(), + ) + .await?; + + // nested without variable shadowing + roundtrip_with_ctx("SELECT array_transform2([[data3.p1]], p0 -> array_transform2(p0, p2 -> p2 * 2)) from data3", ctx.clone()) + .await?; + + // nested with multiple parameters without variable shadowing + roundtrip_with_ctx("SELECT array_transform2([[data3.p1]], (p0, p2) -> array_transform2(p0, (p3, p4) -> p2 * p3 * p4)) from data3", ctx.clone()) + .await?; + + // since substrait doesn't encode lambda parameters names, they got generated, non-conflicting names during consumption + // testing name shadowing requires to assert against the generated plan and check the correct parameter usage instead of round tripping + + // nested with variable shadowing. + let plan = generate_plan_from_sql_with_ctx( + "SELECT array_transform2([[data3.p1]], v -> array_transform2(v, v -> v * 2)) from data3", + true, + true, + &ctx, + ) + .await?; + + assert_snapshot!( + plan, + @" + Projection: array_transform2(make_array(make_array(data3.p1)), (p0) -> array_transform2(p0, (p2) -> p2 * Int64(2))) AS array_transform2(make_array(make_array(data3.p1)),(v) -> array_transform2(v,(v) -> v * Int64(2))) + TableScan: data3 projection=[p1] + " + ); + + // nested with variable shadowing with multiple parameters + let plan = generate_plan_from_sql_with_ctx( + "SELECT array_transform2([[data3.p1]], (v, i) -> array_transform2(v, (v, i) -> v * i)) from data3", + true, + true, + &ctx, + ) + .await?; + + assert_snapshot!( + plan, + @" + Projection: array_transform2(make_array(make_array(data3.p1)), (p0, p2) -> array_transform2(p0, (p3, p4) -> p3 * p4)) AS array_transform2(make_array(make_array(data3.p1)),(v, i) -> array_transform2(v,(v, i) -> v * i)) + TableScan: data3 projection=[p1] + " + ); + + // nested with variable shadowing and later reuse of the shadowed var after exiting the shadowing expression + let plan = generate_plan_from_sql_with_ctx( + "SELECT array_transform2( + [[data3.p1]], + v -> array_concat( + -- when entering this expression, inner v is pushed into the producer and shadows outer v, but after exiting this, + -- it should be removed and unshadow the outer v, so that it can be used in the next expression + array_transform2(v, v -> v * 2), + array_transform2(v, v -> v * 2) + ) + ) from data3", + true, + true, + &ctx, + ) + .await?; + + 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)))) + TableScan: data3 projection=[p1] + " + ); + + Ok(()) +} + +pub(crate) async fn higher_order_function_ctx() -> Result { + let mut ctx = create_context_with_dialect(Some(Dialect::Databricks)).await?; + + ctx.register_higher_order_function(Arc::new(ArrayTransform::new()))?; + + let data3_fields = vec![ + Field::new("p1", DataType::Int64, true), // lambda parameters should not conflict with this column + ]; + let data3 = Schema::new(data3_fields); + let mut data3_options = CsvReadOptions::new(); + data3_options.schema = Some(&data3); + data3_options.has_header = false; + ctx.register_csv("data3", "tests/testdata/empty.csv", data3_options) + .await?; + + Ok(ctx) +} + +// todo use core array_transform when it supports multiple lambda parameters +#[derive(Debug, PartialEq, Eq, Hash)] +struct ArrayTransform { + signature: HigherOrderSignature, +} + +impl ArrayTransform { + fn new() -> Self { + Self { + signature: HigherOrderSignature::variadic_any(Volatility::Immutable), + } + } +} + +impl HigherOrderUDF for ArrayTransform { + fn name(&self) -> &str { + "array_transform2" + } + + fn aliases(&self) -> &[String] { + &[] + } + + fn signature(&self) -> &HigherOrderSignature { + &self.signature + } + + fn lambda_parameters( + &self, + _step: usize, + fields: &[ValueOrLambda>], + ) -> Result { + let [ValueOrLambda::Value(list), ValueOrLambda::Lambda(_)] = fields else { + unreachable!() + }; + + let field = match list.data_type() { + DataType::List(field) => field, + _ => unreachable!(), + }; + + Ok(LambdaParametersProgress::Complete(vec![vec![ + Arc::clone(field), + Arc::new(Field::new("", DataType::Int64, true)), + ]])) + } + + fn return_field_from_args( + &self, + args: HigherOrderReturnFieldArgs, + ) -> Result { + let [ValueOrLambda::Value(list), ValueOrLambda::Lambda(lambda)] = args.arg_fields + else { + unreachable!() + }; + + let field = Arc::new(Field::new( + Field::LIST_FIELD_DEFAULT_NAME, + lambda.data_type().clone(), + lambda.is_nullable(), + )); + + let return_type = match list.data_type() { + DataType::List(_) => DataType::List(field), + _ => unreachable!(), + }; + + Ok(Arc::new(Field::new("", return_type, list.is_nullable()))) + } + + fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result { + // this function is only tested with roundtrip_with_ctx, which only prints the output + // and generate_plan_from_sql_with_ctx which doesn't execute nothing, so the output doesn't matter + Ok(ColumnarValue::Scalar(ScalarValue::new_default( + args.return_type(), + )?)) + } +} + fn check_post_join_filters(rel: &Rel) -> Result<()> { // search for target_rel and field value in proto match &rel.rel_type { @@ -1916,6 +2133,15 @@ async fn generate_plan_from_sql( optimized: bool, ) -> Result { let ctx = create_context().await?; + generate_plan_from_sql_with_ctx(sql, assert_schema, optimized, &ctx).await +} + +async fn generate_plan_from_sql_with_ctx( + sql: &str, + assert_schema: bool, + optimized: bool, + ctx: &SessionContext, +) -> Result { let df: DataFrame = ctx.sql(sql).await?; let plan = if optimized { @@ -2269,8 +2495,18 @@ async fn roundtrip_all_types(sql: &str) -> Result<()> { } async fn create_context() -> Result { + create_context_with_dialect(None).await +} + +async fn create_context_with_dialect(dialect: Option) -> Result { + let mut session_config = SessionConfig::default(); + + if let Some(dialect) = dialect { + session_config.options_mut().sql_parser.dialect = dialect; + } + let mut state = SessionStateBuilder::new() - .with_config(SessionConfig::default()) + .with_config(session_config) .with_runtime_env(Arc::new(RuntimeEnv::default())) .with_default_features() .with_serializer_registry(Arc::new(MockSerializerRegistry)) diff --git a/datafusion/substrait/tests/cases/serialize.rs b/datafusion/substrait/tests/cases/serialize.rs index 2d7257fad3394..1981ef66db377 100644 --- a/datafusion/substrait/tests/cases/serialize.rs +++ b/datafusion/substrait/tests/cases/serialize.rs @@ -28,9 +28,16 @@ mod tests { use insta::assert_snapshot; use std::fs; + use substrait::proto::expression::field_reference::{ReferenceType, RootType}; + use substrait::proto::expression::reference_segment; + use substrait::proto::expression::{ReferenceSegment, RexType}; + use substrait::proto::function_argument::ArgType; use substrait::proto::plan_rel::RelType; use substrait::proto::rel_common::{Emit, EmitKind}; - use substrait::proto::{RelCommon, rel}; + use substrait::proto::r#type::{I64, Kind as TypeKind, List, Nullability, Struct}; + use substrait::proto::{Expression, RelCommon, Type, rel}; + + use crate::cases::roundtrip_logical_plan::higher_order_function_ctx; #[tokio::test] async fn serialize_to_file() -> Result<()> { @@ -196,6 +203,101 @@ mod tests { panic!("plan did not match expected structure") } + #[tokio::test] + async fn higher_order_function() -> Result<()> { + let ctx = higher_order_function_ctx().await?; + let df = ctx + .sql( + "SELECT array_transform2( + [[data3.p1]], + (v, i) -> array_concat( + -- when entering this expression, inner v is pushed into the producer and shadows outer v, but after exiting this, + -- it should be removed and unshadow the outer v, so that it can be used in the next expression + array_transform2(v, (v, j) -> v * i * j), + array_transform2(v, (v, j) -> v * i * j) + ) + ) from data3" + ) + .await?; + let datafusion_plan = df.into_optimized_plan()?; + let plan = to_substrait_plan(&datafusion_plan, &ctx.state())? + .as_ref() + .clone(); + + let relation = plan.relations.first().unwrap().rel_type.as_ref(); + let root_rel = match relation { + Some(RelType::Root(root)) => root.input.as_ref().unwrap(), + _ => panic!("expected Root"), + }; + + let Some(rel::RelType::Project(p)) = root_rel.rel_type.as_ref() else { + panic!("expected Project at top of plan") + }; + + let mut params = vec![]; + let mut lambda_refs = vec![]; + + collect_lambda_ref(&p.expressions[0], &mut params, &mut lambda_refs); + + let nullable_i64 = Type { + kind: Some(TypeKind::I64(I64 { + type_variation_reference: 0, + nullability: Nullability::Nullable as i32, + })), + }; + + let inner_lambda_struct = Struct { + // v, j + types: vec![nullable_i64.clone(); 2], + type_variation_reference: 0, + nullability: Nullability::Required as i32, + }; + + assert_eq!( + params, + vec![ + Struct { + types: vec![ + // v + Type { + kind: Some(TypeKind::List(Box::new(List { + r#type: Some(Box::new(nullable_i64.clone())), + type_variation_reference: 0, + nullability: Nullability::Nullable as i32 + }))) + }, + // i + nullable_i64, + ], + type_variation_reference: 0, + nullability: Nullability::Required as i32, + }, + inner_lambda_struct.clone(), + inner_lambda_struct, + ] + ); + + assert_eq!( + lambda_refs, + vec![ + // first inner array_transform2 argument: outer v + (0, 0), + // first inner lambda body: v * i * j + (0, 0), + (1, 1), + (0, 1), + // second inner array_transform2 argument: outer v + (0, 0), + // second inner lambda body: v * i * j + (0, 0), + (1, 1), + (0, 1), + ] + ); + + Ok(()) + } + fn assert_emit(rel_common: Option<&RelCommon>, output_mapping: Vec) { assert_eq!( rel_common.unwrap().emit_kind.clone(), @@ -211,4 +313,54 @@ mod tests { .await?; Ok(ctx) } + + // Recursively walks a expression tree depth-first, collecting in visit order: + // - `params`: the parameter struct of each Lambda encountered + // - `lambda_refs`: every field reference whose root is a LambdaParameterReference, + // recorded as (steps_out, field_index) so tests can assert which enclosing + // lambda each reference resolves to and which parameter within it. + fn collect_lambda_ref( + expr: &Expression, + params: &mut Vec, + lambda_refs: &mut Vec<(u32, i32)>, + ) { + if let Some(rex_type) = &expr.rex_type { + match rex_type { + RexType::Selection(field_reference) => { + if let ( + Some(ReferenceType::DirectReference(ReferenceSegment { + reference_type: + Some(reference_segment::ReferenceType::StructField( + struct_field, + )), + })), + Some(RootType::LambdaParameterReference(lambda_param_ref)), + ) = (&field_reference.reference_type, &field_reference.root_type) + { + lambda_refs.push((lambda_param_ref.steps_out, struct_field.field)) + } + } + RexType::ScalarFunction(scalar_function) => { + for arg in &scalar_function.arguments { + match &arg.arg_type { + Some(ArgType::Value(value)) => { + collect_lambda_ref(value, params, lambda_refs) + } + _ => unreachable!(), + } + } + } + RexType::Lambda(lambda) => { + if let Some(parameters) = &lambda.parameters { + params.push(parameters.clone()); + } + if let Some(body) = &lambda.body { + collect_lambda_ref(body, params, lambda_refs); + } + } + RexType::Literal(_literal) => {} + _ => unreachable!(), + } + } + } } diff --git a/datafusion/substrait/tests/testdata/test_plans/higher_order_function.json b/datafusion/substrait/tests/testdata/test_plans/higher_order_function.json new file mode 100644 index 0000000000000..da613b2573447 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/higher_order_function.json @@ -0,0 +1,438 @@ +{ + "version": { + "minorNumber": 85, + "producer": "datafusion" + }, + "extensions": [ + { + "extensionFunction": { + "extensionUrnReference": 2, + "functionAnchor": 2, + "name": "array_transform2" + } + }, + { + "extensionFunction": { + "extensionUrnReference": 2, + "name": "make_array" + } + }, + { + "extensionFunction": { + "extensionUrnReference": 2, + "functionAnchor": 3, + "name": "array_concat" + } + }, + { + "extensionFunction": { + "extensionUrnReference": 1, + "functionAnchor": 1, + "name": "multiply" + } + } + ], + "relations": [ + { + "root": { + "input": { + "project": { + "common": { + "emit": { + "outputMapping": [ + 1 + ] + } + }, + "input": { + "read": { + "baseSchema": { + "names": [ + "p1" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "projection": { + "select": { + "structItems": [ + {} + ] + } + }, + "namedTable": { + "names": [ + "data3" + ] + } + } + }, + "expressions": [ + { + "scalarFunction": { + "functionReference": 2, + "arguments": [ + { + "value": { + "scalarFunction": { + "arguments": [ + { + "value": { + "scalarFunction": { + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "rootReference": {} + } + } + } + ], + "outputType": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + }, + { + "value": { + "lambda": { + "parameters": { + "types": [ + { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + }, + "body": { + "scalarFunction": { + "functionReference": 3, + "arguments": [ + { + "value": { + "scalarFunction": { + "functionReference": 2, + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + }, + { + "value": { + "lambda": { + "parameters": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + }, + "body": { + "scalarFunction": { + "functionReference": 1, + "arguments": [ + { + "value": { + "scalarFunction": { + "functionReference": 1, + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + }, + { + "value": { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "lambdaParameterReference": { + "stepsOut": 1 + } + } + } + } + ], + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + }, + { + "value": { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "lambdaParameterReference": {} + } + } + } + ], + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + }, + { + "value": { + "scalarFunction": { + "functionReference": 2, + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + }, + { + "value": { + "lambda": { + "parameters": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + }, + "body": { + "scalarFunction": { + "functionReference": 1, + "arguments": [ + { + "value": { + "scalarFunction": { + "functionReference": 1, + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + }, + { + "value": { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "lambdaParameterReference": { + "stepsOut": 1 + } + } + } + } + ], + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + }, + { + "value": { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "lambdaParameterReference": {} + } + } + } + ], + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + ] + } + }, + "names": [ + "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)))" + ] + } + } + ], + "extensionUrns": [ + { + "extensionUrnAnchor": 1, + "urn": "extension:io.substrait:functions_arithmetic" + }, + { + "extensionUrnAnchor": 2, + "urn": "extension:io.substrait:functions_list" + } + ] +} diff --git a/datafusion/substrait/tests/utils.rs b/datafusion/substrait/tests/utils.rs index 6a6824579b4e8..4d9b5ca83e5e0 100644 --- a/datafusion/substrait/tests/utils.rs +++ b/datafusion/substrait/tests/utils.rs @@ -486,6 +486,7 @@ pub mod test { // Enum is deprecated #[expect(deprecated)] RexType::Enum(_) => {} + RexType::Lambda(_) | RexType::LambdaInvocation(_) => {} } Ok(()) } diff --git a/testing b/testing index 0d60ccae40d0e..7df2b70baf4f0 160000 --- a/testing +++ b/testing @@ -1 +1 @@ -Subproject commit 0d60ccae40d0e8f2d22c15fafb01c5d4be8c63a6 +Subproject commit 7df2b70baf4f081ebf8e0c6bd22745cf3cbfd824