From a064b10c924199bb75fe754c197d1c1f18e43229 Mon Sep 17 00:00:00 2001 From: ninglinLiu Date: Sun, 25 Oct 2026 18:39:58 +0800 Subject: [PATCH] fix(proto): add recursion protection to physical plan and expression serde --- .../proto/src/physical_plan/from_proto.rs | 1 + datafusion/proto/src/physical_plan/mod.rs | 2 + .../proto/src/physical_plan/to_proto.rs | 1 + datafusion/proto/tests/cases/stack_safety.rs | 97 ++++++++++++++++++- 4 files changed, 98 insertions(+), 3 deletions(-) diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index e47c54f2b7a16..a0e46c34ab29b 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -242,6 +242,7 @@ pub fn parse_physical_expr( /// any scoped state needed during recursive deserialization. /// * `proto_converter` - Converter hooks used for recursive physical plan and /// expression deserialization. +#[cfg_attr(feature = "recursive_protection", recursive::recursive)] pub fn parse_physical_expr_with_converter( proto: &protobuf::PhysicalExprNode, input_schema: &Schema, diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index f72a9d8d0607d..0fc7c92fe0b68 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -1203,6 +1203,7 @@ pub trait PhysicalPlanNodeExt: Sized { self.try_into_physical_plan_with_context(&decode_ctx, proto_converter) } + #[cfg_attr(feature = "recursive_protection", recursive::recursive)] fn try_into_physical_plan_with_context( &self, ctx: &PhysicalPlanDecodeContext<'_>, @@ -1369,6 +1370,7 @@ pub trait PhysicalPlanNodeExt: Sized { } } + #[cfg_attr(feature = "recursive_protection", recursive::recursive)] fn try_from_physical_plan_with_converter( plan: Arc, codec: &dyn PhysicalExtensionCodec, diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 97101db471907..ceac55fa0f14a 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -266,6 +266,7 @@ impl datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprE /// serialization of udfs requiring specialized serialization (see [`PhysicalExtensionCodec::try_encode_udf`]). /// A [`PhysicalProtoConverterExtension`] can be provided to handle the /// conversion process (see [`PhysicalProtoConverterExtension::physical_expr_to_proto`]). +#[cfg_attr(feature = "recursive_protection", recursive::recursive)] pub fn serialize_physical_expr_with_converter( value: &Arc, codec: &dyn PhysicalExtensionCodec, diff --git a/datafusion/proto/tests/cases/stack_safety.rs b/datafusion/proto/tests/cases/stack_safety.rs index 5caf4119a7186..3427d085a7ef3 100644 --- a/datafusion/proto/tests/cases/stack_safety.rs +++ b/datafusion/proto/tests/cases/stack_safety.rs @@ -59,9 +59,13 @@ fn serialize_on_two_mib_stack(alias_depth: usize) { } fn run_in_child(test_name: &str, alias_depth: usize) { + run_in_child_with(test_name, alias_depth, serialize_on_two_mib_stack); +} + +fn run_in_child_with(test_name: &str, depth: usize, run: fn(usize)) { if std::env::var_os(CHILD_ENV).is_some() { - let alias_depth = std::env::var(ALIAS_DEPTH_ENV).unwrap().parse().unwrap(); - serialize_on_two_mib_stack(alias_depth); + let depth = std::env::var(ALIAS_DEPTH_ENV).unwrap().parse().unwrap(); + run(depth); return; } @@ -70,7 +74,7 @@ fn run_in_child(test_name: &str, alias_depth: usize) { let output = Command::new(std::env::current_exe().unwrap()) .args(["--exact", test_name, "--nocapture"]) .env(CHILD_ENV, "1") - .env(ALIAS_DEPTH_ENV, alias_depth.to_string()) + .env(ALIAS_DEPTH_ENV, depth.to_string()) .output() .unwrap(); @@ -97,3 +101,90 @@ fn deeply_nested_logical_plan_serialization_uses_a_growable_stack() { // exercising the `recursive` stack-growth checkpoint. run_in_child(GROWABLE_STACK_TEST_NAME, 2_000); } + +// Physical plan serde (#26738): nested plans go through +// `try_from_physical_plan_with_converter` / `try_into_physical_plan_with_context` +// recursively, so a deep plan needs the same stack-growth checkpoint. +#[cfg(feature = "recursive_protection")] +mod physical { + use std::sync::Arc; + + use datafusion::arrow::datatypes::Schema; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::expressions::{NotExpr, lit}; + use datafusion::physical_plan::limit::GlobalLimitExec; + use datafusion::physical_plan::{ExecutionPlan, PhysicalExpr}; + use datafusion::prelude::SessionContext; + use datafusion_proto::physical_plan::from_proto::parse_physical_expr; + use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; + use datafusion_proto::physical_plan::{ + AsExecutionPlan, DefaultPhysicalExtensionCodec, + }; + use datafusion_proto::protobuf::PhysicalPlanNode; + + use super::run_in_child_with; + + const TEST_NAME: &str = "cases::stack_safety::physical::deeply_nested_physical_plan_roundtrip_uses_a_growable_stack"; + + fn deeply_nested_plan(depth: usize) -> Arc { + let mut plan: Arc = + Arc::new(EmptyExec::new(Arc::new(Schema::empty()))); + for _ in 0..depth { + plan = Arc::new(GlobalLimitExec::new(plan, 0, Some(1))); + } + plan + } + + fn roundtrip_on_two_mib_stack(depth: usize) { + std::thread::Builder::new() + .name("two-megabyte-stack".into()) + .stack_size(2 * 1024 * 1024) + .spawn(move || { + let codec = DefaultPhysicalExtensionCodec {}; + let proto = PhysicalPlanNode::try_from_physical_plan( + deeply_nested_plan(depth), + &codec, + ) + .unwrap(); + let ctx = SessionContext::new(); + proto + .try_into_physical_plan(&ctx.task_ctx(), &codec) + .unwrap(); + }) + .unwrap() + .join() + .unwrap(); + } + + #[test] + fn deeply_nested_physical_plan_roundtrip_uses_a_growable_stack() { + run_in_child_with(TEST_NAME, 2_000, roundtrip_on_two_mib_stack); + } + + const EXPR_TEST_NAME: &str = "cases::stack_safety::physical::deeply_nested_physical_expr_roundtrip_uses_a_growable_stack"; + + fn expr_roundtrip_on_two_mib_stack(depth: usize) { + std::thread::Builder::new() + .name("two-megabyte-stack".into()) + .stack_size(2 * 1024 * 1024) + .spawn(move || { + let mut expr: Arc = lit(true); + for _ in 0..depth { + expr = Arc::new(NotExpr::new(expr)); + } + let codec = DefaultPhysicalExtensionCodec {}; + let proto = serialize_physical_expr(&expr, &codec).unwrap(); + let ctx = SessionContext::new(); + parse_physical_expr(&proto, &ctx.task_ctx(), &Schema::empty(), &codec) + .unwrap(); + }) + .unwrap() + .join() + .unwrap(); + } + + #[test] + fn deeply_nested_physical_expr_roundtrip_uses_a_growable_stack() { + run_in_child_with(EXPR_TEST_NAME, 2_000, expr_roundtrip_on_two_mib_stack); + } +}