From 76fce1378712975317437990eb56b373f67884f0 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Wed, 16 Sep 2026 23:53:38 +0800 Subject: [PATCH 1/3] fix: preserve literal metadata and destructure leaf and unary proto hooks --- .../physical-expr/src/expressions/column.rs | 20 ++- .../src/expressions/is_not_null.rs | 11 +- .../physical-expr/src/expressions/is_null.rs | 6 +- .../physical-expr/src/expressions/literal.rs | 76 +++++++++-- .../physical-expr/src/expressions/negative.rs | 6 +- .../physical-expr/src/expressions/not.rs | 7 +- .../src/expressions/unknown_column.rs | 9 +- .../proto-models/proto/datafusion.proto | 6 + .../proto-models/src/generated/pbjson.rs | 124 ++++++++++++++++++ .../proto-models/src/generated/prost.rs | 14 +- .../proto/src/physical_plan/from_proto.rs | 4 +- datafusion/proto/tests/cases/plans/exprs.rs | 48 +++++++ 12 files changed, 293 insertions(+), 38 deletions(-) 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![ From 4095c51f7df922be58f9f73f03884ad68dd568cd Mon Sep 17 00:00:00 2001 From: peterxcli Date: Thu, 17 Sep 2026 00:22:08 +0800 Subject: [PATCH 2/3] refactor: move proto hook imports to module scope --- .../physical-expr/src/expressions/column.rs | 50 +++++++------ .../src/expressions/is_not_null.rs | 56 +++++++-------- .../physical-expr/src/expressions/is_null.rs | 54 +++++++------- .../physical-expr/src/expressions/literal.rs | 71 +++++++++---------- .../physical-expr/src/expressions/negative.rs | 54 +++++++------- .../physical-expr/src/expressions/not.rs | 54 +++++++------- .../src/expressions/unknown_column.rs | 57 ++++++++------- 7 files changed, 185 insertions(+), 211 deletions(-) diff --git a/datafusion/physical-expr/src/expressions/column.rs b/datafusion/physical-expr/src/expressions/column.rs index ea00110ecc516..1f504acb8a5fd 100644 --- a/datafusion/physical-expr/src/expressions/column.rs +++ b/datafusion/physical-expr/src/expressions/column.rs @@ -31,6 +31,18 @@ use datafusion_common::{Result, internal_err, plan_err}; use datafusion_expr::ColumnarValue; use datafusion_expr_common::placement::ExpressionPlacement; +#[cfg(feature = "proto")] +use datafusion_physical_expr_common::{ + expect_expr_variant, + physical_expr::{ + proto_decode::PhysicalExprDecodeCtx, proto_encode::PhysicalExprEncodeCtx, + }, +}; +#[cfg(feature = "proto")] +use datafusion_proto_models::protobuf::{ + PhysicalColumn, PhysicalExprNode, physical_expr_node::ExprType, +}; + /// Represents the column at a given index in a RecordBatch /// /// This is a physical expression that represents a column at a given index in an @@ -150,32 +162,29 @@ impl PhysicalExpr for Column { #[cfg(feature = "proto")] fn try_to_proto( &self, - _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, - ) -> Result> { - use datafusion_proto_models::protobuf; + _ctx: &PhysicalExprEncodeCtx<'_>, + ) -> Result> { let Self { name, index } = self; - Ok(Some(protobuf::PhysicalExprNode { + Ok(Some(PhysicalExprNode { expr_id: None, - expr_type: Some(protobuf::physical_expr_node::ExprType::Column( - protobuf::PhysicalColumn { - name: name.clone(), - index: *index as u32, - }, - )), + expr_type: Some(ExprType::Column(PhysicalColumn { + name: name.clone(), + index: *index as u32, + })), })) } } #[cfg(feature = "proto")] -impl From<&datafusion_proto_models::protobuf::PhysicalColumn> for Column { - fn from(c: &datafusion_proto_models::protobuf::PhysicalColumn) -> Self { - let datafusion_proto_models::protobuf::PhysicalColumn { name, index } = c; +impl From<&PhysicalColumn> for Column { + fn from(c: &PhysicalColumn) -> Self { + let PhysicalColumn { name, index } = c; Column::new(name, *index as usize) } } #[cfg(feature = "proto")] -impl From<&Column> for datafusion_proto_models::protobuf::PhysicalColumn { +impl From<&Column> for PhysicalColumn { fn from(c: &Column) -> Self { let Column { name, index } = c; Self { @@ -199,16 +208,11 @@ impl Column { /// [`PhysicalExpr::try_to_proto`]: datafusion_physical_expr_common::physical_expr::PhysicalExpr::try_to_proto /// [`PhysicalExprDecodeCtx::decode`]: datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx::decode pub fn try_from_proto( - node: &datafusion_proto_models::protobuf::PhysicalExprNode, - _ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + node: &PhysicalExprNode, + _ctx: &PhysicalExprDecodeCtx<'_>, ) -> Result> { - use datafusion_physical_expr_common::expect_expr_variant; - use datafusion_proto_models::protobuf; - let protobuf::PhysicalColumn { name, index } = expect_expr_variant!( - node, - protobuf::physical_expr_node::ExprType::Column, - "Column", - ); + let PhysicalColumn { name, index } = + expect_expr_variant!(node, ExprType::Column, "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 0444ecb724d52..190d85e045f76 100644 --- a/datafusion/physical-expr/src/expressions/is_not_null.rs +++ b/datafusion/physical-expr/src/expressions/is_not_null.rs @@ -27,6 +27,18 @@ use datafusion_expr::ColumnarValue; use std::hash::Hash; use std::sync::Arc; +#[cfg(feature = "proto")] +use datafusion_physical_expr_common::{ + expect_expr_variant, + physical_expr::{ + proto_decode::PhysicalExprDecodeCtx, proto_encode::PhysicalExprEncodeCtx, + }, +}; +#[cfg(feature = "proto")] +use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalIsNotNull, physical_expr_node::ExprType, +}; + /// IS NOT NULL expression #[derive(Debug, Eq)] pub struct IsNotNullExpr { @@ -106,18 +118,14 @@ impl PhysicalExpr for IsNotNullExpr { #[cfg(feature = "proto")] fn try_to_proto( &self, - ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, - ) -> Result> { - use datafusion_proto_models::protobuf; - + ctx: &PhysicalExprEncodeCtx<'_>, + ) -> Result> { let Self { arg } = self; - Ok(Some(protobuf::PhysicalExprNode { + Ok(Some(PhysicalExprNode { expr_id: None, - expr_type: Some(protobuf::physical_expr_node::ExprType::IsNotNullExpr( - Box::new(protobuf::PhysicalIsNotNull { - expr: Some(Box::new(ctx.encode_child(arg)?)), - }), - )), + expr_type: Some(ExprType::IsNotNullExpr(Box::new(PhysicalIsNotNull { + expr: Some(Box::new(ctx.encode_child(arg)?)), + }))), })) } } @@ -126,18 +134,11 @@ impl PhysicalExpr for IsNotNullExpr { impl IsNotNullExpr { /// Reconstruct an [`IsNotNullExpr`] from its protobuf representation. pub fn try_from_proto( - node: &datafusion_proto_models::protobuf::PhysicalExprNode, - ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + node: &PhysicalExprNode, + ctx: &PhysicalExprDecodeCtx<'_>, ) -> Result> { - use datafusion_physical_expr_common::expect_expr_variant; - use datafusion_proto_models::protobuf; - - let node = expect_expr_variant!( - node, - protobuf::physical_expr_node::ExprType::IsNotNullExpr, - "IsNotNullExpr", - ); - let protobuf::PhysicalIsNotNull { expr } = node.as_ref(); + let node = expect_expr_variant!(node, ExprType::IsNotNullExpr, "IsNotNullExpr"); + let PhysicalIsNotNull { expr } = node.as_ref(); let expr = ctx.decode_required_expression(expr.as_deref(), "IsNotNullExpr", "expr")?; @@ -263,18 +264,13 @@ mod proto_tests { }; use arrow::datatypes::Field; use datafusion_common::DataFusionError; - use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; - use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; - use datafusion_proto_models::protobuf::{ - PhysicalExprNode, PhysicalIsNotNull, physical_expr_node, - }; fn is_not_null_node(expr: Option>) -> PhysicalExprNode { PhysicalExprNode { expr_id: None, - expr_type: Some(physical_expr_node::ExprType::IsNotNullExpr(Box::new( - PhysicalIsNotNull { expr }, - ))), + expr_type: Some(ExprType::IsNotNullExpr(Box::new(PhysicalIsNotNull { + expr, + }))), } } @@ -296,7 +292,7 @@ mod proto_tests { assert!(node.expr_id.is_none()); let is_not_null_node = match node.expr_type { - Some(physical_expr_node::ExprType::IsNotNullExpr(boxed)) => *boxed, + Some(ExprType::IsNotNullExpr(boxed)) => *boxed, other => panic!("expected an IsNotNullExpr node, got {other:?}"), }; assert!(is_not_null_node.expr.is_some()); diff --git a/datafusion/physical-expr/src/expressions/is_null.rs b/datafusion/physical-expr/src/expressions/is_null.rs index 08748af2f357c..bfa79d72a2282 100644 --- a/datafusion/physical-expr/src/expressions/is_null.rs +++ b/datafusion/physical-expr/src/expressions/is_null.rs @@ -27,6 +27,18 @@ use datafusion_expr::ColumnarValue; use std::hash::Hash; use std::sync::Arc; +#[cfg(feature = "proto")] +use datafusion_physical_expr_common::{ + expect_expr_variant, + physical_expr::{ + proto_decode::PhysicalExprDecodeCtx, proto_encode::PhysicalExprEncodeCtx, + }, +}; +#[cfg(feature = "proto")] +use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalIsNull, physical_expr_node::ExprType, +}; + /// IS NULL expression #[derive(Debug, Eq)] pub struct IsNullExpr { @@ -105,18 +117,14 @@ impl PhysicalExpr for IsNullExpr { #[cfg(feature = "proto")] fn try_to_proto( &self, - ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, - ) -> Result> { - use datafusion_proto_models::protobuf; - + ctx: &PhysicalExprEncodeCtx<'_>, + ) -> Result> { let Self { arg } = self; - Ok(Some(protobuf::PhysicalExprNode { + Ok(Some(PhysicalExprNode { expr_id: None, - expr_type: Some(protobuf::physical_expr_node::ExprType::IsNullExpr( - Box::new(protobuf::PhysicalIsNull { - expr: Some(Box::new(ctx.encode_child(arg)?)), - }), - )), + expr_type: Some(ExprType::IsNullExpr(Box::new(PhysicalIsNull { + expr: Some(Box::new(ctx.encode_child(arg)?)), + }))), })) } } @@ -125,18 +133,11 @@ impl PhysicalExpr for IsNullExpr { impl IsNullExpr { /// Reconstruct an [`IsNullExpr`] from its protobuf representation. pub fn try_from_proto( - node: &datafusion_proto_models::protobuf::PhysicalExprNode, - ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + node: &PhysicalExprNode, + ctx: &PhysicalExprDecodeCtx<'_>, ) -> Result> { - use datafusion_physical_expr_common::expect_expr_variant; - use datafusion_proto_models::protobuf; - - let node = expect_expr_variant!( - node, - protobuf::physical_expr_node::ExprType::IsNullExpr, - "IsNullExpr", - ); - let protobuf::PhysicalIsNull { expr } = node.as_ref(); + let node = expect_expr_variant!(node, ExprType::IsNullExpr, "IsNullExpr"); + let PhysicalIsNull { expr } = node.as_ref(); let expr = ctx.decode_required_expression(expr.as_deref(), "IsNullExpr", "expr")?; @@ -274,18 +275,11 @@ mod proto_tests { }; use arrow::datatypes::Field; use datafusion_common::DataFusionError; - use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; - use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; - use datafusion_proto_models::protobuf::{ - PhysicalExprNode, PhysicalIsNull, physical_expr_node, - }; fn is_null_node(expr: Option>) -> PhysicalExprNode { PhysicalExprNode { expr_id: None, - expr_type: Some(physical_expr_node::ExprType::IsNullExpr(Box::new( - PhysicalIsNull { expr }, - ))), + expr_type: Some(ExprType::IsNullExpr(Box::new(PhysicalIsNull { expr }))), } } @@ -307,7 +301,7 @@ mod proto_tests { assert!(node.expr_id.is_none()); let is_null_node = match node.expr_type { - Some(physical_expr_node::ExprType::IsNullExpr(boxed)) => *boxed, + Some(ExprType::IsNullExpr(boxed)) => *boxed, other => panic!("expected an IsNullExpr node, got {other:?}"), }; assert!(is_null_node.expr.is_some()); diff --git a/datafusion/physical-expr/src/expressions/literal.rs b/datafusion/physical-expr/src/expressions/literal.rs index 3a480f901427c..e68decbc3e102 100644 --- a/datafusion/physical-expr/src/expressions/literal.rs +++ b/datafusion/physical-expr/src/expressions/literal.rs @@ -35,6 +35,18 @@ use datafusion_expr_common::interval_arithmetic::Interval; use datafusion_expr_common::placement::ExpressionPlacement; use datafusion_expr_common::sort_properties::{ExprProperties, SortProperties}; +#[cfg(feature = "proto")] +use datafusion_common::{internal_datafusion_err, internal_err}; +#[cfg(feature = "proto")] +use datafusion_physical_expr_common::physical_expr::{ + proto_decode::PhysicalExprDecodeCtx, proto_encode::PhysicalExprEncodeCtx, +}; +#[cfg(feature = "proto")] +use datafusion_proto_models::{ + datafusion_common::ScalarValue as ProtoScalarValue, + protobuf::{PhysicalExprNode, PhysicalLiteralNode, physical_expr_node::ExprType}, +}; + /// Represents a literal value #[derive(Debug, PartialEq, Eq, Clone)] pub struct Literal { @@ -139,23 +151,19 @@ impl PhysicalExpr for Literal { #[cfg(feature = "proto")] fn try_to_proto( &self, - _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, - ) -> Result> { - use datafusion_proto_models::protobuf; - + _ctx: &PhysicalExprEncodeCtx<'_>, + ) -> Result> { 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()?) + ExprType::Literal(value.try_into()?) } else { - protobuf::physical_expr_node::ExprType::LiteralWithMetadata( - protobuf::PhysicalLiteralNode { - value: Some(value.try_into()?), - metadata: field.metadata().clone(), - }, - ) + ExprType::LiteralWithMetadata(PhysicalLiteralNode { + value: Some(value.try_into()?), + metadata: field.metadata().clone(), + }) }; - Ok(Some(protobuf::PhysicalExprNode { + Ok(Some(PhysicalExprNode { expr_id: None, expr_type: Some(expr_type), })) @@ -166,27 +174,23 @@ impl PhysicalExpr for Literal { impl Literal { /// Reconstruct a [`Literal`] from its protobuf representation. pub fn try_from_proto( - node: &datafusion_proto_models::protobuf::PhysicalExprNode, - _ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + node: &PhysicalExprNode, + _ctx: &PhysicalExprDecodeCtx<'_>, ) -> Result> { - use datafusion_common::{internal_datafusion_err, internal_err}; - use datafusion_proto_models::protobuf; - use protobuf::physical_expr_node::ExprType; - - let protobuf::PhysicalExprNode { + let 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 { + let ProtoScalarValue { // The scalar payload is decoded by ScalarValue::try_from. value: _, } = scalar; (ScalarValue::try_from(scalar)?, None) } - Some(ExprType::LiteralWithMetadata(protobuf::PhysicalLiteralNode { + Some(ExprType::LiteralWithMetadata(PhysicalLiteralNode { value, metadata, })) => { @@ -266,9 +270,6 @@ mod proto_tests { use super::*; use crate::proto_test_util::{StubEncoder, UnreachableDecoder, column_node}; use datafusion_common::DataFusionError; - use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; - use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; - use datafusion_proto_models::protobuf::physical_expr_node; fn i32_literal() -> Literal { Literal::new(ScalarValue::Int32(Some(42))) @@ -290,10 +291,7 @@ mod proto_tests { // Literal nodes never set expr_id. assert!(node.expr_id.is_none()); // Variant must be Literal, not any other expr type. - assert!(matches!( - node.expr_type, - Some(physical_expr_node::ExprType::Literal(_)) - )); + assert!(matches!(node.expr_type, Some(ExprType::Literal(_)))); } #[test] @@ -307,10 +305,7 @@ mod proto_tests { .unwrap() .expect("null Literal should encode to Some(node)"); - assert!(matches!( - node.expr_type, - Some(physical_expr_node::ExprType::Literal(_)) - )); + assert!(matches!(node.expr_type, Some(ExprType::Literal(_)))); // Decode and verify the null payload round-trips correctly. let schema = Schema::empty(); @@ -349,14 +344,12 @@ mod proto_tests { #[test] fn try_from_proto_rejects_missing_value() { - let node = datafusion_proto_models::protobuf::PhysicalExprNode { + let node = PhysicalExprNode { expr_id: None, - expr_type: Some(physical_expr_node::ExprType::LiteralWithMetadata( - datafusion_proto_models::protobuf::PhysicalLiteralNode { - value: None, - metadata: Default::default(), - }, - )), + expr_type: Some(ExprType::LiteralWithMetadata(PhysicalLiteralNode { + value: None, + metadata: Default::default(), + })), }; let schema = Schema::empty(); let decoder = UnreachableDecoder; diff --git a/datafusion/physical-expr/src/expressions/negative.rs b/datafusion/physical-expr/src/expressions/negative.rs index a932c8d5220f6..fcdedcfedc272 100644 --- a/datafusion/physical-expr/src/expressions/negative.rs +++ b/datafusion/physical-expr/src/expressions/negative.rs @@ -40,6 +40,18 @@ use datafusion_expr::{ type_coercion::{is_interval, is_signed_numeric, is_timestamp}, }; +#[cfg(feature = "proto")] +use datafusion_physical_expr_common::{ + expect_expr_variant, + physical_expr::{ + proto_decode::PhysicalExprDecodeCtx, proto_encode::PhysicalExprEncodeCtx, + }, +}; +#[cfg(feature = "proto")] +use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalNegativeNode, physical_expr_node::ExprType, +}; + /// Negative expression #[derive(Debug, Eq)] pub struct NegativeExpr { @@ -180,18 +192,14 @@ impl PhysicalExpr for NegativeExpr { #[cfg(feature = "proto")] fn try_to_proto( &self, - ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, - ) -> Result> { - use datafusion_proto_models::protobuf; - + ctx: &PhysicalExprEncodeCtx<'_>, + ) -> Result> { let Self { arg } = self; - Ok(Some(protobuf::PhysicalExprNode { + Ok(Some(PhysicalExprNode { expr_id: None, - expr_type: Some(protobuf::physical_expr_node::ExprType::Negative(Box::new( - protobuf::PhysicalNegativeNode { - expr: Some(Box::new(ctx.encode_child(arg)?)), - }, - ))), + expr_type: Some(ExprType::Negative(Box::new(PhysicalNegativeNode { + expr: Some(Box::new(ctx.encode_child(arg)?)), + }))), })) } } @@ -200,18 +208,11 @@ impl PhysicalExpr for NegativeExpr { impl NegativeExpr { /// Reconstruct a [`NegativeExpr`] from its protobuf representation. pub fn try_from_proto( - node: &datafusion_proto_models::protobuf::PhysicalExprNode, - ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + node: &PhysicalExprNode, + ctx: &PhysicalExprDecodeCtx<'_>, ) -> Result> { - use datafusion_physical_expr_common::expect_expr_variant; - use datafusion_proto_models::protobuf; - - let n = expect_expr_variant!( - node, - protobuf::physical_expr_node::ExprType::Negative, - "Negative", - ); - let protobuf::PhysicalNegativeNode { expr } = n.as_ref(); + let n = expect_expr_variant!(node, ExprType::Negative, "Negative"); + let PhysicalNegativeNode { expr } = n.as_ref(); let expr = ctx.decode_required_expression(expr.as_deref(), "NegativeExpr", "expr")?; @@ -455,19 +456,12 @@ mod proto_tests { }; use arrow::datatypes::Field; use datafusion_common::DataFusionError; - use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; - use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; - use datafusion_proto_models::protobuf::{ - PhysicalExprNode, PhysicalNegativeNode, physical_expr_node, - }; /// Build a `NegativeExpr` proto node with the given children. fn negative_node(expr: Option>) -> PhysicalExprNode { PhysicalExprNode { expr_id: None, - expr_type: Some(physical_expr_node::ExprType::Negative(Box::new( - PhysicalNegativeNode { expr }, - ))), + expr_type: Some(ExprType::Negative(Box::new(PhysicalNegativeNode { expr }))), } } @@ -490,7 +484,7 @@ mod proto_tests { assert!(node.expr_id.is_none()); let negative_node = match node.expr_type { - Some(physical_expr_node::ExprType::Negative(boxed)) => *boxed, + Some(ExprType::Negative(boxed)) => *boxed, other => panic!("expected a NegativeExpr node, got {other:?}"), }; assert!(negative_node.expr.is_some()); diff --git a/datafusion/physical-expr/src/expressions/not.rs b/datafusion/physical-expr/src/expressions/not.rs index 7699c8944ea23..4c2ac97771054 100644 --- a/datafusion/physical-expr/src/expressions/not.rs +++ b/datafusion/physical-expr/src/expressions/not.rs @@ -31,6 +31,18 @@ use datafusion_expr::interval_arithmetic::Interval; #[expect(deprecated)] use datafusion_expr::statistics::Distribution::{self, Bernoulli}; +#[cfg(feature = "proto")] +use datafusion_physical_expr_common::{ + expect_expr_variant, + physical_expr::{ + proto_decode::PhysicalExprDecodeCtx, proto_encode::PhysicalExprEncodeCtx, + }, +}; +#[cfg(feature = "proto")] +use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalNot, physical_expr_node::ExprType, +}; + /// Not expression #[derive(Debug, Eq)] pub struct NotExpr { @@ -185,18 +197,14 @@ impl PhysicalExpr for NotExpr { #[cfg(feature = "proto")] fn try_to_proto( &self, - ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, - ) -> Result> { - use datafusion_proto_models::protobuf; - + ctx: &PhysicalExprEncodeCtx<'_>, + ) -> Result> { let Self { arg } = self; - Ok(Some(protobuf::PhysicalExprNode { + Ok(Some(PhysicalExprNode { expr_id: None, - expr_type: Some(protobuf::physical_expr_node::ExprType::NotExpr(Box::new( - protobuf::PhysicalNot { - expr: Some(Box::new(ctx.encode_child(arg)?)), - }, - ))), + expr_type: Some(ExprType::NotExpr(Box::new(PhysicalNot { + expr: Some(Box::new(ctx.encode_child(arg)?)), + }))), })) } } @@ -205,18 +213,11 @@ impl PhysicalExpr for NotExpr { impl NotExpr { /// Reconstruct a [`NotExpr`] from its protobuf representation. pub fn try_from_proto( - node: &datafusion_proto_models::protobuf::PhysicalExprNode, - ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + node: &PhysicalExprNode, + ctx: &PhysicalExprDecodeCtx<'_>, ) -> Result> { - use datafusion_physical_expr_common::expect_expr_variant; - use datafusion_proto_models::protobuf; - - let not_expr = expect_expr_variant!( - node, - protobuf::physical_expr_node::ExprType::NotExpr, - "NotExpr", - ); - let protobuf::PhysicalNot { expr } = not_expr.as_ref(); + let not_expr = expect_expr_variant!(node, ExprType::NotExpr, "NotExpr"); + let PhysicalNot { expr } = not_expr.as_ref(); let expr = ctx.decode_required_expression(expr.as_deref(), "NotExpr", "expr")?; Ok(Arc::new(NotExpr::new(expr))) @@ -408,19 +409,12 @@ mod proto_tests { }; use arrow::datatypes::Field; use datafusion_common::DataFusionError; - use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; - use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; - use datafusion_proto_models::protobuf::{ - PhysicalExprNode, PhysicalNot, physical_expr_node, - }; /// Build a `NotExpr` proto node with the given child. fn not_node(expr: Option>) -> PhysicalExprNode { PhysicalExprNode { expr_id: None, - expr_type: Some(physical_expr_node::ExprType::NotExpr(Box::new( - PhysicalNot { expr }, - ))), + expr_type: Some(ExprType::NotExpr(Box::new(PhysicalNot { expr }))), } } @@ -443,7 +437,7 @@ mod proto_tests { assert!(node.expr_id.is_none()); let not_node = match node.expr_type { - Some(physical_expr_node::ExprType::NotExpr(boxed)) => *boxed, + Some(ExprType::NotExpr(boxed)) => *boxed, other => panic!("expected a NotExpr node, got {other:?}"), }; assert!(not_node.expr.is_some()); diff --git a/datafusion/physical-expr/src/expressions/unknown_column.rs b/datafusion/physical-expr/src/expressions/unknown_column.rs index c3b6dd9dafade..c9df5eefcdf40 100644 --- a/datafusion/physical-expr/src/expressions/unknown_column.rs +++ b/datafusion/physical-expr/src/expressions/unknown_column.rs @@ -30,6 +30,18 @@ use datafusion_common::{Result, internal_err}; use datafusion_expr::ColumnarValue; +#[cfg(feature = "proto")] +use datafusion_physical_expr_common::{ + expect_expr_variant, + physical_expr::{ + proto_decode::PhysicalExprDecodeCtx, proto_encode::PhysicalExprEncodeCtx, + }, +}; +#[cfg(feature = "proto")] +use datafusion_proto_models::protobuf::{ + PhysicalExprNode, UnknownColumn, physical_expr_node::ExprType, +}; + #[derive(Debug, Clone, Eq)] pub struct UnKnownColumn { name: String, @@ -89,16 +101,14 @@ impl PhysicalExpr for UnKnownColumn { #[cfg(feature = "proto")] fn try_to_proto( &self, - _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, - ) -> Result> { - use datafusion_proto_models::protobuf; - + _ctx: &PhysicalExprEncodeCtx<'_>, + ) -> Result> { let Self { name } = self; - Ok(Some(protobuf::PhysicalExprNode { + Ok(Some(PhysicalExprNode { expr_id: None, - expr_type: Some(protobuf::physical_expr_node::ExprType::UnknownColumn( - protobuf::UnknownColumn { name: name.clone() }, - )), + expr_type: Some(ExprType::UnknownColumn(UnknownColumn { + name: name.clone(), + })), })) } } @@ -107,17 +117,11 @@ impl PhysicalExpr for UnKnownColumn { impl UnKnownColumn { /// Reconstruct an [`UnKnownColumn`] from its protobuf representation. pub fn try_from_proto( - node: &datafusion_proto_models::protobuf::PhysicalExprNode, - _ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + node: &PhysicalExprNode, + _ctx: &PhysicalExprDecodeCtx<'_>, ) -> Result> { - use datafusion_physical_expr_common::expect_expr_variant; - use datafusion_proto_models::protobuf; - - let protobuf::UnknownColumn { name } = expect_expr_variant!( - node, - protobuf::physical_expr_node::ExprType::UnknownColumn, - "UnKnownColumn", - ); + let UnknownColumn { name } = + expect_expr_variant!(node, ExprType::UnknownColumn, "UnKnownColumn"); Ok(Arc::new(UnKnownColumn::new(name))) } } @@ -143,9 +147,6 @@ mod proto_tests { use crate::proto_test_util::{StubEncoder, UnreachableDecoder, column_node}; use arrow::datatypes::Schema; use datafusion_common::DataFusionError; - use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; - use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; - use datafusion_proto_models::protobuf::{self, physical_expr_node}; // ── try_to_proto ───────────────────────────────────────────────────────── @@ -164,8 +165,8 @@ mod proto_tests { assert!(node.expr_id.is_none()); // Verify the encoded name matches the original. - let protobuf::UnknownColumn { name } = match node.expr_type { - Some(physical_expr_node::ExprType::UnknownColumn(c)) => c, + let UnknownColumn { name } = match node.expr_type { + Some(ExprType::UnknownColumn(c)) => c, other => panic!("expected UnknownColumn proto node, got {other:?}"), }; assert_eq!(name, "my_col"); @@ -175,13 +176,11 @@ mod proto_tests { #[test] fn try_from_proto_decodes_name() { - let node = protobuf::PhysicalExprNode { + let node = PhysicalExprNode { expr_id: None, - expr_type: Some(physical_expr_node::ExprType::UnknownColumn( - protobuf::UnknownColumn { - name: "my_col".to_string(), - }, - )), + expr_type: Some(ExprType::UnknownColumn(UnknownColumn { + name: "my_col".to_string(), + })), }; let schema = Schema::empty(); // UnKnownColumn has no child exprs so the decoder is never called. From e953405905b63702bedd078059e5742c51790712 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Thu, 17 Sep 2026 00:29:47 +0800 Subject: [PATCH 3/3] style: follow existing proto hook import convention --- .../physical-expr/src/expressions/column.rs | 50 ++++++------- .../src/expressions/is_not_null.rs | 56 ++++++++------- .../physical-expr/src/expressions/is_null.rs | 54 +++++++------- .../physical-expr/src/expressions/literal.rs | 71 ++++++++++--------- .../physical-expr/src/expressions/negative.rs | 54 +++++++------- .../physical-expr/src/expressions/not.rs | 54 +++++++------- .../src/expressions/unknown_column.rs | 57 +++++++-------- 7 files changed, 211 insertions(+), 185 deletions(-) diff --git a/datafusion/physical-expr/src/expressions/column.rs b/datafusion/physical-expr/src/expressions/column.rs index 1f504acb8a5fd..ea00110ecc516 100644 --- a/datafusion/physical-expr/src/expressions/column.rs +++ b/datafusion/physical-expr/src/expressions/column.rs @@ -31,18 +31,6 @@ use datafusion_common::{Result, internal_err, plan_err}; use datafusion_expr::ColumnarValue; use datafusion_expr_common::placement::ExpressionPlacement; -#[cfg(feature = "proto")] -use datafusion_physical_expr_common::{ - expect_expr_variant, - physical_expr::{ - proto_decode::PhysicalExprDecodeCtx, proto_encode::PhysicalExprEncodeCtx, - }, -}; -#[cfg(feature = "proto")] -use datafusion_proto_models::protobuf::{ - PhysicalColumn, PhysicalExprNode, physical_expr_node::ExprType, -}; - /// Represents the column at a given index in a RecordBatch /// /// This is a physical expression that represents a column at a given index in an @@ -162,29 +150,32 @@ impl PhysicalExpr for Column { #[cfg(feature = "proto")] fn try_to_proto( &self, - _ctx: &PhysicalExprEncodeCtx<'_>, - ) -> Result> { + _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; let Self { name, index } = self; - Ok(Some(PhysicalExprNode { + Ok(Some(protobuf::PhysicalExprNode { expr_id: None, - expr_type: Some(ExprType::Column(PhysicalColumn { - name: name.clone(), - index: *index as u32, - })), + expr_type: Some(protobuf::physical_expr_node::ExprType::Column( + protobuf::PhysicalColumn { + name: name.clone(), + index: *index as u32, + }, + )), })) } } #[cfg(feature = "proto")] -impl From<&PhysicalColumn> for Column { - fn from(c: &PhysicalColumn) -> Self { - let PhysicalColumn { name, index } = c; +impl From<&datafusion_proto_models::protobuf::PhysicalColumn> for Column { + fn from(c: &datafusion_proto_models::protobuf::PhysicalColumn) -> Self { + let datafusion_proto_models::protobuf::PhysicalColumn { name, index } = c; Column::new(name, *index as usize) } } #[cfg(feature = "proto")] -impl From<&Column> for PhysicalColumn { +impl From<&Column> for datafusion_proto_models::protobuf::PhysicalColumn { fn from(c: &Column) -> Self { let Column { name, index } = c; Self { @@ -208,11 +199,16 @@ impl Column { /// [`PhysicalExpr::try_to_proto`]: datafusion_physical_expr_common::physical_expr::PhysicalExpr::try_to_proto /// [`PhysicalExprDecodeCtx::decode`]: datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx::decode pub fn try_from_proto( - node: &PhysicalExprNode, - _ctx: &PhysicalExprDecodeCtx<'_>, + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, ) -> Result> { - let PhysicalColumn { name, index } = - expect_expr_variant!(node, ExprType::Column, "Column"); + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + let protobuf::PhysicalColumn { name, index } = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::Column, + "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 190d85e045f76..0444ecb724d52 100644 --- a/datafusion/physical-expr/src/expressions/is_not_null.rs +++ b/datafusion/physical-expr/src/expressions/is_not_null.rs @@ -27,18 +27,6 @@ use datafusion_expr::ColumnarValue; use std::hash::Hash; use std::sync::Arc; -#[cfg(feature = "proto")] -use datafusion_physical_expr_common::{ - expect_expr_variant, - physical_expr::{ - proto_decode::PhysicalExprDecodeCtx, proto_encode::PhysicalExprEncodeCtx, - }, -}; -#[cfg(feature = "proto")] -use datafusion_proto_models::protobuf::{ - PhysicalExprNode, PhysicalIsNotNull, physical_expr_node::ExprType, -}; - /// IS NOT NULL expression #[derive(Debug, Eq)] pub struct IsNotNullExpr { @@ -118,14 +106,18 @@ impl PhysicalExpr for IsNotNullExpr { #[cfg(feature = "proto")] fn try_to_proto( &self, - ctx: &PhysicalExprEncodeCtx<'_>, - ) -> Result> { + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let Self { arg } = self; - Ok(Some(PhysicalExprNode { + Ok(Some(protobuf::PhysicalExprNode { expr_id: None, - expr_type: Some(ExprType::IsNotNullExpr(Box::new(PhysicalIsNotNull { - expr: Some(Box::new(ctx.encode_child(arg)?)), - }))), + expr_type: Some(protobuf::physical_expr_node::ExprType::IsNotNullExpr( + Box::new(protobuf::PhysicalIsNotNull { + expr: Some(Box::new(ctx.encode_child(arg)?)), + }), + )), })) } } @@ -134,11 +126,18 @@ impl PhysicalExpr for IsNotNullExpr { impl IsNotNullExpr { /// Reconstruct an [`IsNotNullExpr`] from its protobuf representation. pub fn try_from_proto( - node: &PhysicalExprNode, - ctx: &PhysicalExprDecodeCtx<'_>, + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, ) -> Result> { - let node = expect_expr_variant!(node, ExprType::IsNotNullExpr, "IsNotNullExpr"); - let PhysicalIsNotNull { expr } = node.as_ref(); + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let node = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::IsNotNullExpr, + "IsNotNullExpr", + ); + let protobuf::PhysicalIsNotNull { expr } = node.as_ref(); let expr = ctx.decode_required_expression(expr.as_deref(), "IsNotNullExpr", "expr")?; @@ -264,13 +263,18 @@ mod proto_tests { }; use arrow::datatypes::Field; use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalIsNotNull, physical_expr_node, + }; fn is_not_null_node(expr: Option>) -> PhysicalExprNode { PhysicalExprNode { expr_id: None, - expr_type: Some(ExprType::IsNotNullExpr(Box::new(PhysicalIsNotNull { - expr, - }))), + expr_type: Some(physical_expr_node::ExprType::IsNotNullExpr(Box::new( + PhysicalIsNotNull { expr }, + ))), } } @@ -292,7 +296,7 @@ mod proto_tests { assert!(node.expr_id.is_none()); let is_not_null_node = match node.expr_type { - Some(ExprType::IsNotNullExpr(boxed)) => *boxed, + Some(physical_expr_node::ExprType::IsNotNullExpr(boxed)) => *boxed, other => panic!("expected an IsNotNullExpr node, got {other:?}"), }; assert!(is_not_null_node.expr.is_some()); diff --git a/datafusion/physical-expr/src/expressions/is_null.rs b/datafusion/physical-expr/src/expressions/is_null.rs index bfa79d72a2282..08748af2f357c 100644 --- a/datafusion/physical-expr/src/expressions/is_null.rs +++ b/datafusion/physical-expr/src/expressions/is_null.rs @@ -27,18 +27,6 @@ use datafusion_expr::ColumnarValue; use std::hash::Hash; use std::sync::Arc; -#[cfg(feature = "proto")] -use datafusion_physical_expr_common::{ - expect_expr_variant, - physical_expr::{ - proto_decode::PhysicalExprDecodeCtx, proto_encode::PhysicalExprEncodeCtx, - }, -}; -#[cfg(feature = "proto")] -use datafusion_proto_models::protobuf::{ - PhysicalExprNode, PhysicalIsNull, physical_expr_node::ExprType, -}; - /// IS NULL expression #[derive(Debug, Eq)] pub struct IsNullExpr { @@ -117,14 +105,18 @@ impl PhysicalExpr for IsNullExpr { #[cfg(feature = "proto")] fn try_to_proto( &self, - ctx: &PhysicalExprEncodeCtx<'_>, - ) -> Result> { + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let Self { arg } = self; - Ok(Some(PhysicalExprNode { + Ok(Some(protobuf::PhysicalExprNode { expr_id: None, - expr_type: Some(ExprType::IsNullExpr(Box::new(PhysicalIsNull { - expr: Some(Box::new(ctx.encode_child(arg)?)), - }))), + expr_type: Some(protobuf::physical_expr_node::ExprType::IsNullExpr( + Box::new(protobuf::PhysicalIsNull { + expr: Some(Box::new(ctx.encode_child(arg)?)), + }), + )), })) } } @@ -133,11 +125,18 @@ impl PhysicalExpr for IsNullExpr { impl IsNullExpr { /// Reconstruct an [`IsNullExpr`] from its protobuf representation. pub fn try_from_proto( - node: &PhysicalExprNode, - ctx: &PhysicalExprDecodeCtx<'_>, + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, ) -> Result> { - let node = expect_expr_variant!(node, ExprType::IsNullExpr, "IsNullExpr"); - let PhysicalIsNull { expr } = node.as_ref(); + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let node = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::IsNullExpr, + "IsNullExpr", + ); + let protobuf::PhysicalIsNull { expr } = node.as_ref(); let expr = ctx.decode_required_expression(expr.as_deref(), "IsNullExpr", "expr")?; @@ -275,11 +274,18 @@ mod proto_tests { }; use arrow::datatypes::Field; use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalIsNull, physical_expr_node, + }; fn is_null_node(expr: Option>) -> PhysicalExprNode { PhysicalExprNode { expr_id: None, - expr_type: Some(ExprType::IsNullExpr(Box::new(PhysicalIsNull { expr }))), + expr_type: Some(physical_expr_node::ExprType::IsNullExpr(Box::new( + PhysicalIsNull { expr }, + ))), } } @@ -301,7 +307,7 @@ mod proto_tests { assert!(node.expr_id.is_none()); let is_null_node = match node.expr_type { - Some(ExprType::IsNullExpr(boxed)) => *boxed, + Some(physical_expr_node::ExprType::IsNullExpr(boxed)) => *boxed, other => panic!("expected an IsNullExpr node, got {other:?}"), }; assert!(is_null_node.expr.is_some()); diff --git a/datafusion/physical-expr/src/expressions/literal.rs b/datafusion/physical-expr/src/expressions/literal.rs index e68decbc3e102..3a480f901427c 100644 --- a/datafusion/physical-expr/src/expressions/literal.rs +++ b/datafusion/physical-expr/src/expressions/literal.rs @@ -35,18 +35,6 @@ use datafusion_expr_common::interval_arithmetic::Interval; use datafusion_expr_common::placement::ExpressionPlacement; use datafusion_expr_common::sort_properties::{ExprProperties, SortProperties}; -#[cfg(feature = "proto")] -use datafusion_common::{internal_datafusion_err, internal_err}; -#[cfg(feature = "proto")] -use datafusion_physical_expr_common::physical_expr::{ - proto_decode::PhysicalExprDecodeCtx, proto_encode::PhysicalExprEncodeCtx, -}; -#[cfg(feature = "proto")] -use datafusion_proto_models::{ - datafusion_common::ScalarValue as ProtoScalarValue, - protobuf::{PhysicalExprNode, PhysicalLiteralNode, physical_expr_node::ExprType}, -}; - /// Represents a literal value #[derive(Debug, PartialEq, Eq, Clone)] pub struct Literal { @@ -151,19 +139,23 @@ impl PhysicalExpr for Literal { #[cfg(feature = "proto")] fn try_to_proto( &self, - _ctx: &PhysicalExprEncodeCtx<'_>, - ) -> Result> { + _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> 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() { - ExprType::Literal(value.try_into()?) + protobuf::physical_expr_node::ExprType::Literal(value.try_into()?) } else { - ExprType::LiteralWithMetadata(PhysicalLiteralNode { - value: Some(value.try_into()?), - metadata: field.metadata().clone(), - }) + protobuf::physical_expr_node::ExprType::LiteralWithMetadata( + protobuf::PhysicalLiteralNode { + value: Some(value.try_into()?), + metadata: field.metadata().clone(), + }, + ) }; - Ok(Some(PhysicalExprNode { + Ok(Some(protobuf::PhysicalExprNode { expr_id: None, expr_type: Some(expr_type), })) @@ -174,23 +166,27 @@ impl PhysicalExpr for Literal { impl Literal { /// Reconstruct a [`Literal`] from its protobuf representation. pub fn try_from_proto( - node: &PhysicalExprNode, - _ctx: &PhysicalExprDecodeCtx<'_>, + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, ) -> Result> { - let PhysicalExprNode { + use datafusion_common::{internal_datafusion_err, internal_err}; + use datafusion_proto_models::protobuf; + 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 ProtoScalarValue { + 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(PhysicalLiteralNode { + Some(ExprType::LiteralWithMetadata(protobuf::PhysicalLiteralNode { value, metadata, })) => { @@ -270,6 +266,9 @@ mod proto_tests { use super::*; use crate::proto_test_util::{StubEncoder, UnreachableDecoder, column_node}; use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::physical_expr_node; fn i32_literal() -> Literal { Literal::new(ScalarValue::Int32(Some(42))) @@ -291,7 +290,10 @@ mod proto_tests { // Literal nodes never set expr_id. assert!(node.expr_id.is_none()); // Variant must be Literal, not any other expr type. - assert!(matches!(node.expr_type, Some(ExprType::Literal(_)))); + assert!(matches!( + node.expr_type, + Some(physical_expr_node::ExprType::Literal(_)) + )); } #[test] @@ -305,7 +307,10 @@ mod proto_tests { .unwrap() .expect("null Literal should encode to Some(node)"); - assert!(matches!(node.expr_type, Some(ExprType::Literal(_)))); + assert!(matches!( + node.expr_type, + Some(physical_expr_node::ExprType::Literal(_)) + )); // Decode and verify the null payload round-trips correctly. let schema = Schema::empty(); @@ -344,12 +349,14 @@ mod proto_tests { #[test] fn try_from_proto_rejects_missing_value() { - let node = PhysicalExprNode { + let node = datafusion_proto_models::protobuf::PhysicalExprNode { expr_id: None, - expr_type: Some(ExprType::LiteralWithMetadata(PhysicalLiteralNode { - value: None, - metadata: Default::default(), - })), + 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; diff --git a/datafusion/physical-expr/src/expressions/negative.rs b/datafusion/physical-expr/src/expressions/negative.rs index fcdedcfedc272..a932c8d5220f6 100644 --- a/datafusion/physical-expr/src/expressions/negative.rs +++ b/datafusion/physical-expr/src/expressions/negative.rs @@ -40,18 +40,6 @@ use datafusion_expr::{ type_coercion::{is_interval, is_signed_numeric, is_timestamp}, }; -#[cfg(feature = "proto")] -use datafusion_physical_expr_common::{ - expect_expr_variant, - physical_expr::{ - proto_decode::PhysicalExprDecodeCtx, proto_encode::PhysicalExprEncodeCtx, - }, -}; -#[cfg(feature = "proto")] -use datafusion_proto_models::protobuf::{ - PhysicalExprNode, PhysicalNegativeNode, physical_expr_node::ExprType, -}; - /// Negative expression #[derive(Debug, Eq)] pub struct NegativeExpr { @@ -192,14 +180,18 @@ impl PhysicalExpr for NegativeExpr { #[cfg(feature = "proto")] fn try_to_proto( &self, - ctx: &PhysicalExprEncodeCtx<'_>, - ) -> Result> { + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let Self { arg } = self; - Ok(Some(PhysicalExprNode { + Ok(Some(protobuf::PhysicalExprNode { expr_id: None, - expr_type: Some(ExprType::Negative(Box::new(PhysicalNegativeNode { - expr: Some(Box::new(ctx.encode_child(arg)?)), - }))), + expr_type: Some(protobuf::physical_expr_node::ExprType::Negative(Box::new( + protobuf::PhysicalNegativeNode { + expr: Some(Box::new(ctx.encode_child(arg)?)), + }, + ))), })) } } @@ -208,11 +200,18 @@ impl PhysicalExpr for NegativeExpr { impl NegativeExpr { /// Reconstruct a [`NegativeExpr`] from its protobuf representation. pub fn try_from_proto( - node: &PhysicalExprNode, - ctx: &PhysicalExprDecodeCtx<'_>, + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, ) -> Result> { - let n = expect_expr_variant!(node, ExprType::Negative, "Negative"); - let PhysicalNegativeNode { expr } = n.as_ref(); + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let n = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::Negative, + "Negative", + ); + let protobuf::PhysicalNegativeNode { expr } = n.as_ref(); let expr = ctx.decode_required_expression(expr.as_deref(), "NegativeExpr", "expr")?; @@ -456,12 +455,19 @@ mod proto_tests { }; use arrow::datatypes::Field; use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalNegativeNode, physical_expr_node, + }; /// Build a `NegativeExpr` proto node with the given children. fn negative_node(expr: Option>) -> PhysicalExprNode { PhysicalExprNode { expr_id: None, - expr_type: Some(ExprType::Negative(Box::new(PhysicalNegativeNode { expr }))), + expr_type: Some(physical_expr_node::ExprType::Negative(Box::new( + PhysicalNegativeNode { expr }, + ))), } } @@ -484,7 +490,7 @@ mod proto_tests { assert!(node.expr_id.is_none()); let negative_node = match node.expr_type { - Some(ExprType::Negative(boxed)) => *boxed, + Some(physical_expr_node::ExprType::Negative(boxed)) => *boxed, other => panic!("expected a NegativeExpr node, got {other:?}"), }; assert!(negative_node.expr.is_some()); diff --git a/datafusion/physical-expr/src/expressions/not.rs b/datafusion/physical-expr/src/expressions/not.rs index 4c2ac97771054..7699c8944ea23 100644 --- a/datafusion/physical-expr/src/expressions/not.rs +++ b/datafusion/physical-expr/src/expressions/not.rs @@ -31,18 +31,6 @@ use datafusion_expr::interval_arithmetic::Interval; #[expect(deprecated)] use datafusion_expr::statistics::Distribution::{self, Bernoulli}; -#[cfg(feature = "proto")] -use datafusion_physical_expr_common::{ - expect_expr_variant, - physical_expr::{ - proto_decode::PhysicalExprDecodeCtx, proto_encode::PhysicalExprEncodeCtx, - }, -}; -#[cfg(feature = "proto")] -use datafusion_proto_models::protobuf::{ - PhysicalExprNode, PhysicalNot, physical_expr_node::ExprType, -}; - /// Not expression #[derive(Debug, Eq)] pub struct NotExpr { @@ -197,14 +185,18 @@ impl PhysicalExpr for NotExpr { #[cfg(feature = "proto")] fn try_to_proto( &self, - ctx: &PhysicalExprEncodeCtx<'_>, - ) -> Result> { + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let Self { arg } = self; - Ok(Some(PhysicalExprNode { + Ok(Some(protobuf::PhysicalExprNode { expr_id: None, - expr_type: Some(ExprType::NotExpr(Box::new(PhysicalNot { - expr: Some(Box::new(ctx.encode_child(arg)?)), - }))), + expr_type: Some(protobuf::physical_expr_node::ExprType::NotExpr(Box::new( + protobuf::PhysicalNot { + expr: Some(Box::new(ctx.encode_child(arg)?)), + }, + ))), })) } } @@ -213,11 +205,18 @@ impl PhysicalExpr for NotExpr { impl NotExpr { /// Reconstruct a [`NotExpr`] from its protobuf representation. pub fn try_from_proto( - node: &PhysicalExprNode, - ctx: &PhysicalExprDecodeCtx<'_>, + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, ) -> Result> { - let not_expr = expect_expr_variant!(node, ExprType::NotExpr, "NotExpr"); - let PhysicalNot { expr } = not_expr.as_ref(); + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let not_expr = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::NotExpr, + "NotExpr", + ); + 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))) @@ -409,12 +408,19 @@ mod proto_tests { }; use arrow::datatypes::Field; use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalNot, physical_expr_node, + }; /// Build a `NotExpr` proto node with the given child. fn not_node(expr: Option>) -> PhysicalExprNode { PhysicalExprNode { expr_id: None, - expr_type: Some(ExprType::NotExpr(Box::new(PhysicalNot { expr }))), + expr_type: Some(physical_expr_node::ExprType::NotExpr(Box::new( + PhysicalNot { expr }, + ))), } } @@ -437,7 +443,7 @@ mod proto_tests { assert!(node.expr_id.is_none()); let not_node = match node.expr_type { - Some(ExprType::NotExpr(boxed)) => *boxed, + Some(physical_expr_node::ExprType::NotExpr(boxed)) => *boxed, other => panic!("expected a NotExpr node, got {other:?}"), }; assert!(not_node.expr.is_some()); diff --git a/datafusion/physical-expr/src/expressions/unknown_column.rs b/datafusion/physical-expr/src/expressions/unknown_column.rs index c9df5eefcdf40..c3b6dd9dafade 100644 --- a/datafusion/physical-expr/src/expressions/unknown_column.rs +++ b/datafusion/physical-expr/src/expressions/unknown_column.rs @@ -30,18 +30,6 @@ use datafusion_common::{Result, internal_err}; use datafusion_expr::ColumnarValue; -#[cfg(feature = "proto")] -use datafusion_physical_expr_common::{ - expect_expr_variant, - physical_expr::{ - proto_decode::PhysicalExprDecodeCtx, proto_encode::PhysicalExprEncodeCtx, - }, -}; -#[cfg(feature = "proto")] -use datafusion_proto_models::protobuf::{ - PhysicalExprNode, UnknownColumn, physical_expr_node::ExprType, -}; - #[derive(Debug, Clone, Eq)] pub struct UnKnownColumn { name: String, @@ -101,14 +89,16 @@ impl PhysicalExpr for UnKnownColumn { #[cfg(feature = "proto")] fn try_to_proto( &self, - _ctx: &PhysicalExprEncodeCtx<'_>, - ) -> Result> { + _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let Self { name } = self; - Ok(Some(PhysicalExprNode { + Ok(Some(protobuf::PhysicalExprNode { expr_id: None, - expr_type: Some(ExprType::UnknownColumn(UnknownColumn { - name: name.clone(), - })), + expr_type: Some(protobuf::physical_expr_node::ExprType::UnknownColumn( + protobuf::UnknownColumn { name: name.clone() }, + )), })) } } @@ -117,11 +107,17 @@ impl PhysicalExpr for UnKnownColumn { impl UnKnownColumn { /// Reconstruct an [`UnKnownColumn`] from its protobuf representation. pub fn try_from_proto( - node: &PhysicalExprNode, - _ctx: &PhysicalExprDecodeCtx<'_>, + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, ) -> Result> { - let UnknownColumn { name } = - expect_expr_variant!(node, ExprType::UnknownColumn, "UnKnownColumn"); + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let protobuf::UnknownColumn { name } = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::UnknownColumn, + "UnKnownColumn", + ); Ok(Arc::new(UnKnownColumn::new(name))) } } @@ -147,6 +143,9 @@ mod proto_tests { use crate::proto_test_util::{StubEncoder, UnreachableDecoder, column_node}; use arrow::datatypes::Schema; use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::{self, physical_expr_node}; // ── try_to_proto ───────────────────────────────────────────────────────── @@ -165,8 +164,8 @@ mod proto_tests { assert!(node.expr_id.is_none()); // Verify the encoded name matches the original. - let UnknownColumn { name } = match node.expr_type { - Some(ExprType::UnknownColumn(c)) => c, + let protobuf::UnknownColumn { name } = match node.expr_type { + Some(physical_expr_node::ExprType::UnknownColumn(c)) => c, other => panic!("expected UnknownColumn proto node, got {other:?}"), }; assert_eq!(name, "my_col"); @@ -176,11 +175,13 @@ mod proto_tests { #[test] fn try_from_proto_decodes_name() { - let node = PhysicalExprNode { + let node = protobuf::PhysicalExprNode { expr_id: None, - expr_type: Some(ExprType::UnknownColumn(UnknownColumn { - name: "my_col".to_string(), - })), + expr_type: Some(physical_expr_node::ExprType::UnknownColumn( + protobuf::UnknownColumn { + name: "my_col".to_string(), + }, + )), }; let schema = Schema::empty(); // UnKnownColumn has no child exprs so the decoder is never called.