diff --git a/datafusion/physical-expr/src/expressions/column.rs b/datafusion/physical-expr/src/expressions/column.rs index 482ab6ef1e787..ea00110ecc516 100644 --- a/datafusion/physical-expr/src/expressions/column.rs +++ b/datafusion/physical-expr/src/expressions/column.rs @@ -153,9 +153,15 @@ impl PhysicalExpr for Column { _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, ) -> Result> { use datafusion_proto_models::protobuf; + let Self { name, index } = self; Ok(Some(protobuf::PhysicalExprNode { expr_id: None, - expr_type: Some(protobuf::physical_expr_node::ExprType::Column(self.into())), + expr_type: Some(protobuf::physical_expr_node::ExprType::Column( + protobuf::PhysicalColumn { + name: name.clone(), + index: *index as u32, + }, + )), })) } } @@ -163,16 +169,18 @@ impl PhysicalExpr for Column { #[cfg(feature = "proto")] impl From<&datafusion_proto_models::protobuf::PhysicalColumn> for Column { fn from(c: &datafusion_proto_models::protobuf::PhysicalColumn) -> Self { - Column::new(&c.name, c.index as usize) + let datafusion_proto_models::protobuf::PhysicalColumn { name, index } = c; + Column::new(name, *index as usize) } } #[cfg(feature = "proto")] impl From<&Column> for datafusion_proto_models::protobuf::PhysicalColumn { fn from(c: &Column) -> Self { + let Column { name, index } = c; Self { - name: c.name.clone(), - index: c.index as u32, + name: name.clone(), + index: *index as u32, } } } @@ -196,12 +204,12 @@ impl Column { ) -> Result> { use datafusion_physical_expr_common::expect_expr_variant; use datafusion_proto_models::protobuf; - let column = expect_expr_variant!( + let protobuf::PhysicalColumn { name, index } = expect_expr_variant!( node, protobuf::physical_expr_node::ExprType::Column, "Column", ); - Ok(Arc::new(Column::from(column))) + Ok(Arc::new(Column::new(name, *index as usize))) } } diff --git a/datafusion/physical-expr/src/expressions/is_not_null.rs b/datafusion/physical-expr/src/expressions/is_not_null.rs index 3f3b7d16e543a..0444ecb724d52 100644 --- a/datafusion/physical-expr/src/expressions/is_not_null.rs +++ b/datafusion/physical-expr/src/expressions/is_not_null.rs @@ -110,11 +110,12 @@ impl PhysicalExpr for IsNotNullExpr { ) -> Result> { use datafusion_proto_models::protobuf; + let Self { arg } = self; Ok(Some(protobuf::PhysicalExprNode { expr_id: None, expr_type: Some(protobuf::physical_expr_node::ExprType::IsNotNullExpr( Box::new(protobuf::PhysicalIsNotNull { - expr: Some(Box::new(ctx.encode_child(&self.arg)?)), + expr: Some(Box::new(ctx.encode_child(arg)?)), }), )), })) @@ -136,11 +137,9 @@ impl IsNotNullExpr { protobuf::physical_expr_node::ExprType::IsNotNullExpr, "IsNotNullExpr", ); - let expr = ctx.decode_required_expression( - node.expr.as_deref(), - "IsNotNullExpr", - "expr", - )?; + let protobuf::PhysicalIsNotNull { expr } = node.as_ref(); + let expr = + ctx.decode_required_expression(expr.as_deref(), "IsNotNullExpr", "expr")?; Ok(Arc::new(IsNotNullExpr::new(expr))) } diff --git a/datafusion/physical-expr/src/expressions/is_null.rs b/datafusion/physical-expr/src/expressions/is_null.rs index da008a1cfb821..08748af2f357c 100644 --- a/datafusion/physical-expr/src/expressions/is_null.rs +++ b/datafusion/physical-expr/src/expressions/is_null.rs @@ -109,11 +109,12 @@ impl PhysicalExpr for IsNullExpr { ) -> Result> { use datafusion_proto_models::protobuf; + let Self { arg } = self; Ok(Some(protobuf::PhysicalExprNode { expr_id: None, expr_type: Some(protobuf::physical_expr_node::ExprType::IsNullExpr( Box::new(protobuf::PhysicalIsNull { - expr: Some(Box::new(ctx.encode_child(&self.arg)?)), + expr: Some(Box::new(ctx.encode_child(arg)?)), }), )), })) @@ -135,8 +136,9 @@ impl IsNullExpr { protobuf::physical_expr_node::ExprType::IsNullExpr, "IsNullExpr", ); + let protobuf::PhysicalIsNull { expr } = node.as_ref(); let expr = - ctx.decode_required_expression(node.expr.as_deref(), "IsNullExpr", "expr")?; + ctx.decode_required_expression(expr.as_deref(), "IsNullExpr", "expr")?; Ok(Arc::new(IsNullExpr::new(expr))) } diff --git a/datafusion/physical-expr/src/expressions/literal.rs b/datafusion/physical-expr/src/expressions/literal.rs index a7af824230780..3a480f901427c 100644 --- a/datafusion/physical-expr/src/expressions/literal.rs +++ b/datafusion/physical-expr/src/expressions/literal.rs @@ -143,11 +143,21 @@ impl PhysicalExpr for Literal { ) -> Result> { use datafusion_proto_models::protobuf; + let Self { value, field } = self; + // The field name, type, and nullability are reconstructed by new_with_metadata. + let expr_type = if field.metadata().is_empty() { + protobuf::physical_expr_node::ExprType::Literal(value.try_into()?) + } else { + protobuf::physical_expr_node::ExprType::LiteralWithMetadata( + protobuf::PhysicalLiteralNode { + value: Some(value.try_into()?), + metadata: field.metadata().clone(), + }, + ) + }; Ok(Some(protobuf::PhysicalExprNode { expr_id: None, - expr_type: Some(protobuf::physical_expr_node::ExprType::Literal( - (&self.value).try_into()?, - )), + expr_type: Some(expr_type), })) } } @@ -159,16 +169,38 @@ impl Literal { node: &datafusion_proto_models::protobuf::PhysicalExprNode, _ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, ) -> Result> { - use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_common::{internal_datafusion_err, internal_err}; use datafusion_proto_models::protobuf; - - let scalar_proto = expect_expr_variant!( - node, - protobuf::physical_expr_node::ExprType::Literal, - "Literal", - ); - let value = ScalarValue::try_from(scalar_proto)?; - Ok(Arc::new(Literal::new(value))) + use protobuf::physical_expr_node::ExprType; + + let protobuf::PhysicalExprNode { + expr_type, + // Expression IDs are handled by the enclosing proto converter. + expr_id: _, + } = node; + let (value, metadata) = match expr_type { + Some(ExprType::Literal(scalar)) => { + let datafusion_proto_models::datafusion_common::ScalarValue { + // The scalar payload is decoded by ScalarValue::try_from. + value: _, + } = scalar; + (ScalarValue::try_from(scalar)?, None) + } + Some(ExprType::LiteralWithMetadata(protobuf::PhysicalLiteralNode { + value, + metadata, + })) => { + let value = value.as_ref().ok_or_else(|| { + internal_datafusion_err!("Literal is missing required field 'value'") + })?; + ( + ScalarValue::try_from(value)?, + Some(FieldMetadata::from(metadata)), + ) + } + _ => return internal_err!("PhysicalExprNode is not a Literal"), + }; + Ok(Arc::new(Literal::new_with_metadata(value, metadata))) } } @@ -315,6 +347,26 @@ mod proto_tests { assert_eq!(lit.value(), &ScalarValue::Int32(Some(42))); } + #[test] + fn try_from_proto_rejects_missing_value() { + let node = datafusion_proto_models::protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::LiteralWithMetadata( + datafusion_proto_models::protobuf::PhysicalLiteralNode { + value: None, + metadata: Default::default(), + }, + )), + }; + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = Literal::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("Literal is missing required field 'value'")) + ); + } + #[test] fn try_from_proto_rejects_non_literal_node() { let node = column_node("a"); diff --git a/datafusion/physical-expr/src/expressions/negative.rs b/datafusion/physical-expr/src/expressions/negative.rs index c894c12784dc5..a932c8d5220f6 100644 --- a/datafusion/physical-expr/src/expressions/negative.rs +++ b/datafusion/physical-expr/src/expressions/negative.rs @@ -184,11 +184,12 @@ impl PhysicalExpr for NegativeExpr { ) -> Result> { use datafusion_proto_models::protobuf; + let Self { arg } = self; Ok(Some(protobuf::PhysicalExprNode { expr_id: None, expr_type: Some(protobuf::physical_expr_node::ExprType::Negative(Box::new( protobuf::PhysicalNegativeNode { - expr: Some(Box::new(ctx.encode_child(&self.arg)?)), + expr: Some(Box::new(ctx.encode_child(arg)?)), }, ))), })) @@ -210,8 +211,9 @@ impl NegativeExpr { protobuf::physical_expr_node::ExprType::Negative, "Negative", ); + let protobuf::PhysicalNegativeNode { expr } = n.as_ref(); let expr = - ctx.decode_required_expression(n.expr.as_deref(), "NegativeExpr", "expr")?; + ctx.decode_required_expression(expr.as_deref(), "NegativeExpr", "expr")?; Ok(Arc::new(NegativeExpr::new(expr))) } diff --git a/datafusion/physical-expr/src/expressions/not.rs b/datafusion/physical-expr/src/expressions/not.rs index f856dd568a8da..7699c8944ea23 100644 --- a/datafusion/physical-expr/src/expressions/not.rs +++ b/datafusion/physical-expr/src/expressions/not.rs @@ -189,11 +189,12 @@ impl PhysicalExpr for NotExpr { ) -> Result> { use datafusion_proto_models::protobuf; + let Self { arg } = self; Ok(Some(protobuf::PhysicalExprNode { expr_id: None, expr_type: Some(protobuf::physical_expr_node::ExprType::NotExpr(Box::new( protobuf::PhysicalNot { - expr: Some(Box::new(ctx.encode_child(&self.arg)?)), + expr: Some(Box::new(ctx.encode_child(arg)?)), }, ))), })) @@ -215,8 +216,8 @@ impl NotExpr { protobuf::physical_expr_node::ExprType::NotExpr, "NotExpr", ); - let expr = - ctx.decode_required_expression(not_expr.expr.as_deref(), "NotExpr", "expr")?; + let protobuf::PhysicalNot { expr } = not_expr.as_ref(); + let expr = ctx.decode_required_expression(expr.as_deref(), "NotExpr", "expr")?; Ok(Arc::new(NotExpr::new(expr))) } diff --git a/datafusion/physical-expr/src/expressions/unknown_column.rs b/datafusion/physical-expr/src/expressions/unknown_column.rs index ed85f20dd274b..c3b6dd9dafade 100644 --- a/datafusion/physical-expr/src/expressions/unknown_column.rs +++ b/datafusion/physical-expr/src/expressions/unknown_column.rs @@ -93,12 +93,11 @@ impl PhysicalExpr for UnKnownColumn { ) -> Result> { use datafusion_proto_models::protobuf; + let Self { name } = self; Ok(Some(protobuf::PhysicalExprNode { expr_id: None, expr_type: Some(protobuf::physical_expr_node::ExprType::UnknownColumn( - protobuf::UnknownColumn { - name: self.name.clone(), - }, + protobuf::UnknownColumn { name: name.clone() }, )), })) } @@ -114,12 +113,12 @@ impl UnKnownColumn { use datafusion_physical_expr_common::expect_expr_variant; use datafusion_proto_models::protobuf; - let unknown_col = expect_expr_variant!( + let protobuf::UnknownColumn { name } = expect_expr_variant!( node, protobuf::physical_expr_node::ExprType::UnknownColumn, "UnKnownColumn", ); - Ok(Arc::new(UnKnownColumn::new(&unknown_col.name))) + Ok(Arc::new(UnKnownColumn::new(name))) } } diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index fac5ff27191cd..ccf0b4a630517 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1074,9 +1074,15 @@ message PhysicalExprNode { PhysicalLambdaVariableExprNode lambda_variable = 26; PhysicalRangeExprNode range_expr = 27; PhysicalSqlSimilarToPatternNode sql_similar_to_pattern = 28; + PhysicalLiteralNode literal_with_metadata = 29; } } +message PhysicalLiteralNode { + datafusion_common.ScalarValue value = 1; + map metadata = 2; +} + message PhysicalDynamicFilterNode { repeated PhysicalExprNode children = 1; repeated PhysicalExprNode remapped_children = 2; diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index e06f5b7011504..548755ff0c0cd 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -19253,6 +19253,9 @@ impl serde::Serialize for PhysicalExprNode { physical_expr_node::ExprType::SqlSimilarToPattern(v) => { struct_ser.serialize_field("sqlSimilarToPattern", v)?; } + physical_expr_node::ExprType::LiteralWithMetadata(v) => { + struct_ser.serialize_field("literalWithMetadata", v)?; + } } } struct_ser.end() @@ -19312,6 +19315,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { "rangeExpr", "sql_similar_to_pattern", "sqlSimilarToPattern", + "literal_with_metadata", + "literalWithMetadata", ]; #[allow(clippy::enum_variant_names)] @@ -19343,6 +19348,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { LambdaVariable, RangeExpr, SqlSimilarToPattern, + LiteralWithMetadata, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -19391,6 +19397,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { "lambdaVariable" | "lambda_variable" => Ok(GeneratedField::LambdaVariable), "rangeExpr" | "range_expr" => Ok(GeneratedField::RangeExpr), "sqlSimilarToPattern" | "sql_similar_to_pattern" => Ok(GeneratedField::SqlSimilarToPattern), + "literalWithMetadata" | "literal_with_metadata" => Ok(GeneratedField::LiteralWithMetadata), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -19602,6 +19609,13 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { return Err(serde::de::Error::duplicate_field("sqlSimilarToPattern")); } expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::SqlSimilarToPattern) +; + } + GeneratedField::LiteralWithMetadata => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("literalWithMetadata")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::LiteralWithMetadata) ; } } @@ -20881,6 +20895,116 @@ impl<'de> serde::Deserialize<'de> for PhysicalLikeExprNode { deserializer.deserialize_struct("datafusion.PhysicalLikeExprNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for PhysicalLiteralNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.value.is_some() { + len += 1; + } + if !self.metadata.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalLiteralNode", len)?; + if let Some(v) = self.value.as_ref() { + struct_ser.serialize_field("value", v)?; + } + if !self.metadata.is_empty() { + struct_ser.serialize_field("metadata", &self.metadata)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalLiteralNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "value", + "metadata", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Value, + Metadata, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + 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(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "value" => Ok(GeneratedField::Value), + "metadata" => Ok(GeneratedField::Metadata), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalLiteralNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalLiteralNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut value__ = None; + let mut metadata__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Value => { + if value__.is_some() { + return Err(serde::de::Error::duplicate_field("value")); + } + value__ = map_.next_value()?; + } + GeneratedField::Metadata => { + if metadata__.is_some() { + return Err(serde::de::Error::duplicate_field("metadata")); + } + metadata__ = Some( + map_.next_value::>()? + ); + } + } + } + Ok(PhysicalLiteralNode { + value: value__, + metadata: metadata__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalLiteralNode", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for PhysicalNegativeNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 149839dacc967..f08298cfa897e 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1588,7 +1588,7 @@ pub struct PhysicalExprNode { pub expr_id: ::core::option::Option, #[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, 24, 25, 26, 27, 28" + 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, 27, 28, 29" )] pub expr_type: ::core::option::Option, } @@ -1657,9 +1657,21 @@ pub mod physical_expr_node { SqlSimilarToPattern( ::prost::alloc::boxed::Box, ), + #[prost(message, tag = "29")] + LiteralWithMetadata(super::PhysicalLiteralNode), } } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalLiteralNode { + #[prost(message, optional, tag = "1")] + pub value: ::core::option::Option, + #[prost(map = "string, string", tag = "2")] + pub metadata: ::std::collections::HashMap< + ::prost::alloc::string::String, + ::prost::alloc::string::String, + >, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalDynamicFilterNode { #[prost(message, repeated, tag = "1")] pub children: ::prost::alloc::vec::Vec, diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index bc443149df413..8c3416ec98d50 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -269,7 +269,9 @@ pub fn parse_physical_expr_with_converter( // to the right constructor. ExprType::Column(_) => Column::try_from_proto(proto, &decode_ctx)?, ExprType::UnknownColumn(_) => UnKnownColumn::try_from_proto(proto, &decode_ctx)?, - ExprType::Literal(_) => Literal::try_from_proto(proto, &decode_ctx)?, + ExprType::Literal(_) | ExprType::LiteralWithMetadata(_) => { + Literal::try_from_proto(proto, &decode_ctx)? + } ExprType::BinaryExpr(_) => BinaryExpr::try_from_proto(proto, &decode_ctx)?, ExprType::AggregateExpr(_) => { return not_impl_err!( diff --git a/datafusion/proto/tests/cases/plans/exprs.rs b/datafusion/proto/tests/cases/plans/exprs.rs index f2b14b043959f..9518e23732a1f 100644 --- a/datafusion/proto/tests/cases/plans/exprs.rs +++ b/datafusion/proto/tests/cases/plans/exprs.rs @@ -46,6 +46,54 @@ use datafusion_proto::protobuf::PhysicalPlanNode; use std::sync::Arc; use std::vec; +#[test] +fn roundtrip_literal_metadata() -> Result<()> { + use datafusion_common::metadata::FieldMetadata; + use datafusion_proto::bytes::{physical_plan_from_bytes, physical_plan_to_bytes}; + + let schema = Arc::new(Schema::empty()); + let metadata = FieldMetadata::from(std::collections::HashMap::from([ + ("ARROW:extension:name".to_string(), "example.id".to_string()), + ("ARROW:extension:metadata".to_string(), "{}".to_string()), + ("description".to_string(), "identifier".to_string()), + ])); + let ctx = SessionContext::new(); + for value in [ScalarValue::Int32(Some(42)), ScalarValue::Int32(None)] { + let literal: Arc = Arc::new(Literal::new_with_metadata( + value.clone(), + Some(metadata.clone()), + )); + let plan: Arc = Arc::new(ProjectionExec::try_new( + vec![ProjectionExpr::new(Arc::clone(&literal), "result")], + Arc::new(EmptyExec::new(Arc::clone(&schema))), + )?); + let bytes = physical_plan_to_bytes(Arc::clone(&plan))?; + let decoded = physical_plan_from_bytes(&bytes, ctx.task_ctx().as_ref())?; + assert_eq!(decoded.schema(), plan.schema()); + let projection = decoded.downcast_ref::().unwrap(); + let decoded_literal = &projection.expr()[0].expr; + assert_eq!( + decoded_literal.return_field(&schema)?, + literal.return_field(&schema)? + ); + assert_eq!( + decoded_literal.downcast_ref::().unwrap().value(), + &value + ); + + #[cfg(feature = "json")] + { + use datafusion_proto::bytes::{ + physical_plan_from_json, physical_plan_to_json, + }; + let json = physical_plan_to_json(Arc::clone(&plan))?; + let decoded = physical_plan_from_json(&json, ctx.task_ctx().as_ref())?; + assert_eq!(decoded.schema(), plan.schema()); + } + } + Ok(()) +} + #[test] fn roundtrip_date_time_interval() -> Result<()> { let schema = Schema::new(vec![