Skip to content

Commit 96ff37a

Browse files
fix(proto): honor ExecutionPlan downcast_delegate during serialization (apache#23154)
## Which issue does this PR close? N/A. This is a follow-up to apache#22559 and apache#21263, but there is no dedicated issue for this specific bug. ## Rationale for this change apache#21263 removed `ExecutionPlan::as_any()` now that `ExecutionPlan` has `Any` as a supertrait. Most call sites were migrated from code like this: ```rust plan.as_any().downcast_ref::<SomeExec>() ``` to the new helper: ```rust plan.downcast_ref::<SomeExec>() ``` That distinction matters after apache#22559. `ExecutionPlan::downcast_ref` is now the public downcast operation for execution plans: it follows `ExecutionPlan::downcast_delegate()` before falling back to raw `Any`. This lets wrapper plans keep their own concrete type private while still presenting the wrapped plan's normal public identity. For example, tracing or instrumentation wrappers can implement `downcast_delegate()` so callers that ask "is this a `FilterExec` / `EmptyExec` / `ProjectionExec`?" get the same answer they would have received without the wrapper. The wrapper remains visible only to code that intentionally performs a raw concrete-type `Any` check. This problem came up while trying to instrument distributed plans through [`datafusion-distributed`](https://github.com/datafusion-contrib/datafusion-distributed). In that setting, physical plans may be wrapped for tracing/instrumentation and then serialized for distributed execution. The wrapper is meant to be transparent for normal plan inspection, but serialization still needs to recognize the delegated built-in plan underneath it. The physical plan proto serializer still had one leftover mechanical migration from the old `as_any()` world: ```rust let plan = plan.as_ref() as &dyn Any; ``` That line bypasses `ExecutionPlan::downcast_ref` entirely. As a result, a delegating wrapper around a built-in execution plan is not serialized as the built-in plan. Instead, the serializer sees only the wrapper's concrete type, fails all built-in `Exec` checks, and falls through to the extension codec. With the default extension codec this produces an unsupported-plan error, even though the wrapped plan itself is serializable. This is particularly relevant for transparent wrappers such as [`InstrumentedExec`](https://github.com/datafusion-contrib/datafusion-tracing/blob/main/datafusion-tracing/src/instrumented_exec.rs) in [`datafusion-contrib/datafusion-tracing`](https://github.com/datafusion-contrib/datafusion-tracing), which intentionally delegate public downcasts to the wrapped plan. ## What changes are included in this PR? - Changes physical plan proto serialization to bind the plan as `&dyn ExecutionPlan` instead of `&dyn Any`. - Leaves the existing serializer downcast chain intact, so each `plan.downcast_ref::<...>()` now dispatches through the `ExecutionPlan` helper and therefore honors `downcast_delegate()`. - Adds a regression test with a small wrapper execution plan that delegates public downcasts to an inner `EmptyExec`. - Audited other `as_ref() as &dyn Any` patterns. The remaining hits are either non-`ExecutionPlan` traits, `PhysicalExpr` checks, UDF tests, or the intentional raw `Any` fallback inside the `ExecutionPlan` helper itself. ## Are these changes tested? Yes. ```bash cargo fmt --all cargo fmt --all --check cargo test -p datafusion-proto --test proto_integration serialize_uses_downcast_delegate cargo clippy --all-targets --all-features -- -D warnings ``` ## Are there any user-facing changes? Yes, as a bug fix. Physical plan proto serialization now honors the documented `ExecutionPlan::downcast_delegate()` behavior. Transparent wrapper plans can serialize as their delegated built-in execution plan when appropriate instead of requiring an extension codec for the wrapper itself.
1 parent 994b926 commit 96ff37a

2 files changed

Lines changed: 71 additions & 3 deletions

File tree

datafusion/proto/src/physical_plan/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -442,7 +442,7 @@ pub trait PhysicalPlanNodeExt: Sized {
442442
proto_converter: &dyn PhysicalProtoConverterExtension,
443443
) -> Result<protobuf::PhysicalPlanNode> {
444444
let plan_clone = Arc::clone(&plan);
445-
let plan = plan.as_ref() as &dyn Any;
445+
let plan = plan.as_ref();
446446

447447
if let Some(exec) = plan.downcast_ref::<ExplainExec>() {
448448
return protobuf::PhysicalPlanNode::try_from_explain_exec(exec, codec);

datafusion/proto/tests/cases/roundtrip_physical_plan.rs

Lines changed: 70 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -89,8 +89,8 @@ use datafusion::physical_plan::windows::{
8989
};
9090
use datafusion::physical_plan::{
9191
DisplayAs, DisplayFormatType, ExecutionPlan, InputOrderMode, Partitioning,
92-
PhysicalExpr, RangePartitioning, SendableRecordBatchStream, SplitPoint, Statistics,
93-
displayable,
92+
PhysicalExpr, PlanProperties, RangePartitioning, SendableRecordBatchStream,
93+
SplitPoint, Statistics, displayable,
9494
};
9595
use datafusion::prelude::{ParquetReadOptions, SessionContext};
9696
use datafusion::scalar::ScalarValue;
@@ -229,6 +229,74 @@ fn roundtrip_empty() -> Result<()> {
229229
roundtrip_test(Arc::new(EmptyExec::new(Arc::new(Schema::empty()))))
230230
}
231231

232+
#[derive(Debug)]
233+
struct DowncastDelegatingExec {
234+
inner: Arc<dyn ExecutionPlan>,
235+
}
236+
237+
impl DowncastDelegatingExec {
238+
fn new(inner: Arc<dyn ExecutionPlan>) -> Self {
239+
Self { inner }
240+
}
241+
}
242+
243+
impl DisplayAs for DowncastDelegatingExec {
244+
fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
245+
self.inner.fmt_as(t, f)
246+
}
247+
}
248+
249+
impl ExecutionPlan for DowncastDelegatingExec {
250+
fn name(&self) -> &str {
251+
self.inner.name()
252+
}
253+
254+
fn properties(&self) -> &Arc<PlanProperties> {
255+
self.inner.properties()
256+
}
257+
258+
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
259+
self.inner.children()
260+
}
261+
262+
fn with_new_children(
263+
self: Arc<Self>,
264+
children: Vec<Arc<dyn ExecutionPlan>>,
265+
) -> Result<Arc<dyn ExecutionPlan>> {
266+
let inner = Arc::clone(&self.inner).with_new_children(children)?;
267+
Ok(Arc::new(Self::new(inner)))
268+
}
269+
270+
fn downcast_delegate(&self) -> Option<&dyn ExecutionPlan> {
271+
Some(self.inner.as_ref())
272+
}
273+
274+
fn execute(
275+
&self,
276+
partition: usize,
277+
context: Arc<TaskContext>,
278+
) -> Result<SendableRecordBatchStream> {
279+
self.inner.execute(partition, context)
280+
}
281+
}
282+
283+
#[test]
284+
fn serialize_uses_downcast_delegate() -> Result<()> {
285+
let inner: Arc<dyn ExecutionPlan> =
286+
Arc::new(EmptyExec::new(Arc::new(Schema::empty())));
287+
let plan: Arc<dyn ExecutionPlan> = Arc::new(DowncastDelegatingExec::new(inner));
288+
let codec = DefaultPhysicalExtensionCodec {};
289+
290+
let proto = PhysicalPlanNode::try_from_physical_plan(plan, &codec)?;
291+
292+
assert!(matches!(
293+
proto.physical_plan_type,
294+
Some(protobuf::physical_plan_node::PhysicalPlanType::Empty(_))
295+
));
296+
297+
Ok(())
298+
}
299+
232300
#[test]
233301
fn roundtrip_date_time_interval() -> Result<()> {
234302
let schema = Schema::new(vec![

0 commit comments

Comments
 (0)