This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-22362-b79d3965b7a2c29853e83def04656408d9d7ee7f in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 7ac784f9ea0ba7386bb8aa4cccbf7d0fac2f9510 Author: gstvg <[email protected]> AuthorDate: Fri Jul 10 13:45:15 2026 -0300 Add protobuf support for lambdas (#22362) ## Which issue does this PR close? Part of #21172 ## Rationale for this change Protobuf support wasn't implemented in main lambda PR to not make it even bigger ## What changes are included in this PR? Protobuf encoding and decoding (~1000 LOC in generated files, ~210 impl, ~400 tests) ## Are these changes tested? Unit tests, similar to the existing ones for scalar functions ## Are there any user-facing changes? Proto `ExprType` has new variants --- datafusion/physical-expr/src/expressions/lambda.rs | 39 + .../src/expressions/lambda_variable.rs | 44 + datafusion/proto-models/proto/datafusion.proto | 40 + datafusion/proto-models/src/generated/pbjson.rs | 1101 +++++++++++++++++--- datafusion/proto-models/src/generated/prost.rs | 62 +- datafusion/proto/src/logical_plan/from_proto.rs | 30 +- datafusion/proto/src/logical_plan/to_proto.rs | 39 +- datafusion/proto/src/physical_plan/from_proto.rs | 34 +- datafusion/proto/src/physical_plan/mod.rs | 20 +- datafusion/proto/src/physical_plan/to_proto.rs | 15 +- datafusion/proto/tests/cases/mod.rs | 69 +- .../proto/tests/cases/roundtrip_logical_plan.rs | 159 ++- .../proto/tests/cases/roundtrip_physical_plan.rs | 131 ++- datafusion/proto/tests/cases/serialize.rs | 97 +- 14 files changed, 1691 insertions(+), 189 deletions(-) diff --git a/datafusion/physical-expr/src/expressions/lambda.rs b/datafusion/physical-expr/src/expressions/lambda.rs index 9275821ae9..cab2eea64f 100644 --- a/datafusion/physical-expr/src/expressions/lambda.rs +++ b/datafusion/physical-expr/src/expressions/lambda.rs @@ -142,6 +142,27 @@ impl LambdaExpr { &self.body } + #[cfg(feature = "proto")] + /// Reconstruct a [`LambdaExpr`] from a proto node. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result<Arc<dyn PhysicalExpr>> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let lambda = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::Lambda, + "LambdaExpr", + ); + + Ok(Arc::new(LambdaExpr::try_new( + lambda.params.clone(), + ctx.decode_required_expression(lambda.body.as_deref(), "LambdaExpr", "body")?, + )?)) + } + pub(crate) fn projection(&self) -> &[usize] { &self.projection } @@ -193,6 +214,24 @@ impl PhysicalExpr for LambdaExpr { fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "({}) -> {}", self.params.join(", "), self.body) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Lambda(Box::new( + protobuf::PhysicalLambdaExprNode { + params: self.params().to_vec(), + body: Some(Box::new(ctx.encode_child(self.body())?)), + }, + ))), + })) + } } /// Create a lambda expression diff --git a/datafusion/physical-expr/src/expressions/lambda_variable.rs b/datafusion/physical-expr/src/expressions/lambda_variable.rs index 1c130ab12e..f7e6910020 100644 --- a/datafusion/physical-expr/src/expressions/lambda_variable.rs +++ b/datafusion/physical-expr/src/expressions/lambda_variable.rs @@ -72,6 +72,32 @@ impl LambdaVariable { pub fn field(&self) -> &FieldRef { &self.field } + + #[cfg(feature = "proto")] + /// Reconstruct a [`LambdaVariable`] from a proto node. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result<Arc<dyn PhysicalExpr>> { + use datafusion_physical_expr_common::{ + expect_expr_variant, physical_expr::proto_decode::require_proto_field, + }; + use datafusion_proto_models::protobuf; + + let var = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::LambdaVariable, + "LambdaVariable", + ); + + Ok(Arc::new(LambdaVariable::new( + var.index as usize, + Arc::new( + require_proto_field(var.field.as_ref(), "LambdaVariable", "field")? + .try_into()?, + ), + ))) + } } impl std::fmt::Display for LambdaVariable { @@ -135,6 +161,24 @@ impl PhysicalExpr for LambdaVariable { fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}@{}", self.name(), self.index) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::LambdaVariable( + protobuf::PhysicalLambdaVariableExprNode { + index: self.index() as u32, + field: Some(self.field().as_ref().try_into()?), + }, + )), + })) + } } /// Create a lambda variable expression diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index d68973c44e..b22fad9c0e 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -501,6 +501,10 @@ message LogicalExprNode { // Subquery expressions ScalarSubqueryExprNode scalar_subquery_expr = 36; + + HigherOrderUDFExprNode higher_order_udf_expr = 37; + Lambda lambda = 38; + LambdaVariable lambda_variable = 39; } } @@ -638,6 +642,22 @@ message ScalarUDFExprNode { optional bytes fun_definition = 3; } +message HigherOrderUDFExprNode { + string fun_name = 1; + repeated LogicalExprNode args = 2; + optional bytes fun_definition = 3; +} + +message Lambda { + repeated string params = 1; + LogicalExprNode body = 2; +} + +message LambdaVariable { + string name = 1; + datafusion_common.Field field = 2; +} + message WindowExprNode { oneof window_function { // BuiltInWindowFunction built_in_function = 2; @@ -1007,6 +1027,10 @@ message PhysicalExprNode { PhysicalScalarSubqueryExprNode scalar_subquery = 22; PhysicalDynamicFilterNode dynamic_filter = 23; + + PhysicalHigherOrderUdfNode higher_order_udf = 24; + PhysicalLambdaExprNode lambda = 25; + PhysicalLambdaVariableExprNode lambda_variable = 26; } } @@ -1027,6 +1051,22 @@ message PhysicalScalarUdfNode { string return_field_name = 6; } +message PhysicalHigherOrderUdfNode { + string name = 1; + repeated PhysicalExprNode args = 2; + optional bytes fun_definition = 3; +} + +message PhysicalLambdaExprNode { + repeated string params = 1; + PhysicalExprNode body = 2; +} + +message PhysicalLambdaVariableExprNode { + uint32 index = 1; + datafusion_common.Field field = 2; +} + message PhysicalAggregateExprNode { oneof AggregateFunction { string user_defined_aggr_function = 4; diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index f8e2103035..c334eac2f5 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -9342,6 +9342,137 @@ impl<'de> serde::Deserialize<'de> for HashRepartition { deserializer.deserialize_struct("datafusion.HashRepartition", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for HigherOrderUdfExprNode { + #[allow(deprecated)] + fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.fun_name.is_empty() { + len += 1; + } + if !self.args.is_empty() { + len += 1; + } + if self.fun_definition.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.HigherOrderUDFExprNode", len)?; + if !self.fun_name.is_empty() { + struct_ser.serialize_field("funName", &self.fun_name)?; + } + if !self.args.is_empty() { + struct_ser.serialize_field("args", &self.args)?; + } + if let Some(v) = self.fun_definition.as_ref() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("funDefinition", pbjson::private::base64::encode(&v).as_str())?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for HigherOrderUdfExprNode { + #[allow(deprecated)] + fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error> + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "fun_name", + "funName", + "args", + "fun_definition", + "funDefinition", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + FunName, + Args, + FunDefinition, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error> + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str<E>(self, value: &str) -> std::result::Result<GeneratedField, E> + where + E: serde::de::Error, + { + match value { + "funName" | "fun_name" => Ok(GeneratedField::FunName), + "args" => Ok(GeneratedField::Args), + "funDefinition" | "fun_definition" => Ok(GeneratedField::FunDefinition), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = HigherOrderUdfExprNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.HigherOrderUDFExprNode") + } + + fn visit_map<V>(self, mut map_: V) -> std::result::Result<HigherOrderUdfExprNode, V::Error> + where + V: serde::de::MapAccess<'de>, + { + let mut fun_name__ = None; + let mut args__ = None; + let mut fun_definition__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::FunName => { + if fun_name__.is_some() { + return Err(serde::de::Error::duplicate_field("funName")); + } + fun_name__ = Some(map_.next_value()?); + } + GeneratedField::Args => { + if args__.is_some() { + return Err(serde::de::Error::duplicate_field("args")); + } + args__ = Some(map_.next_value()?); + } + GeneratedField::FunDefinition => { + if fun_definition__.is_some() { + return Err(serde::de::Error::duplicate_field("funDefinition")); + } + fun_definition__ = + map_.next_value::<::std::option::Option<::pbjson::private::BytesDeserialize<_>>>()?.map(|x| x.0) + ; + } + } + } + Ok(HigherOrderUdfExprNode { + fun_name: fun_name__.unwrap_or_default(), + args: args__.unwrap_or_default(), + fun_definition: fun_definition__, + }) + } + } + deserializer.deserialize_struct("datafusion.HigherOrderUDFExprNode", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for ILikeNode { #[allow(deprecated)] fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> @@ -11321,7 +11452,7 @@ impl<'de> serde::Deserialize<'de> for JsonSinkExecNode { deserializer.deserialize_struct("datafusion.JsonSinkExecNode", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for LikeNode { +impl serde::Serialize for Lambda { #[allow(deprecated)] fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> where @@ -11329,54 +11460,37 @@ impl serde::Serialize for LikeNode { { use serde::ser::SerializeStruct; let mut len = 0; - if self.negated { - len += 1; - } - if self.expr.is_some() { + if !self.params.is_empty() { len += 1; } - if self.pattern.is_some() { - len += 1; - } - if !self.escape_char.is_empty() { + if self.body.is_some() { len += 1; } - let mut struct_ser = serializer.serialize_struct("datafusion.LikeNode", len)?; - if self.negated { - struct_ser.serialize_field("negated", &self.negated)?; - } - if let Some(v) = self.expr.as_ref() { - struct_ser.serialize_field("expr", v)?; - } - if let Some(v) = self.pattern.as_ref() { - struct_ser.serialize_field("pattern", v)?; + let mut struct_ser = serializer.serialize_struct("datafusion.Lambda", len)?; + if !self.params.is_empty() { + struct_ser.serialize_field("params", &self.params)?; } - if !self.escape_char.is_empty() { - struct_ser.serialize_field("escapeChar", &self.escape_char)?; + if let Some(v) = self.body.as_ref() { + struct_ser.serialize_field("body", v)?; } struct_ser.end() } } -impl<'de> serde::Deserialize<'de> for LikeNode { +impl<'de> serde::Deserialize<'de> for Lambda { #[allow(deprecated)] fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error> where D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "negated", - "expr", - "pattern", - "escape_char", - "escapeChar", + "params", + "body", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { - Negated, - Expr, - Pattern, - EscapeChar, + Params, + Body, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error> @@ -11398,10 +11512,8 @@ impl<'de> serde::Deserialize<'de> for LikeNode { E: serde::de::Error, { match value { - "negated" => Ok(GeneratedField::Negated), - "expr" => Ok(GeneratedField::Expr), - "pattern" => Ok(GeneratedField::Pattern), - "escapeChar" | "escape_char" => Ok(GeneratedField::EscapeChar), + "params" => Ok(GeneratedField::Params), + "body" => Ok(GeneratedField::Body), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -11411,60 +11523,44 @@ impl<'de> serde::Deserialize<'de> for LikeNode { } struct GeneratedVisitor; impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = LikeNode; + type Value = Lambda; fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.LikeNode") + formatter.write_str("struct datafusion.Lambda") } - fn visit_map<V>(self, mut map_: V) -> std::result::Result<LikeNode, V::Error> + fn visit_map<V>(self, mut map_: V) -> std::result::Result<Lambda, V::Error> where V: serde::de::MapAccess<'de>, { - let mut negated__ = None; - let mut expr__ = None; - let mut pattern__ = None; - let mut escape_char__ = None; + let mut params__ = None; + let mut body__ = None; while let Some(k) = map_.next_key()? { match k { - GeneratedField::Negated => { - if negated__.is_some() { - return Err(serde::de::Error::duplicate_field("negated")); - } - negated__ = Some(map_.next_value()?); - } - GeneratedField::Expr => { - if expr__.is_some() { - return Err(serde::de::Error::duplicate_field("expr")); - } - expr__ = map_.next_value()?; - } - GeneratedField::Pattern => { - if pattern__.is_some() { - return Err(serde::de::Error::duplicate_field("pattern")); + GeneratedField::Params => { + if params__.is_some() { + return Err(serde::de::Error::duplicate_field("params")); } - pattern__ = map_.next_value()?; + params__ = Some(map_.next_value()?); } - GeneratedField::EscapeChar => { - if escape_char__.is_some() { - return Err(serde::de::Error::duplicate_field("escapeChar")); + GeneratedField::Body => { + if body__.is_some() { + return Err(serde::de::Error::duplicate_field("body")); } - escape_char__ = Some(map_.next_value()?); + body__ = map_.next_value()?; } } } - Ok(LikeNode { - negated: negated__.unwrap_or_default(), - expr: expr__, - pattern: pattern__, - escape_char: escape_char__.unwrap_or_default(), + Ok(Lambda { + params: params__.unwrap_or_default(), + body: body__, }) } } - deserializer.deserialize_struct("datafusion.LikeNode", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.Lambda", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for LimitNode { +impl serde::Serialize for LambdaVariable { #[allow(deprecated)] fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> where @@ -11472,49 +11568,37 @@ impl serde::Serialize for LimitNode { { use serde::ser::SerializeStruct; let mut len = 0; - if self.input.is_some() { - len += 1; - } - if self.skip != 0 { + if !self.name.is_empty() { len += 1; } - if self.fetch != 0 { + if self.field.is_some() { len += 1; } - let mut struct_ser = serializer.serialize_struct("datafusion.LimitNode", len)?; - if let Some(v) = self.input.as_ref() { - struct_ser.serialize_field("input", v)?; - } - if self.skip != 0 { - #[allow(clippy::needless_borrow)] - #[allow(clippy::needless_borrows_for_generic_args)] - struct_ser.serialize_field("skip", ToString::to_string(&self.skip).as_str())?; + let mut struct_ser = serializer.serialize_struct("datafusion.LambdaVariable", len)?; + if !self.name.is_empty() { + struct_ser.serialize_field("name", &self.name)?; } - if self.fetch != 0 { - #[allow(clippy::needless_borrow)] - #[allow(clippy::needless_borrows_for_generic_args)] - struct_ser.serialize_field("fetch", ToString::to_string(&self.fetch).as_str())?; + if let Some(v) = self.field.as_ref() { + struct_ser.serialize_field("field", v)?; } struct_ser.end() } } -impl<'de> serde::Deserialize<'de> for LimitNode { +impl<'de> serde::Deserialize<'de> for LambdaVariable { #[allow(deprecated)] fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error> where D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "input", - "skip", - "fetch", + "name", + "field", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { - Input, - Skip, - Fetch, + Name, + Field, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error> @@ -11536,9 +11620,8 @@ impl<'de> serde::Deserialize<'de> for LimitNode { E: serde::de::Error, { match value { - "input" => Ok(GeneratedField::Input), - "skip" => Ok(GeneratedField::Skip), - "fetch" => Ok(GeneratedField::Fetch), + "name" => Ok(GeneratedField::Name), + "field" => Ok(GeneratedField::Field), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -11548,48 +11631,312 @@ impl<'de> serde::Deserialize<'de> for LimitNode { } struct GeneratedVisitor; impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = LimitNode; + type Value = LambdaVariable; fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.LimitNode") + formatter.write_str("struct datafusion.LambdaVariable") } - fn visit_map<V>(self, mut map_: V) -> std::result::Result<LimitNode, V::Error> + fn visit_map<V>(self, mut map_: V) -> std::result::Result<LambdaVariable, V::Error> where V: serde::de::MapAccess<'de>, { - let mut input__ = None; - let mut skip__ = None; - let mut fetch__ = None; + let mut name__ = None; + let mut field__ = None; while let Some(k) = map_.next_key()? { match k { - GeneratedField::Input => { - if input__.is_some() { - return Err(serde::de::Error::duplicate_field("input")); - } - input__ = map_.next_value()?; - } - GeneratedField::Skip => { - if skip__.is_some() { - return Err(serde::de::Error::duplicate_field("skip")); + GeneratedField::Name => { + if name__.is_some() { + return Err(serde::de::Error::duplicate_field("name")); } - skip__ = - Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) - ; + name__ = Some(map_.next_value()?); } - GeneratedField::Fetch => { - if fetch__.is_some() { - return Err(serde::de::Error::duplicate_field("fetch")); + GeneratedField::Field => { + if field__.is_some() { + return Err(serde::de::Error::duplicate_field("field")); } - fetch__ = - Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) - ; + field__ = map_.next_value()?; } } } - Ok(LimitNode { - input: input__, - skip: skip__.unwrap_or_default(), + Ok(LambdaVariable { + name: name__.unwrap_or_default(), + field: field__, + }) + } + } + deserializer.deserialize_struct("datafusion.LambdaVariable", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for LikeNode { + #[allow(deprecated)] + fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.negated { + len += 1; + } + if self.expr.is_some() { + len += 1; + } + if self.pattern.is_some() { + len += 1; + } + if !self.escape_char.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.LikeNode", len)?; + if self.negated { + struct_ser.serialize_field("negated", &self.negated)?; + } + if let Some(v) = self.expr.as_ref() { + struct_ser.serialize_field("expr", v)?; + } + if let Some(v) = self.pattern.as_ref() { + struct_ser.serialize_field("pattern", v)?; + } + if !self.escape_char.is_empty() { + struct_ser.serialize_field("escapeChar", &self.escape_char)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for LikeNode { + #[allow(deprecated)] + fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error> + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "negated", + "expr", + "pattern", + "escape_char", + "escapeChar", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Negated, + Expr, + Pattern, + EscapeChar, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error> + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str<E>(self, value: &str) -> std::result::Result<GeneratedField, E> + where + E: serde::de::Error, + { + match value { + "negated" => Ok(GeneratedField::Negated), + "expr" => Ok(GeneratedField::Expr), + "pattern" => Ok(GeneratedField::Pattern), + "escapeChar" | "escape_char" => Ok(GeneratedField::EscapeChar), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = LikeNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.LikeNode") + } + + fn visit_map<V>(self, mut map_: V) -> std::result::Result<LikeNode, V::Error> + where + V: serde::de::MapAccess<'de>, + { + let mut negated__ = None; + let mut expr__ = None; + let mut pattern__ = None; + let mut escape_char__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Negated => { + if negated__.is_some() { + return Err(serde::de::Error::duplicate_field("negated")); + } + negated__ = Some(map_.next_value()?); + } + GeneratedField::Expr => { + if expr__.is_some() { + return Err(serde::de::Error::duplicate_field("expr")); + } + expr__ = map_.next_value()?; + } + GeneratedField::Pattern => { + if pattern__.is_some() { + return Err(serde::de::Error::duplicate_field("pattern")); + } + pattern__ = map_.next_value()?; + } + GeneratedField::EscapeChar => { + if escape_char__.is_some() { + return Err(serde::de::Error::duplicate_field("escapeChar")); + } + escape_char__ = Some(map_.next_value()?); + } + } + } + Ok(LikeNode { + negated: negated__.unwrap_or_default(), + expr: expr__, + pattern: pattern__, + escape_char: escape_char__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.LikeNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for LimitNode { + #[allow(deprecated)] + fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.input.is_some() { + len += 1; + } + if self.skip != 0 { + len += 1; + } + if self.fetch != 0 { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.LimitNode", len)?; + if let Some(v) = self.input.as_ref() { + struct_ser.serialize_field("input", v)?; + } + if self.skip != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("skip", ToString::to_string(&self.skip).as_str())?; + } + if self.fetch != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("fetch", ToString::to_string(&self.fetch).as_str())?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for LimitNode { + #[allow(deprecated)] + fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error> + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "input", + "skip", + "fetch", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Input, + Skip, + Fetch, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error> + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str<E>(self, value: &str) -> std::result::Result<GeneratedField, E> + where + E: serde::de::Error, + { + match value { + "input" => Ok(GeneratedField::Input), + "skip" => Ok(GeneratedField::Skip), + "fetch" => Ok(GeneratedField::Fetch), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = LimitNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.LimitNode") + } + + fn visit_map<V>(self, mut map_: V) -> std::result::Result<LimitNode, V::Error> + where + V: serde::de::MapAccess<'de>, + { + let mut input__ = None; + let mut skip__ = None; + let mut fetch__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Input => { + if input__.is_some() { + return Err(serde::de::Error::duplicate_field("input")); + } + input__ = map_.next_value()?; + } + GeneratedField::Skip => { + if skip__.is_some() { + return Err(serde::de::Error::duplicate_field("skip")); + } + skip__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + GeneratedField::Fetch => { + if fetch__.is_some() { + return Err(serde::de::Error::duplicate_field("fetch")); + } + fetch__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + } + } + Ok(LimitNode { + input: input__, + skip: skip__.unwrap_or_default(), fetch: fetch__.unwrap_or_default(), }) } @@ -12525,6 +12872,15 @@ impl serde::Serialize for LogicalExprNode { logical_expr_node::ExprType::ScalarSubqueryExpr(v) => { struct_ser.serialize_field("scalarSubqueryExpr", v)?; } + logical_expr_node::ExprType::HigherOrderUdfExpr(v) => { + struct_ser.serialize_field("higherOrderUdfExpr", v)?; + } + logical_expr_node::ExprType::Lambda(v) => { + struct_ser.serialize_field("lambda", v)?; + } + logical_expr_node::ExprType::LambdaVariable(v) => { + struct_ser.serialize_field("lambdaVariable", v)?; + } } } struct_ser.end() @@ -12588,6 +12944,11 @@ impl<'de> serde::Deserialize<'de> for LogicalExprNode { "unnest", "scalar_subquery_expr", "scalarSubqueryExpr", + "higher_order_udf_expr", + "higherOrderUdfExpr", + "lambda", + "lambda_variable", + "lambdaVariable", ]; #[allow(clippy::enum_variant_names)] @@ -12624,6 +12985,9 @@ impl<'de> serde::Deserialize<'de> for LogicalExprNode { Placeholder, Unnest, ScalarSubqueryExpr, + HigherOrderUdfExpr, + Lambda, + LambdaVariable, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error> @@ -12677,6 +13041,9 @@ impl<'de> serde::Deserialize<'de> for LogicalExprNode { "placeholder" => Ok(GeneratedField::Placeholder), "unnest" => Ok(GeneratedField::Unnest), "scalarSubqueryExpr" | "scalar_subquery_expr" => Ok(GeneratedField::ScalarSubqueryExpr), + "higherOrderUdfExpr" | "higher_order_udf_expr" => Ok(GeneratedField::HigherOrderUdfExpr), + "lambda" => Ok(GeneratedField::Lambda), + "lambdaVariable" | "lambda_variable" => Ok(GeneratedField::LambdaVariable), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -12921,6 +13288,27 @@ impl<'de> serde::Deserialize<'de> for LogicalExprNode { return Err(serde::de::Error::duplicate_field("scalarSubqueryExpr")); } expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_expr_node::ExprType::ScalarSubqueryExpr) +; + } + GeneratedField::HigherOrderUdfExpr => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("higherOrderUdfExpr")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_expr_node::ExprType::HigherOrderUdfExpr) +; + } + GeneratedField::Lambda => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("lambda")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_expr_node::ExprType::Lambda) +; + } + GeneratedField::LambdaVariable => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("lambdaVariable")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_expr_node::ExprType::LambdaVariable) ; } } @@ -17982,6 +18370,15 @@ impl serde::Serialize for PhysicalExprNode { physical_expr_node::ExprType::DynamicFilter(v) => { struct_ser.serialize_field("dynamicFilter", v)?; } + physical_expr_node::ExprType::HigherOrderUdf(v) => { + struct_ser.serialize_field("higherOrderUdf", v)?; + } + physical_expr_node::ExprType::Lambda(v) => { + struct_ser.serialize_field("lambda", v)?; + } + physical_expr_node::ExprType::LambdaVariable(v) => { + struct_ser.serialize_field("lambdaVariable", v)?; + } } } struct_ser.end() @@ -18032,6 +18429,11 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { "scalarSubquery", "dynamic_filter", "dynamicFilter", + "higher_order_udf", + "higherOrderUdf", + "lambda", + "lambda_variable", + "lambdaVariable", ]; #[allow(clippy::enum_variant_names)] @@ -18058,6 +18460,9 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { HashExpr, ScalarSubquery, DynamicFilter, + HigherOrderUdf, + Lambda, + LambdaVariable, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error> @@ -18101,6 +18506,9 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { "hashExpr" | "hash_expr" => Ok(GeneratedField::HashExpr), "scalarSubquery" | "scalar_subquery" => Ok(GeneratedField::ScalarSubquery), "dynamicFilter" | "dynamic_filter" => Ok(GeneratedField::DynamicFilter), + "higherOrderUdf" | "higher_order_udf" => Ok(GeneratedField::HigherOrderUdf), + "lambda" => Ok(GeneratedField::Lambda), + "lambdaVariable" | "lambda_variable" => Ok(GeneratedField::LambdaVariable), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -18277,6 +18685,27 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { return Err(serde::de::Error::duplicate_field("dynamicFilter")); } expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::DynamicFilter) +; + } + GeneratedField::HigherOrderUdf => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("higherOrderUdf")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::HigherOrderUdf) +; + } + GeneratedField::Lambda => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("lambda")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::Lambda) +; + } + GeneratedField::LambdaVariable => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("lambdaVariable")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::LambdaVariable) ; } } @@ -18634,17 +19063,131 @@ impl<'de> serde::Deserialize<'de> for PhysicalHashExprNode { } } } - Ok(PhysicalHashExprNode { - on_columns: on_columns__.unwrap_or_default(), - seed0: seed0__.unwrap_or_default(), - description: description__.unwrap_or_default(), + Ok(PhysicalHashExprNode { + on_columns: on_columns__.unwrap_or_default(), + seed0: seed0__.unwrap_or_default(), + description: description__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalHashExprNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for PhysicalHashRepartition { + #[allow(deprecated)] + fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.hash_expr.is_empty() { + len += 1; + } + if self.partition_count != 0 { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalHashRepartition", len)?; + if !self.hash_expr.is_empty() { + struct_ser.serialize_field("hashExpr", &self.hash_expr)?; + } + if self.partition_count != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("partitionCount", ToString::to_string(&self.partition_count).as_str())?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalHashRepartition { + #[allow(deprecated)] + fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error> + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "hash_expr", + "hashExpr", + "partition_count", + "partitionCount", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + HashExpr, + PartitionCount, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error> + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str<E>(self, value: &str) -> std::result::Result<GeneratedField, E> + where + E: serde::de::Error, + { + match value { + "hashExpr" | "hash_expr" => Ok(GeneratedField::HashExpr), + "partitionCount" | "partition_count" => Ok(GeneratedField::PartitionCount), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalHashRepartition; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalHashRepartition") + } + + fn visit_map<V>(self, mut map_: V) -> std::result::Result<PhysicalHashRepartition, V::Error> + where + V: serde::de::MapAccess<'de>, + { + let mut hash_expr__ = None; + let mut partition_count__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::HashExpr => { + if hash_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("hashExpr")); + } + hash_expr__ = Some(map_.next_value()?); + } + GeneratedField::PartitionCount => { + if partition_count__.is_some() { + return Err(serde::de::Error::duplicate_field("partitionCount")); + } + partition_count__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + } + } + Ok(PhysicalHashRepartition { + hash_expr: hash_expr__.unwrap_or_default(), + partition_count: partition_count__.unwrap_or_default(), }) } } - deserializer.deserialize_struct("datafusion.PhysicalHashExprNode", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.PhysicalHashRepartition", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for PhysicalHashRepartition { +impl serde::Serialize for PhysicalHigherOrderUdfNode { #[allow(deprecated)] fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> where @@ -18652,41 +19195,48 @@ impl serde::Serialize for PhysicalHashRepartition { { use serde::ser::SerializeStruct; let mut len = 0; - if !self.hash_expr.is_empty() { + if !self.name.is_empty() { len += 1; } - if self.partition_count != 0 { + if !self.args.is_empty() { len += 1; } - let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalHashRepartition", len)?; - if !self.hash_expr.is_empty() { - struct_ser.serialize_field("hashExpr", &self.hash_expr)?; + if self.fun_definition.is_some() { + len += 1; } - if self.partition_count != 0 { + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalHigherOrderUdfNode", len)?; + if !self.name.is_empty() { + struct_ser.serialize_field("name", &self.name)?; + } + if !self.args.is_empty() { + struct_ser.serialize_field("args", &self.args)?; + } + if let Some(v) = self.fun_definition.as_ref() { #[allow(clippy::needless_borrow)] #[allow(clippy::needless_borrows_for_generic_args)] - struct_ser.serialize_field("partitionCount", ToString::to_string(&self.partition_count).as_str())?; + struct_ser.serialize_field("funDefinition", pbjson::private::base64::encode(&v).as_str())?; } struct_ser.end() } } -impl<'de> serde::Deserialize<'de> for PhysicalHashRepartition { +impl<'de> serde::Deserialize<'de> for PhysicalHigherOrderUdfNode { #[allow(deprecated)] fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error> where D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "hash_expr", - "hashExpr", - "partition_count", - "partitionCount", + "name", + "args", + "fun_definition", + "funDefinition", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { - HashExpr, - PartitionCount, + Name, + Args, + FunDefinition, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error> @@ -18708,8 +19258,9 @@ impl<'de> serde::Deserialize<'de> for PhysicalHashRepartition { E: serde::de::Error, { match value { - "hashExpr" | "hash_expr" => Ok(GeneratedField::HashExpr), - "partitionCount" | "partition_count" => Ok(GeneratedField::PartitionCount), + "name" => Ok(GeneratedField::Name), + "args" => Ok(GeneratedField::Args), + "funDefinition" | "fun_definition" => Ok(GeneratedField::FunDefinition), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -18719,43 +19270,51 @@ impl<'de> serde::Deserialize<'de> for PhysicalHashRepartition { } struct GeneratedVisitor; impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = PhysicalHashRepartition; + type Value = PhysicalHigherOrderUdfNode; fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.PhysicalHashRepartition") + formatter.write_str("struct datafusion.PhysicalHigherOrderUdfNode") } - fn visit_map<V>(self, mut map_: V) -> std::result::Result<PhysicalHashRepartition, V::Error> + fn visit_map<V>(self, mut map_: V) -> std::result::Result<PhysicalHigherOrderUdfNode, V::Error> where V: serde::de::MapAccess<'de>, { - let mut hash_expr__ = None; - let mut partition_count__ = None; + let mut name__ = None; + let mut args__ = None; + let mut fun_definition__ = None; while let Some(k) = map_.next_key()? { match k { - GeneratedField::HashExpr => { - if hash_expr__.is_some() { - return Err(serde::de::Error::duplicate_field("hashExpr")); + GeneratedField::Name => { + if name__.is_some() { + return Err(serde::de::Error::duplicate_field("name")); } - hash_expr__ = Some(map_.next_value()?); + name__ = Some(map_.next_value()?); } - GeneratedField::PartitionCount => { - if partition_count__.is_some() { - return Err(serde::de::Error::duplicate_field("partitionCount")); + GeneratedField::Args => { + if args__.is_some() { + return Err(serde::de::Error::duplicate_field("args")); } - partition_count__ = - Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + args__ = Some(map_.next_value()?); + } + GeneratedField::FunDefinition => { + if fun_definition__.is_some() { + return Err(serde::de::Error::duplicate_field("funDefinition")); + } + fun_definition__ = + map_.next_value::<::std::option::Option<::pbjson::private::BytesDeserialize<_>>>()?.map(|x| x.0) ; } } } - Ok(PhysicalHashRepartition { - hash_expr: hash_expr__.unwrap_or_default(), - partition_count: partition_count__.unwrap_or_default(), + Ok(PhysicalHigherOrderUdfNode { + name: name__.unwrap_or_default(), + args: args__.unwrap_or_default(), + fun_definition: fun_definition__, }) } } - deserializer.deserialize_struct("datafusion.PhysicalHashRepartition", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.PhysicalHigherOrderUdfNode", FIELDS, GeneratedVisitor) } } impl serde::Serialize for PhysicalInListNode { @@ -19065,6 +19624,224 @@ impl<'de> serde::Deserialize<'de> for PhysicalIsNull { deserializer.deserialize_struct("datafusion.PhysicalIsNull", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for PhysicalLambdaExprNode { + #[allow(deprecated)] + fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.params.is_empty() { + len += 1; + } + if self.body.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalLambdaExprNode", len)?; + if !self.params.is_empty() { + struct_ser.serialize_field("params", &self.params)?; + } + if let Some(v) = self.body.as_ref() { + struct_ser.serialize_field("body", v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalLambdaExprNode { + #[allow(deprecated)] + fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error> + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "params", + "body", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Params, + Body, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error> + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str<E>(self, value: &str) -> std::result::Result<GeneratedField, E> + where + E: serde::de::Error, + { + match value { + "params" => Ok(GeneratedField::Params), + "body" => Ok(GeneratedField::Body), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalLambdaExprNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalLambdaExprNode") + } + + fn visit_map<V>(self, mut map_: V) -> std::result::Result<PhysicalLambdaExprNode, V::Error> + where + V: serde::de::MapAccess<'de>, + { + let mut params__ = None; + let mut body__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Params => { + if params__.is_some() { + return Err(serde::de::Error::duplicate_field("params")); + } + params__ = Some(map_.next_value()?); + } + GeneratedField::Body => { + if body__.is_some() { + return Err(serde::de::Error::duplicate_field("body")); + } + body__ = map_.next_value()?; + } + } + } + Ok(PhysicalLambdaExprNode { + params: params__.unwrap_or_default(), + body: body__, + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalLambdaExprNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for PhysicalLambdaVariableExprNode { + #[allow(deprecated)] + fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.index != 0 { + len += 1; + } + if self.field.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalLambdaVariableExprNode", len)?; + if self.index != 0 { + struct_ser.serialize_field("index", &self.index)?; + } + if let Some(v) = self.field.as_ref() { + struct_ser.serialize_field("field", v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalLambdaVariableExprNode { + #[allow(deprecated)] + fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error> + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "index", + "field", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Index, + Field, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error> + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str<E>(self, value: &str) -> std::result::Result<GeneratedField, E> + where + E: serde::de::Error, + { + match value { + "index" => Ok(GeneratedField::Index), + "field" => Ok(GeneratedField::Field), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalLambdaVariableExprNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalLambdaVariableExprNode") + } + + fn visit_map<V>(self, mut map_: V) -> std::result::Result<PhysicalLambdaVariableExprNode, V::Error> + where + V: serde::de::MapAccess<'de>, + { + let mut index__ = None; + let mut field__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Index => { + if index__.is_some() { + return Err(serde::de::Error::duplicate_field("index")); + } + index__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + GeneratedField::Field => { + if field__.is_some() { + return Err(serde::de::Error::duplicate_field("field")); + } + field__ = map_.next_value()?; + } + } + } + Ok(PhysicalLambdaVariableExprNode { + index: index__.unwrap_or_default(), + field: field__, + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalLambdaVariableExprNode", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for PhysicalLikeExprNode { #[allow(deprecated)] fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 675ead23f4..db51edfd5d 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -721,7 +721,7 @@ pub struct SubqueryAliasNode { pub struct LogicalExprNode { #[prost( oneof = "logical_expr_node::ExprType", - tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 13, 14, 15, 17, 18, 19, 20, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36" + tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 13, 14, 15, 17, 18, 19, 20, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39" )] pub expr_type: ::core::option::Option<logical_expr_node::ExprType>, } @@ -802,6 +802,12 @@ pub mod logical_expr_node { /// Subquery expressions #[prost(message, tag = "36")] ScalarSubqueryExpr(::prost::alloc::boxed::Box<super::ScalarSubqueryExprNode>), + #[prost(message, tag = "37")] + HigherOrderUdfExpr(super::HigherOrderUdfExprNode), + #[prost(message, tag = "38")] + Lambda(::prost::alloc::boxed::Box<super::Lambda>), + #[prost(message, tag = "39")] + LambdaVariable(super::LambdaVariable), } } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] @@ -991,6 +997,29 @@ pub struct ScalarUdfExprNode { pub fun_definition: ::core::option::Option<::prost::alloc::vec::Vec<u8>>, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct HigherOrderUdfExprNode { + #[prost(string, tag = "1")] + pub fun_name: ::prost::alloc::string::String, + #[prost(message, repeated, tag = "2")] + pub args: ::prost::alloc::vec::Vec<LogicalExprNode>, + #[prost(bytes = "vec", optional, tag = "3")] + pub fun_definition: ::core::option::Option<::prost::alloc::vec::Vec<u8>>, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct Lambda { + #[prost(string, repeated, tag = "1")] + pub params: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + #[prost(message, optional, boxed, tag = "2")] + pub body: ::core::option::Option<::prost::alloc::boxed::Box<LogicalExprNode>>, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct LambdaVariable { + #[prost(string, tag = "1")] + pub name: ::prost::alloc::string::String, + #[prost(message, optional, tag = "2")] + pub field: ::core::option::Option<super::datafusion_common::Field>, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct WindowExprNode { #[prost(message, repeated, tag = "4")] pub exprs: ::prost::alloc::vec::Vec<LogicalExprNode>, @@ -1475,7 +1504,7 @@ pub struct PhysicalExprNode { pub expr_id: ::core::option::Option<u64>, #[prost( oneof = "physical_expr_node::ExprType", - tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15, 16, 18, 19, 20, 21, 22, 23" + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15, 16, 18, 19, 20, 21, 22, 23, 24, 25, 26" )] pub expr_type: ::core::option::Option<physical_expr_node::ExprType>, } @@ -1532,6 +1561,12 @@ pub mod physical_expr_node { ScalarSubquery(super::PhysicalScalarSubqueryExprNode), #[prost(message, tag = "23")] DynamicFilter(::prost::alloc::boxed::Box<super::PhysicalDynamicFilterNode>), + #[prost(message, tag = "24")] + HigherOrderUdf(super::PhysicalHigherOrderUdfNode), + #[prost(message, tag = "25")] + Lambda(::prost::alloc::boxed::Box<super::PhysicalLambdaExprNode>), + #[prost(message, tag = "26")] + LambdaVariable(super::PhysicalLambdaVariableExprNode), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -1563,6 +1598,29 @@ pub struct PhysicalScalarUdfNode { pub return_field_name: ::prost::alloc::string::String, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalHigherOrderUdfNode { + #[prost(string, tag = "1")] + pub name: ::prost::alloc::string::String, + #[prost(message, repeated, tag = "2")] + pub args: ::prost::alloc::vec::Vec<PhysicalExprNode>, + #[prost(bytes = "vec", optional, tag = "3")] + pub fun_definition: ::core::option::Option<::prost::alloc::vec::Vec<u8>>, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalLambdaExprNode { + #[prost(string, repeated, tag = "1")] + pub params: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + #[prost(message, optional, boxed, tag = "2")] + pub body: ::core::option::Option<::prost::alloc::boxed::Box<PhysicalExprNode>>, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalLambdaVariableExprNode { + #[prost(uint32, tag = "1")] + pub index: u32, + #[prost(message, optional, tag = "2")] + pub field: ::core::option::Option<super::datafusion_common::Field>, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalAggregateExprNode { #[prost(message, repeated, tag = "2")] pub expr: ::prost::alloc::vec::Vec<PhysicalExprNode>, diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index 372213e38f..6d9a73e06f 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -28,7 +28,9 @@ use datafusion_execution::registry::FunctionRegistry; use datafusion_expr::dml::{ InsertOp, MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, MergeIntoOp, }; -use datafusion_expr::expr::{Alias, NullTreatment, Placeholder, Sort}; +use datafusion_expr::expr::{ + Alias, Lambda, LambdaVariable, NullTreatment, Placeholder, Sort, +}; use datafusion_expr::expr::{Unnest, WildcardOptions}; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::{ @@ -718,6 +720,22 @@ pub fn parse_expr( parse_exprs(args, ctx, codec)?, ))) } + ExprType::HigherOrderUdfExpr(protobuf::HigherOrderUdfExprNode { + fun_name, + args, + fun_definition, + }) => { + let hof_fn = match fun_definition { + Some(buf) => codec.try_decode_higher_order_function(fun_name, buf)?, + None => ctx + .higher_order_function(fun_name.as_str()) + .or_else(|_| codec.try_decode_higher_order_function(fun_name, &[]))?, + }; + Ok(Expr::HigherOrderFunction(expr::HigherOrderFunction::new( + hof_fn, + parse_exprs(args, ctx, codec)?, + ))) + } ExprType::AggregateUdfExpr(pb) => { let agg_fn = match &pb.fun_definition { Some(buf) => codec.try_decode_udaf(&pb.fun_name, buf)?, @@ -791,6 +809,16 @@ pub fn parse_expr( )?; Ok(Expr::ScalarSubquery(subquery)) } + ExprType::Lambda(lambda) => Ok(Expr::Lambda(Lambda::new( + lambda.params.clone(), + parse_required_expr(lambda.body.as_deref(), ctx, "body", codec)?, + ))), + ExprType::LambdaVariable(lambda_variable) => { + Ok(Expr::LambdaVariable(LambdaVariable::new( + lambda_variable.name.clone(), + lambda_variable.field.as_ref().optional()?.map(Arc::new), + ))) + } } } diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index b2b035af88..23ce254e99 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -26,8 +26,9 @@ use datafusion_expr::dml::{ MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, MergeIntoOp, }; use datafusion_expr::expr::{ - self, AggregateFunctionParams, Alias, Between, BinaryExpr, Cast, GroupingSet, InList, - Like, NullTreatment, Placeholder, ScalarFunction, Unnest, + self, AggregateFunctionParams, Alias, Between, BinaryExpr, Cast, GroupingSet, + HigherOrderFunction, InList, Lambda, LambdaVariable, Like, NullTreatment, + Placeholder, ScalarFunction, Unnest, }; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::{ @@ -413,6 +414,19 @@ pub fn serialize_expr( })), } } + Expr::HigherOrderFunction(HigherOrderFunction { func, args }) => { + let mut buf = Vec::new(); + let _ = codec.try_encode_higher_order_function(func.as_ref(), &mut buf); + protobuf::LogicalExprNode { + expr_type: Some(ExprType::HigherOrderUdfExpr( + protobuf::HigherOrderUdfExprNode { + fun_name: func.name().to_string(), + fun_definition: (!buf.is_empty()).then_some(buf), + args: serialize_exprs(args, codec)?, + }, + )), + } + } Expr::Not(expr) => { let expr = Box::new(protobuf::Not { expr: Some(Box::new(serialize_expr(expr.as_ref(), codec)?)), @@ -640,11 +654,22 @@ pub fn serialize_expr( .unwrap_or(HashMap::new()), })), }, - Expr::HigherOrderFunction(_) | Expr::Lambda(_) | Expr::LambdaVariable(_) => { - return Err(Error::General( - "Proto serialization error: Lambda not implemented".to_string(), - )); - } + Expr::Lambda(Lambda { params, body }) => protobuf::LogicalExprNode { + expr_type: Some(ExprType::Lambda(Box::new(protobuf::Lambda { + params: params.clone(), + body: Some(Box::new(serialize_expr(body, codec)?)), + }))), + }, + Expr::LambdaVariable(LambdaVariable { + name, + field, + spans: _, + }) => protobuf::LogicalExprNode { + expr_type: Some(ExprType::LambdaVariable(protobuf::LambdaVariable { + name: name.clone(), + field: field.as_deref().map(|v| v.try_into()).transpose()?, + })), + }, }; Ok(expr_node) diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 19e49f3cb8..b908b504bb 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -43,9 +43,12 @@ use datafusion_execution::{FunctionRegistry, TaskContext}; use datafusion_expr::WindowFunctionDefinition; use datafusion_expr::dml::InsertOp; use datafusion_expr::execution_props::SubqueryIndex; +use datafusion_physical_expr::expressions::{LambdaExpr, LambdaVariable}; use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; -use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, ScalarFunctionExpr}; +use datafusion_physical_expr::{ + HigherOrderFunctionExpr, LexOrdering, PhysicalSortExpr, ScalarFunctionExpr, +}; use datafusion_physical_plan::expressions::{ BinaryExpr, CaseExpr, CastExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, LikeExpr, Literal, NegativeExpr, NotExpr, TryCastExpr, UnKnownColumn, @@ -336,6 +339,31 @@ pub fn parse_physical_expr_with_converter( .with_nullable(e.nullable), ) } + ExprType::HigherOrderUdf(e) => { + let func = match &e.fun_definition { + Some(buf) => { + ctx.codec().try_decode_higher_order_function(&e.name, buf)? + } + None => ctx + .task_ctx() + .higher_order_function(e.name.as_str()) + .or_else(|_| { + ctx.codec().try_decode_higher_order_function(&e.name, &[]) + })?, + }; + let func_def = Arc::clone(&func); + + let args = parse_physical_exprs(&e.args, ctx, input_schema, proto_converter)?; + + let config_options = Arc::clone(ctx.task_ctx().session_config().options()); + + Arc::new(HigherOrderFunctionExpr::try_new_with_schema( + func_def, + args, + input_schema, + config_options, + )?) + } ExprType::LikeExpr(_) => LikeExpr::try_from_proto(proto, &decode_ctx)?, ExprType::HashExpr(_) => HashExpr::try_from_proto(proto, &decode_ctx)?, ExprType::ScalarSubquery(sq) => { @@ -371,6 +399,10 @@ pub fn parse_physical_expr_with_converter( ctx.codec() .try_decode_expr(extension.expr.as_slice(), &inputs)? as _ } + ExprType::Lambda(_) => LambdaExpr::try_from_proto(proto, &decode_ctx)?, + ExprType::LambdaVariable(_) => { + LambdaVariable::try_from_proto(proto, &decode_ctx)? + } }; Ok(pexpr) diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 72f6e5af5b..0744a94dce 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -54,7 +54,7 @@ use datafusion_datasource_parquet::source::ParquetSource; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{FunctionRegistry, TaskContext}; use datafusion_expr::execution_props::{ScalarSubqueryResults, SubqueryIndex}; -use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; +use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, WindowUDF}; use datafusion_functions_table::generate_series::{ Empty, GenSeriesArgs, GenerateSeriesTable, GenericSeriesState, TimestampValue, }; @@ -3946,6 +3946,24 @@ pub trait PhysicalExtensionCodec: Debug + Send + Sync + Any { Ok(()) } + fn try_decode_higher_order_function( + &self, + name: &str, + _buf: &[u8], + ) -> Result<Arc<HigherOrderUDF>> { + not_impl_err!( + "PhysicalExtensionCodec is not provided for higher order function {name}" + ) + } + + fn try_encode_higher_order_function( + &self, + _node: &HigherOrderUDF, + _buf: &mut Vec<u8>, + ) -> Result<()> { + Ok(()) + } + fn try_decode_expr( &self, _buf: &[u8], diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index e93bb9cc5f..4025b580e8 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -31,9 +31,9 @@ use datafusion_datasource_json::file_format::JsonSink; #[cfg(feature = "parquet")] use datafusion_datasource_parquet::file_format::ParquetSink; use datafusion_expr::WindowFrame; -use datafusion_physical_expr::ScalarFunctionExpr; use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; use datafusion_physical_expr::window::{SlidingAggregateWindowExpr, StandardWindowExpr}; +use datafusion_physical_expr::{HigherOrderFunctionExpr, ScalarFunctionExpr}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_physical_plan::udaf::AggregateFunctionExpr; use datafusion_physical_plan::windows::{PlainAggregateWindowExpr, WindowUDFExpr}; @@ -318,6 +318,19 @@ pub fn serialize_physical_expr_with_converter( }, )), }) + } else if let Some(expr) = expr.downcast_ref::<HigherOrderFunctionExpr>() { + let mut buf = Vec::new(); + codec.try_encode_higher_order_function(expr.fun(), &mut buf)?; + Ok(protobuf::PhysicalExprNode { + expr_id, + expr_type: Some(protobuf::physical_expr_node::ExprType::HigherOrderUdf( + protobuf::PhysicalHigherOrderUdfNode { + name: expr.name().to_string(), + args: serialize_physical_exprs(expr.args(), codec, proto_converter)?, + fun_definition: (!buf.is_empty()).then_some(buf), + }, + )), + }) } else if let Some(expr) = expr.downcast_ref::<ScalarSubqueryExpr>() { Ok(protobuf::PhysicalExprNode { expr_id, diff --git a/datafusion/proto/tests/cases/mod.rs b/datafusion/proto/tests/cases/mod.rs index 3abbaccf79..7a95ee0c29 100644 --- a/datafusion/proto/tests/cases/mod.rs +++ b/datafusion/proto/tests/cases/mod.rs @@ -21,8 +21,10 @@ use datafusion::physical_expr::PhysicalExpr; use datafusion_common::plan_err; use datafusion_expr::function::AccumulatorArgs; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, LimitEffect, PartitionEvaluator, ScalarFunctionArgs, - ScalarUDFImpl, Signature, Volatility, WindowUDFImpl, + Accumulator, AggregateUDFImpl, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, + HigherOrderSignature, HigherOrderUDFImpl, LambdaParametersProgress, LimitEffect, + PartitionEvaluator, ScalarFunctionArgs, ScalarUDFImpl, Signature, ValueOrLambda, + Volatility, WindowUDFImpl, }; use datafusion_functions_window_common::field::WindowUDFFieldArgs; use datafusion_functions_window_common::partition::PartitionEvaluatorArgs; @@ -180,3 +182,66 @@ pub(in crate::cases) struct CustomUDWFNode { #[prost(string, tag = "1")] pub payload: String, } + +#[derive(Debug, PartialEq, Eq, Hash)] +pub(in crate::cases) struct MyHigherOrderUDF { + signature: HigherOrderSignature, + pub payload: String, +} + +impl MyHigherOrderUDF { + pub fn new(payload: String) -> Self { + Self { + signature: HigherOrderSignature::any(2, Volatility::Immutable), + payload, + } + } +} + +impl HigherOrderUDFImpl for MyHigherOrderUDF { + fn name(&self) -> &str { + "higher_order_udf" + } + + fn signature(&self) -> &HigherOrderSignature { + &self.signature + } + + fn lambda_parameters( + &self, + _step: usize, + fields: &[ValueOrLambda<FieldRef, Option<FieldRef>>], + ) -> datafusion_common::Result<LambdaParametersProgress> { + let list = match fields.first() { + Some(ValueOrLambda::Value(field)) => field, + _ => return plan_err!("higher_order_udf expects a list as first argument"), + }; + let element = match list.data_type() { + DataType::List(field) | DataType::LargeList(field) => Arc::clone(field), + other => { + return plan_err!("higher_order_udf expected a list, got {other}"); + } + }; + Ok(LambdaParametersProgress::Complete(vec![vec![element]])) + } + + fn return_field_from_args( + &self, + _args: HigherOrderReturnFieldArgs, + ) -> datafusion_common::Result<FieldRef> { + Ok(Arc::new(Field::new("", DataType::Int64, true))) + } + + fn invoke_with_args( + &self, + _args: HigherOrderFunctionArgs, + ) -> datafusion_common::Result<ColumnarValue> { + unimplemented!() + } +} + +#[derive(Clone, PartialEq, ::prost::Message)] +pub(in crate::cases) struct MyHigherOrderUdfNode { + #[prost(string, tag = "1")] + pub payload: String, +} diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index 431b49dc8b..b3edc0f5ce 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -83,17 +83,17 @@ use datafusion_expr::dml::{ MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, MergeIntoOp, }; use datafusion_expr::expr::{ - self, Between, BinaryExpr, Case, Cast, GroupingSet, InList, Like, NullTreatment, - ScalarFunction, Unnest, WildcardOptions, + self, Between, BinaryExpr, Case, Cast, GroupingSet, InList, LambdaVariable, Like, + NullTreatment, ScalarFunction, Unnest, WildcardOptions, }; use datafusion_expr::logical_plan::{ ExplainOption, Extension, UserDefinedLogicalNodeCore, }; use datafusion_expr::{ Accumulator, AggregateUDF, ColumnarValue, DmlStatement, ExprFunctionExt, - ExprSchemable, LimitEffect, Literal, LogicalPlan, LogicalPlanBuilder, Operator, - PartitionEvaluator, RangePartitioning, Repartition, ScalarUDF, Signature, TryCast, - Volatility, WindowFrame, WindowFrameBound, WindowFrameUnits, + ExprSchemable, HigherOrderUDF, LimitEffect, Literal, LogicalPlan, LogicalPlanBuilder, + Operator, PartitionEvaluator, RangePartitioning, Repartition, ScalarUDF, Signature, + TryCast, Volatility, WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, WindowUDF, WindowUDFImpl, WriteOp, }; use datafusion_functions_aggregate::average::avg_udaf; @@ -118,7 +118,10 @@ use datafusion_proto::logical_plan::{ }; use datafusion_proto::protobuf; -use crate::cases::{MyAggregateUDF, MyAggregateUdfNode, MyRegexUdf, MyRegexUdfNode}; +use crate::cases::{ + MyAggregateUDF, MyAggregateUdfNode, MyHigherOrderUDF, MyHigherOrderUdfNode, + MyRegexUdf, MyRegexUdfNode, +}; #[cfg(feature = "json")] fn roundtrip_json_test(proto: &protobuf::LogicalExprNode) { @@ -1806,6 +1809,41 @@ impl LogicalExtensionCodec for UDFExtensionCodec { .map_err(|err| internal_datafusion_err!("failed to encode udf: {err}"))?; Ok(()) } + + fn try_decode_higher_order_function( + &self, + name: &str, + buf: &[u8], + ) -> Result<Arc<HigherOrderUDF>> { + if name == "higher_order_udf" { + let proto = MyHigherOrderUdfNode::decode(buf).map_err(|err| { + internal_datafusion_err!("failed to decode higher_order_udf: {err}") + })?; + + Ok(Arc::new(HigherOrderUDF::new_from_impl( + MyHigherOrderUDF::new(proto.payload), + ))) + } else { + not_impl_err!("unrecognized higher order UDF implementation, cannot decode") + } + } + + fn try_encode_higher_order_function( + &self, + node: &HigherOrderUDF, + buf: &mut Vec<u8>, + ) -> Result<()> { + let hof = (node.inner().as_ref() as &dyn Any) + .downcast_ref::<MyHigherOrderUDF>() + .unwrap(); + let proto = MyHigherOrderUdfNode { + payload: hof.payload.clone(), + }; + proto + .encode(buf) + .map_err(|err| internal_datafusion_err!("failed to encode hof: {err}"))?; + Ok(()) + } } #[test] @@ -2867,6 +2905,115 @@ fn roundtrip_aggregate_udf_extension_codec() { roundtrip_json_test(&proto); } +fn dummy_higher_order_function_args() -> Vec<Expr> { + let list = ScalarValue::List(ScalarValue::new_list_nullable( + &[ScalarValue::Int32(Some(1))], + &DataType::Int32, + )); + let lambda_var_with_field = Expr::LambdaVariable(LambdaVariable::new( + "x".to_string(), + Some(Arc::new(Field::new("x", DataType::Int32, true))), + )); + let lambda_var_without_field = + Expr::LambdaVariable(LambdaVariable::new("x".into(), None)); + let lambda = lambda(["x"], lambda_var_with_field + lambda_var_without_field); + vec![Expr::Literal(list, None), lambda] +} + +#[test] +fn roundtrip_higher_order_function() { + let hof = Arc::new(HigherOrderUDF::new_from_impl(MyHigherOrderUDF::new( + "payload".to_string(), + ))); + + let test_expr = Expr::HigherOrderFunction(expr::HigherOrderFunction::new( + Arc::clone(&hof), + dummy_higher_order_function_args(), + )); + + let ctx = SessionContext::new(); + ctx.register_higher_order_function(hof); + + roundtrip_expr_test(test_expr.clone(), ctx); + + // Now test loading the HOF without registering it in the context, but rather creating it + // in the extension codec. + #[derive(Debug)] + struct DummyHigherOrderUDFExtensionCodec; + + impl LogicalExtensionCodec for DummyHigherOrderUDFExtensionCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[LogicalPlan], + _ctx: &TaskContext, + ) -> Result<Extension> { + not_impl_err!("LogicalExtensionCodec is not provided") + } + + fn try_encode(&self, _node: &Extension, _buf: &mut Vec<u8>) -> Result<()> { + not_impl_err!("LogicalExtensionCodec is not provided") + } + + fn try_decode_table_provider( + &self, + _buf: &[u8], + _table_ref: &TableReference, + _schema: SchemaRef, + _ctx: &TaskContext, + ) -> Result<Arc<dyn TableProvider>> { + not_impl_err!("LogicalExtensionCodec is not provided") + } + + fn try_encode_table_provider( + &self, + _table_ref: &TableReference, + _node: Arc<dyn TableProvider>, + _buf: &mut Vec<u8>, + ) -> Result<()> { + not_impl_err!("LogicalExtensionCodec is not provided") + } + + fn try_decode_higher_order_function( + &self, + name: &str, + _buf: &[u8], + ) -> Result<Arc<HigherOrderUDF>> { + if name == "higher_order_udf" { + Ok(Arc::new(HigherOrderUDF::new_from_impl( + MyHigherOrderUDF::new("payload".to_string()), + ))) + } else { + Err(internal_datafusion_err!("HOF {name} not found")) + } + } + } + + let ctx = SessionContext::new(); + roundtrip_expr_test_with_codec(test_expr, ctx, &DummyHigherOrderUDFExtensionCodec) +} + +#[test] +fn roundtrip_higher_order_udf_extension_codec() { + let hof = Arc::new(HigherOrderUDF::new_from_impl(MyHigherOrderUDF::new( + "payload".to_string(), + ))); + + let test_expr = Expr::HigherOrderFunction(expr::HigherOrderFunction::new( + hof, + dummy_higher_order_function_args(), + )); + + let ctx = SessionContext::new(); + let proto = serialize_expr(&test_expr, &UDFExtensionCodec).expect("serialize expr"); + let round_trip = + from_proto::parse_expr(&proto, ctx.task_ctx().as_ref(), &UDFExtensionCodec) + .expect("parse expr"); + + assert_eq!(format!("{:?}", test_expr), format!("{round_trip:?}")); + roundtrip_json_test(&proto); +} + #[test] fn roundtrip_grouping_sets() { let test_expr = Expr::GroupingSet(GroupingSet::GroupingSets(vec![ diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index bb99a1ecbf..6ede6fc0e9 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -52,7 +52,7 @@ use datafusion::physical_expr::aggregate::{AggregateExprBuilder, AggregateFuncti use datafusion::physical_expr::expressions::Literal; use datafusion::physical_expr::window::{SlidingAggregateWindowExpr, StandardWindowExpr}; use datafusion::physical_expr::{ - LexOrdering, PhysicalSortRequirement, ScalarFunctionExpr, + HigherOrderFunctionExpr, LexOrdering, PhysicalSortRequirement, ScalarFunctionExpr, }; use datafusion::physical_optimizer::PhysicalOptimizerRule; use datafusion::physical_optimizer::filter_pushdown::FilterPushdown; @@ -108,7 +108,7 @@ use datafusion_datasource::{TableSchema, TableSchemaBuilder}; use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl}; use datafusion_expr::dml::InsertOp; use datafusion_expr::{ - Accumulator, AccumulatorFactoryFunction, AggregateUDF, ColumnarValue, + Accumulator, AccumulatorFactoryFunction, AggregateUDF, ColumnarValue, HigherOrderUDF, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, SimpleAggregateUDF, WindowFrame, WindowFrameBound, WindowUDF, execution_props::{ScalarSubqueryResults, SubqueryIndex}, @@ -140,9 +140,10 @@ use datafusion_proto::protobuf::{PhysicalExprNode, PhysicalPlanNode}; use prost::Message; use crate::cases::{ - CustomUDWF, CustomUDWFNode, MyAggregateUDF, MyAggregateUdfNode, MyRegexUdf, - MyRegexUdfNode, + CustomUDWF, CustomUDWFNode, MyAggregateUDF, MyAggregateUdfNode, MyHigherOrderUDF, + MyHigherOrderUdfNode, MyRegexUdf, MyRegexUdfNode, }; +use datafusion_physical_expr::expressions::{LambdaVariable, is_not_null, lambda}; use datafusion_physical_expr::utils::reassign_expr_columns; /// Perform a serde roundtrip and assert that the string representation of the before and after plans @@ -1472,6 +1473,42 @@ impl PhysicalExtensionCodec for UDFExtensionCodec { } Ok(()) } + + fn try_decode_higher_order_function( + &self, + name: &str, + buf: &[u8], + ) -> Result<Arc<HigherOrderUDF>> { + if name == "higher_order_udf" { + let proto = MyHigherOrderUdfNode::decode(buf).map_err(|err| { + internal_datafusion_err!("failed to decode higher_order_udf: {err}") + })?; + + Ok(Arc::new(HigherOrderUDF::new_from_impl( + MyHigherOrderUDF::new(proto.payload), + ))) + } else { + not_impl_err!("unrecognized higher order UDF implementation, cannot decode") + } + } + + fn try_encode_higher_order_function( + &self, + node: &HigherOrderUDF, + buf: &mut Vec<u8>, + ) -> Result<()> { + if let Some(hof) = (node.inner().as_ref() as &dyn std::any::Any) + .downcast_ref::<MyHigherOrderUDF>() + { + let proto = MyHigherOrderUdfNode { + payload: hof.payload.clone(), + }; + proto.encode(buf).map_err(|err| { + internal_datafusion_err!("failed to encode hof: {err:?}") + })?; + } + Ok(()) + } } #[test] @@ -1532,6 +1569,92 @@ fn roundtrip_scalar_udf_extension_codec() -> Result<()> { Ok(()) } +#[test] +fn roundtrip_higher_order_udf() -> Result<()> { + let element_field = Arc::new(Field::new("v", DataType::Int32, true)); + let list_field = Field::new( + "list_col", + DataType::List(Arc::clone(&element_field)), + false, + ); + let schema = Arc::new(Schema::new(vec![list_field])); + + let input = Arc::new(EmptyExec::new(schema.clone())); + + let hof = Arc::new(HigherOrderUDF::new_from_impl(MyHigherOrderUDF::new( + "payload".to_string(), + ))); + + let expr = HigherOrderFunctionExpr::try_new_with_schema( + Arc::clone(&hof), + vec![ + col("list_col", &schema)?, + lambda( + ["v"], + is_not_null(Arc::new(LambdaVariable::new(1, element_field)))?, + )?, + ], + &schema, + Arc::new(ConfigOptions::default()), + )?; + + let project = ProjectionExec::try_new( + vec![ProjectionExpr { + expr: Arc::new(expr), + alias: "a".to_string(), + }], + input, + )?; + + let ctx = SessionContext::new(); + ctx.register_higher_order_function(hof); + + roundtrip_test_with_context(Arc::new(project), &ctx) +} + +#[test] +fn roundtrip_higher_order_udf_extension_codec() -> Result<()> { + let element_field = Arc::new(Field::new("v", DataType::Int32, true)); + let list_field = Field::new( + "list_col", + DataType::List(Arc::clone(&element_field)), + false, + ); + let schema = Arc::new(Schema::new(vec![list_field])); + let input = Arc::new(EmptyExec::new(schema.clone())); + + let lambda_body = Arc::new(LambdaVariable::new(1, Arc::clone(&element_field))); + let lambda_expr = lambda(["v"], lambda_body)?; + + let hof = Arc::new(HigherOrderUDF::new_from_impl(MyHigherOrderUDF::new( + "payload".to_string(), + ))); + let hof_expr = Arc::new(HigherOrderFunctionExpr::try_new_with_schema( + hof, + vec![col("list_col", &schema)?, lambda_expr], + &schema, + Arc::new(ConfigOptions::default()), + )?); + + let project = ProjectionExec::try_new( + vec![ProjectionExpr { + expr: hof_expr, + alias: "out".to_string(), + }], + input, + )?; + + let ctx = SessionContext::new(); + let proto_converter = DefaultPhysicalProtoConverter {}; + roundtrip_test_and_return( + Arc::new(project), + &ctx, + &UDFExtensionCodec, + &proto_converter, + )?; + Ok(()) +} + #[test] fn roundtrip_udwf_extension_codec() -> Result<()> { let field_a = Field::new("a", DataType::Int64, false); diff --git a/datafusion/proto/tests/cases/serialize.rs b/datafusion/proto/tests/cases/serialize.rs index 850fd42ce1..a0a917e323 100644 --- a/datafusion/proto/tests/cases/serialize.rs +++ b/datafusion/proto/tests/cases/serialize.rs @@ -22,14 +22,18 @@ use arrow::datatypes::{DataType, Field}; use datafusion::execution::FunctionRegistry; use datafusion::prelude::SessionContext; -use datafusion_expr::expr::Placeholder; -use datafusion_expr::{ColumnarValue, col, create_udf, lit}; +use datafusion_common::ScalarValue; +use datafusion_expr::expr::{HigherOrderFunction, LambdaVariable, Placeholder}; +use datafusion_expr::{ColumnarValue, HigherOrderUDF, col, create_udf, lambda, lit}; use datafusion_expr::{Expr, Volatility}; use datafusion_functions::string; use datafusion_proto::bytes::Serializeable; use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec; +use datafusion_proto::logical_plan::from_proto::parse_expr; use datafusion_proto::logical_plan::to_proto::serialize_expr; +use crate::cases::MyHigherOrderUDF; + #[test] #[should_panic( expected = "Error decoding expr as protobuf: failed to decode Protobuf message" @@ -298,3 +302,92 @@ fn test_expression_serialization_roundtrip() { name.split('(').next().unwrap().to_string() } } + +/// return a `SessionContext` with `MyHigherOrderUDF` registered as a higher-order UDF +fn context_with_higher_order_function() -> SessionContext { + let ctx = SessionContext::new(); + let hof = Arc::new(HigherOrderUDF::new_from_impl(MyHigherOrderUDF::new( + "payload".to_string(), + ))); + ctx.register_higher_order_function(hof); + ctx +} + +fn dummy_higher_order_function_call(hof: Arc<HigherOrderUDF>) -> Expr { + let list = ScalarValue::List(ScalarValue::new_list_nullable( + &[ScalarValue::Int32(Some(1))], + &DataType::Int32, + )); + let lambda_var_with_field = Expr::LambdaVariable(LambdaVariable::new( + "x".to_string(), + Some(Arc::new(Field::new("x", DataType::Int32, true))), + )); + let lambda_var_without_field = + Expr::LambdaVariable(LambdaVariable::new("x".into(), None)); + let lambda = lambda(["x"], lambda_var_with_field + lambda_var_without_field); + Expr::HigherOrderFunction(HigherOrderFunction::new( + hof, + vec![Expr::Literal(list, None), lambda], + )) +} + +#[test] +fn hof_roundtrip_with_registry() { + let ctx = context_with_higher_order_function(); + let hof = ctx + .higher_order_function("higher_order_udf") + .expect("could not find higher order udf"); + + let expr = dummy_higher_order_function_call(hof); + + let bytes = expr.to_bytes().unwrap(); + let deserialized_expr = + Expr::from_bytes_with_ctx(&bytes, ctx.task_ctx().as_ref()).unwrap(); + + assert_eq!(expr, deserialized_expr); +} + +#[test] +#[should_panic( + expected = "LogicalExtensionCodec is not provided for higher order function higher_order_udf" +)] +fn hof_roundtrip_without_registry() { + let ctx = context_with_higher_order_function(); + let hof = ctx + .higher_order_function("higher_order_udf") + .expect("could not find higher order udf"); + + let expr = dummy_higher_order_function_call(hof); + + let bytes = expr.to_bytes().unwrap(); + Expr::from_bytes(&bytes).unwrap(); +} + +#[test] +fn test_higher_order_serialization_roundtrip() { + let ctx = SessionContext::new(); + let list = ScalarValue::List(ScalarValue::new_list_nullable( + &[ScalarValue::Int32(Some(1))], + &DataType::Int32, + )); + let lambda_var_with_field = Expr::LambdaVariable(LambdaVariable::new( + "x".to_string(), + Some(Arc::new(Field::new("x", DataType::Int32, true))), + )); + let lambda_var_without_field = + Expr::LambdaVariable(LambdaVariable::new("x".into(), None)); + let lambda = lambda(["x"], lambda_var_with_field + lambda_var_without_field); + let args = vec![Expr::Literal(list, None), lambda]; + + for function in datafusion::functions_nested::all_default_higher_order_functions() { + let expr = + Expr::HigherOrderFunction(HigherOrderFunction::new(function, args.clone())); + + let extension_codec = DefaultLogicalExtensionCodec {}; + let proto = serialize_expr(&expr, &extension_codec).unwrap(); + let deserialize = + parse_expr(&proto, ctx.task_ctx().as_ref(), &extension_codec).unwrap(); + + assert_eq!(deserialize, expr); + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
