From ef526120f517076ae59b5289fa1f6b4a58ca2485 Mon Sep 17 00:00:00 2001 From: eric_song Date: Thu, 28 May 2026 01:04:23 +1000 Subject: [PATCH 1/4] Port CastExpr to proto hooks --- .../physical-expr/src/expressions/cast.rs | 57 +++++++++++++++++++ .../proto/src/physical_plan/from_proto.rs | 12 +--- .../proto/src/physical_plan/to_proto.rs | 16 +----- 3 files changed, 60 insertions(+), 25 deletions(-) diff --git a/datafusion/physical-expr/src/expressions/cast.rs b/datafusion/physical-expr/src/expressions/cast.rs index ad214a89ceb71..2e097deabca3a 100644 --- a/datafusion/physical-expr/src/expressions/cast.rs +++ b/datafusion/physical-expr/src/expressions/cast.rs @@ -298,6 +298,63 @@ impl PhysicalExpr for CastExpr { write!(f, ")") } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Cast(Box::new( + protobuf::PhysicalCastNode { + expr: Some(Box::new(ctx.encode_child(self.expr())?)), + arrow_type: Some(self.cast_type().try_into()?), + }, + ))), + })) + } +} + +#[cfg(feature = "proto")] +impl CastExpr { + /// Reconstruct a [`CastExpr`] from its protobuf representation. + /// + /// Takes the whole [`PhysicalExprNode`] so the decode signature matches + /// other migrated expressions and can inspect outer-node metadata if + /// needed in the future. + /// + /// [`PhysicalExprNode`]: datafusion_proto_models::protobuf::PhysicalExprNode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::DataFusionError; + use datafusion_common::internal_err; + use datafusion_proto_models::protobuf; + + let cast_expr = match &node.expr_type { + Some(protobuf::physical_expr_node::ExprType::Cast(cast_expr)) => { + cast_expr.as_ref() + } + _ => return internal_err!("PhysicalExprNode is not a CastExpr"), + }; + + let expr = ctx.decode_required_expression( + cast_expr.expr.as_deref(), + "CastExpr", + "expr", + )?; + let arrow_type = cast_expr.arrow_type.as_ref().ok_or_else(|| { + DataFusionError::Internal( + "CastExpr is missing required field 'arrow_type'".to_string(), + ) + })?; + + Ok(Arc::new(CastExpr::new(expr, arrow_type.try_into()?, None))) + } } /// Return a PhysicalExpression representing `expr` casted to diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 7d2e68d810959..1299575a0576b 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -360,17 +360,7 @@ pub fn parse_physical_expr_with_converter( }) .transpose()?, )?), - ExprType::Cast(e) => Arc::new(CastExpr::new( - parse_required_physical_expr( - e.expr.as_deref(), - ctx, - "expr", - input_schema, - proto_converter, - )?, - convert_required!(e.arrow_type)?, - None, - )), + ExprType::Cast(_) => CastExpr::try_from_proto(proto, &decode_ctx)?, ExprType::TryCast(e) => Arc::new(TryCastExpr::new( parse_required_physical_expr( e.expr.as_deref(), diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 9cb9e897605be..0813a5aaf1c00 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -36,8 +36,8 @@ use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; use datafusion_physical_expr::window::{SlidingAggregateWindowExpr, StandardWindowExpr}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_physical_plan::expressions::{ - CaseExpr, CastExpr, DynamicFilterPhysicalExpr, IsNotNullExpr, IsNullExpr, Literal, - NotExpr, TryCastExpr, UnKnownColumn, + CaseExpr, DynamicFilterPhysicalExpr, IsNotNullExpr, IsNullExpr, Literal, NotExpr, + TryCastExpr, UnKnownColumn, }; use datafusion_physical_plan::joins::HashExpr; use datafusion_physical_plan::udaf::AggregateFunctionExpr; @@ -396,18 +396,6 @@ pub fn serialize_physical_expr_with_converter( lit.value().try_into()?, )), }) - } else if let Some(cast) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::Cast(Box::new( - protobuf::PhysicalCastNode { - expr: Some(Box::new( - proto_converter.physical_expr_to_proto(cast.expr(), codec)?, - )), - arrow_type: Some(cast.cast_type().try_into()?), - }, - ))), - }) } else if let Some(cast) = expr.downcast_ref::() { Ok(protobuf::PhysicalExprNode { expr_id, From 13846925243bcad8f9358bfb84e77b251c1b2fd4 Mon Sep 17 00:00:00 2001 From: eric_song Date: Thu, 28 May 2026 01:04:36 +1000 Subject: [PATCH 2/4] Add CastExpr proto hook tests --- .../physical-expr/src/expressions/cast.rs | 179 ++++++++++++++++++ 1 file changed, 179 insertions(+) diff --git a/datafusion/physical-expr/src/expressions/cast.rs b/datafusion/physical-expr/src/expressions/cast.rs index 2e097deabca3a..76a17dae4fd3f 100644 --- a/datafusion/physical-expr/src/expressions/cast.rs +++ b/datafusion/physical-expr/src/expressions/cast.rs @@ -1210,4 +1210,183 @@ mod tests { Ok(()) } + + #[cfg(feature = "proto")] + fn proto_cast_fixture() -> CastExpr { + let schema = Schema::new(vec![Field::new("a", Int32, false)]); + CastExpr::new(col("a", &schema).unwrap(), Int64, None) + } + + #[cfg(feature = "proto")] + fn proto_int64_arrow_type() -> datafusion_proto_models::datafusion_common::ArrowType { + (&Int64).try_into().unwrap() + } + + #[cfg(feature = "proto")] + fn proto_cast_node( + expr: Option>, + arrow_type: Option, + ) -> datafusion_proto_models::protobuf::PhysicalExprNode { + use datafusion_proto_models::protobuf; + + protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Cast(Box::new( + protobuf::PhysicalCastNode { expr, arrow_type }, + ))), + } + } + + #[test] + #[cfg(feature = "proto")] + fn try_to_proto_encodes_cast_expr() { + use crate::proto_test_util::StubEncoder; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::physical_expr_node; + + let cast = proto_cast_fixture(); + let encoder = StubEncoder::ok(); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = cast + .try_to_proto(&ctx) + .unwrap() + .expect("CastExpr should encode to Some(node)"); + + assert!(node.expr_id.is_none()); + let cast_node = match node.expr_type { + Some(physical_expr_node::ExprType::Cast(cast_node)) => *cast_node, + other => panic!("expected a Cast node, got {other:?}"), + }; + assert!(cast_node.expr.is_some()); + + let arrow_type = cast_node + .arrow_type + .as_ref() + .expect("cast type should be encoded"); + let data_type: DataType = arrow_type.try_into().unwrap(); + assert_eq!(data_type, Int64); + } + + #[test] + #[cfg(feature = "proto")] + fn try_to_proto_propagates_child_encode_error() { + use crate::proto_test_util::StubEncoder; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + + let cast = proto_cast_fixture(); + let encoder = StubEncoder::failing_on(1); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let err = cast.try_to_proto(&ctx).unwrap_err(); + assert!(matches!( + err, + datafusion_common::DataFusionError::Internal(msg) if msg.contains("call 1") + )); + } + + #[test] + #[cfg(feature = "proto")] + fn try_from_proto_decodes_cast_expr() { + use crate::proto_test_util::{StubDecoder, column_node}; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + + let node = proto_cast_node( + Some(Box::new(column_node("a"))), + Some(proto_int64_arrow_type()), + ); + let schema = Schema::empty(); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let decoded = CastExpr::try_from_proto(&node, &ctx).unwrap(); + let cast = decoded + .downcast_ref::() + .expect("decoded expr should be a CastExpr"); + + assert_eq!(cast.cast_type(), &Int64); + assert!( + cast.expr() + .downcast_ref::() + .is_some() + ); + } + + #[test] + #[cfg(feature = "proto")] + fn try_from_proto_rejects_non_cast_node() { + use crate::proto_test_util::{UnreachableDecoder, column_node}; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + + let node = column_node("a"); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CastExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!( + err, + datafusion_common::DataFusionError::Internal(msg) + if msg.contains("PhysicalExprNode is not a CastExpr") + )); + } + + #[test] + #[cfg(feature = "proto")] + fn try_from_proto_rejects_missing_expr() { + use crate::proto_test_util::UnreachableDecoder; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + + let node = proto_cast_node(None, Some(proto_int64_arrow_type())); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CastExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!( + err, + datafusion_common::DataFusionError::Internal(msg) + if msg.contains("CastExpr is missing required field 'expr'") + )); + } + + #[test] + #[cfg(feature = "proto")] + fn try_from_proto_rejects_missing_arrow_type() { + use crate::proto_test_util::{StubDecoder, column_node}; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + + let node = proto_cast_node(Some(Box::new(column_node("a"))), None); + let schema = Schema::empty(); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CastExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!( + err, + datafusion_common::DataFusionError::Internal(msg) + if msg.contains("CastExpr is missing required field 'arrow_type'") + )); + } + + #[test] + #[cfg(feature = "proto")] + fn try_from_proto_propagates_child_decode_error() { + use crate::proto_test_util::{StubDecoder, column_node}; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + + let node = proto_cast_node( + Some(Box::new(column_node("a"))), + Some(proto_int64_arrow_type()), + ); + let schema = Schema::empty(); + let decoder = StubDecoder::failing_on(1); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CastExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!( + err, + datafusion_common::DataFusionError::Internal(msg) if msg.contains("call 1") + )); + } } From cc992d1857e123d2d7e9df5745ef24d36dbd2f0a Mon Sep 17 00:00:00 2001 From: eric_song Date: Thu, 28 May 2026 01:16:06 +1000 Subject: [PATCH 3/4] Use internal error helper for CastExpr proto --- datafusion/physical-expr/src/expressions/cast.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-expr/src/expressions/cast.rs b/datafusion/physical-expr/src/expressions/cast.rs index 76a17dae4fd3f..4e1bf5a88709c 100644 --- a/datafusion/physical-expr/src/expressions/cast.rs +++ b/datafusion/physical-expr/src/expressions/cast.rs @@ -331,7 +331,7 @@ impl CastExpr { node: &datafusion_proto_models::protobuf::PhysicalExprNode, ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, ) -> Result> { - use datafusion_common::DataFusionError; + use datafusion_common::internal_datafusion_err; use datafusion_common::internal_err; use datafusion_proto_models::protobuf; @@ -348,9 +348,7 @@ impl CastExpr { "expr", )?; let arrow_type = cast_expr.arrow_type.as_ref().ok_or_else(|| { - DataFusionError::Internal( - "CastExpr is missing required field 'arrow_type'".to_string(), - ) + internal_datafusion_err!("CastExpr is missing required field 'arrow_type'") })?; Ok(Arc::new(CastExpr::new(expr, arrow_type.try_into()?, None))) From b3f786d541395197fac49cf9e1089a73f5f9c418 Mon Sep 17 00:00:00 2001 From: eric_song Date: Fri, 29 May 2026 01:24:54 +1000 Subject: [PATCH 4/4] Move CastExpr proto tests to module --- .../physical-expr/src/expressions/cast.rs | 84 ++++++++----------- 1 file changed, 33 insertions(+), 51 deletions(-) diff --git a/datafusion/physical-expr/src/expressions/cast.rs b/datafusion/physical-expr/src/expressions/cast.rs index 4e1bf5a88709c..26f06b546ad1d 100644 --- a/datafusion/physical-expr/src/expressions/cast.rs +++ b/datafusion/physical-expr/src/expressions/cast.rs @@ -1208,40 +1208,50 @@ mod tests { Ok(()) } +} - #[cfg(feature = "proto")] +/// Tests for the `try_to_proto` / `try_from_proto` hooks. +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::*; + use crate::expressions::{Column, col}; + use crate::proto_test_util::{ + StubDecoder, StubEncoder, UnreachableDecoder, column_node, + }; + 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::datafusion_common::ArrowType; + use datafusion_proto_models::protobuf::{ + PhysicalCastNode, PhysicalExprNode, physical_expr_node, + }; + + /// A `CastExpr` over an `Int32` column, casting to `Int64`. fn proto_cast_fixture() -> CastExpr { let schema = Schema::new(vec![Field::new("a", Int32, false)]); CastExpr::new(col("a", &schema).unwrap(), Int64, None) } - #[cfg(feature = "proto")] - fn proto_int64_arrow_type() -> datafusion_proto_models::datafusion_common::ArrowType { + fn proto_int64_arrow_type() -> ArrowType { (&Int64).try_into().unwrap() } - #[cfg(feature = "proto")] + /// Build a `CastExpr` proto node with the given child and target type. fn proto_cast_node( - expr: Option>, - arrow_type: Option, - ) -> datafusion_proto_models::protobuf::PhysicalExprNode { - use datafusion_proto_models::protobuf; - - protobuf::PhysicalExprNode { + expr: Option>, + arrow_type: Option, + ) -> PhysicalExprNode { + PhysicalExprNode { expr_id: None, - expr_type: Some(protobuf::physical_expr_node::ExprType::Cast(Box::new( - protobuf::PhysicalCastNode { expr, arrow_type }, + expr_type: Some(physical_expr_node::ExprType::Cast(Box::new( + PhysicalCastNode { expr, arrow_type }, ))), } } #[test] - #[cfg(feature = "proto")] fn try_to_proto_encodes_cast_expr() { - use crate::proto_test_util::StubEncoder; - use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; - use datafusion_proto_models::protobuf::physical_expr_node; - let cast = proto_cast_fixture(); let encoder = StubEncoder::ok(); let ctx = PhysicalExprEncodeCtx::new(&encoder); @@ -1267,11 +1277,7 @@ mod tests { } #[test] - #[cfg(feature = "proto")] fn try_to_proto_propagates_child_encode_error() { - use crate::proto_test_util::StubEncoder; - use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; - let cast = proto_cast_fixture(); let encoder = StubEncoder::failing_on(1); let ctx = PhysicalExprEncodeCtx::new(&encoder); @@ -1279,16 +1285,12 @@ mod tests { let err = cast.try_to_proto(&ctx).unwrap_err(); assert!(matches!( err, - datafusion_common::DataFusionError::Internal(msg) if msg.contains("call 1") + DataFusionError::Internal(msg) if msg.contains("call 1") )); } #[test] - #[cfg(feature = "proto")] fn try_from_proto_decodes_cast_expr() { - use crate::proto_test_util::{StubDecoder, column_node}; - use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; - let node = proto_cast_node( Some(Box::new(column_node("a"))), Some(proto_int64_arrow_type()), @@ -1303,19 +1305,11 @@ mod tests { .expect("decoded expr should be a CastExpr"); assert_eq!(cast.cast_type(), &Int64); - assert!( - cast.expr() - .downcast_ref::() - .is_some() - ); + assert!(cast.expr().downcast_ref::().is_some()); } #[test] - #[cfg(feature = "proto")] fn try_from_proto_rejects_non_cast_node() { - use crate::proto_test_util::{UnreachableDecoder, column_node}; - use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; - let node = column_node("a"); let schema = Schema::empty(); let decoder = UnreachableDecoder; @@ -1324,17 +1318,13 @@ mod tests { let err = CastExpr::try_from_proto(&node, &ctx).unwrap_err(); assert!(matches!( err, - datafusion_common::DataFusionError::Internal(msg) + DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a CastExpr") )); } #[test] - #[cfg(feature = "proto")] fn try_from_proto_rejects_missing_expr() { - use crate::proto_test_util::UnreachableDecoder; - use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; - let node = proto_cast_node(None, Some(proto_int64_arrow_type())); let schema = Schema::empty(); let decoder = UnreachableDecoder; @@ -1343,17 +1333,13 @@ mod tests { let err = CastExpr::try_from_proto(&node, &ctx).unwrap_err(); assert!(matches!( err, - datafusion_common::DataFusionError::Internal(msg) + DataFusionError::Internal(msg) if msg.contains("CastExpr is missing required field 'expr'") )); } #[test] - #[cfg(feature = "proto")] fn try_from_proto_rejects_missing_arrow_type() { - use crate::proto_test_util::{StubDecoder, column_node}; - use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; - let node = proto_cast_node(Some(Box::new(column_node("a"))), None); let schema = Schema::empty(); let decoder = StubDecoder::ok(); @@ -1362,17 +1348,13 @@ mod tests { let err = CastExpr::try_from_proto(&node, &ctx).unwrap_err(); assert!(matches!( err, - datafusion_common::DataFusionError::Internal(msg) + DataFusionError::Internal(msg) if msg.contains("CastExpr is missing required field 'arrow_type'") )); } #[test] - #[cfg(feature = "proto")] fn try_from_proto_propagates_child_decode_error() { - use crate::proto_test_util::{StubDecoder, column_node}; - use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; - let node = proto_cast_node( Some(Box::new(column_node("a"))), Some(proto_int64_arrow_type()), @@ -1384,7 +1366,7 @@ mod tests { let err = CastExpr::try_from_proto(&node, &ctx).unwrap_err(); assert!(matches!( err, - datafusion_common::DataFusionError::Internal(msg) if msg.contains("call 1") + DataFusionError::Internal(msg) if msg.contains("call 1") )); } }