Refactor protobuf ser/de traits to allow decorator pattern

[adriangb]

Refactor PhysicalExtensionCodec with &dyn dispatch (Alternative to #19234)

Which issue does this PR close?

Closes#18477

Related to #19234 - this PR provides an alternative implementation approach.

Rationale for this change

DataFusion's protobuf serialization system currently has limited interception points. The existing PhysicalExtensionCodec
trait only provides hooks for custom extensions (try_decode/try_encode) and unknown expressions
(try_decode_expr/try_encode_expr), but users cannot intercept serialization of all plan and expression nodes.

This limitation prevents important use cases:

• Caching: Reusing previously deserialized expressions to avoid redundant parsing
• Transformation: Modifying nodes during serialization/deserialization
• Metadata injection: Preserving custom state not captured in the protobuf schema
• Decorator patterns: Wrapping standard serialization with custom pre/post-processing

PR #19234 addresses this by switching from &dyn PhysicalExtensionCodec to generics (&C where C: PhysicalExtensionCodec +
?Sized). However, #19234 (comment) raised concerns about FFI
compatibility, since distributed systems and FFI boundaries often require dynamic dispatch.

What changes are included in this PR?

This PR takes an alternative approach: keep &dyn dispatch but add 4 new required methods that intercept every
plan/expression node during serialization:

pubtraitPhysicalExtensionCodec:Debug + Send + Sync{// ... existing methods unchanged ...fndeserialize_physical_plan(&self,proto:&PhysicalPlanNode,ctx:&TaskContext)
-> Result<Arc<dynExecutionPlan>>;fnserialize_physical_plan(&self,plan:Arc<dynExecutionPlan>)
-> Result<PhysicalPlanNode>;fndeserialize_physical_expr(&self,proto:&PhysicalExprNode,ctx:&TaskContext,input_schema:&Schema)
-> Result<Arc<dynPhysicalExpr>>;fnserialize_physical_expr(&self,expr:&Arc<dynPhysicalExpr>)
-> Result<PhysicalExprNode>;}

Key implementation details:

• All recursive calls in from_proto.rs and to_proto.rs now go through the codec methods
• Public default_* helper functions allow implementations to delegate to standard behavior
• No default implementations - users must explicitly implement all methods (with documented examples)

Comparison with PR #19234

Aspect	PR #19234	This PR
Dispatch	Generic &C where C: PhysicalExtensionCodec + ?Sized	Keep &dyn PhysicalExtensionCodec
Default implementations	Yes - methods have defaults calling helpers	No - all methods required
FFI compatibility	Requires concrete types at boundaries	Fully compatible with dynamic dispatch
Migration effort	Minimal (defaults handle it)	Explicit (must add 4 methods, but with simple delegation)
Monomorphization	Yes - separate code paths per codec type	No - single code path via vtable
Compile time impact	Potentially increased	None

Pros of this approach:

• Maintains FFI compatibility for distributed systems using dynamic codec resolution
• No monomorphization overhead - single code path for all codec implementations
• Explicit is better than implicit - users acknowledge the new API surface
• Preserves existing &dyn patterns used throughout the codebase

Cons of this approach:

• Breaking change requires updating all PhysicalExtensionCodec implementations
• Users must add 4 method implementations (though simple delegation to default_* helpers works)

Files changed

• datafusion/proto/src/physical_plan/mod.rs - Added 4 new trait methods + helper functions
• datafusion/proto/src/physical_plan/from_proto.rs - Recursive calls go through codec
• datafusion/proto/src/physical_plan/to_proto.rs - Recursive calls go through codec
• datafusion/ffi/src/proto/physical_extension_codec.rs - Updated FFI codec implementations
• datafusion-examples/examples/proto/composed_extension_codec.rs - Updated example codecs
• datafusion/proto/tests/cases/roundtrip_physical_plan.rs - Updated test codecs
• docs/source/library-user-guide/upgrading.md - Added migration guide

Migration example

For existing implementations that don't need custom interception:

use datafusion_proto::physical_plan::{
default_deserialize_physical_plan, default_serialize_physical_plan,
default_deserialize_physical_expr, default_serialize_physical_expr,};implPhysicalExtensionCodecforMyCodec{// ... existing methods ...fndeserialize_physical_plan(&self,proto:&PhysicalPlanNode,ctx:&TaskContext)
-> Result<Arc<dynExecutionPlan>>{default_deserialize_physical_plan(proto, ctx,self)}fnserialize_physical_plan(&self,plan:Arc<dynExecutionPlan>)
-> Result<PhysicalPlanNode>{default_serialize_physical_plan(plan,self)}fndeserialize_physical_expr(&self,proto:&PhysicalExprNode,ctx:&TaskContext,input_schema:&Schema)
-> Result<Arc<dynPhysicalExpr>>{default_deserialize_physical_expr(proto, ctx, input_schema,self)}fnserialize_physical_expr(&self,expr:&Arc<dynPhysicalExpr>)
-> Result<PhysicalExprNode>{default_serialize_physical_expr(expr,self)}}

[milenkovicm]

I find the requirement to implement those four methods a bit confusing and the increasing API complexity unnecessary, ("we have defaults for you but we force you to implement it yourself"). No perfect solutions apparently

anyway, lets give @timsaucer chance to have a look

[adriangb]

I find the requirement to implement those four methods a bit confusing and the increasing API complexity unnecessary, ("we have defaults for you but we force you to implement it yourself"). No perfect solutions apparently

Yes agreed it's unfortunate but I don't see any way around it. It's just a limitation of Rust.

[milenkovicm]

@adriangb do you need before and after interception points for encode and decode or you could work with before_encode_* and after_decode_*? Not sure if that would do what you intent to do and remove need for recursive calls in ExtensionCodec

[adriangb]

I think it’s useful to be able to control if the default serializer is called or not. For example in the case of caching Arc’ed stuff you want to check the cache, delegate to the default, and then cache after.

I also don’t think that would help the recursiveness: you’d still have to in the default implementation pass a reference to &self where self is unsized to &dyn Trait.

Would it help if we provided a macro to implement the defaults?

[milenkovicm]

I agree with you it's hard problem to crack.
I'm not sure about macros, maybe. If we need to explain how to use the API, maybe we overcomplicate it.

Would [before|after]_[encode|decode]_[plan|expr] where before can return either proceed or stop result which would then trigger default method recurssion? They could have a default implementation, plus users who do not need such functionality will never need to implement it. We would have like eight new methods, which increase API complexity but not necessarily by too much as they will have a default implementation.

or if [encode|decode]_[plan|expr] return some kind of "aspect" instance which would go to recursion rather than codec itself would help

I'm just throwing ideas, trying to help, will have a look tomorrow

[adriangb]

Maybe something similar to the TreeNode API would work?

[milenkovicm]

just rough idea,

this goes to encoder:

fnbefore_encode_expr(&self,_node:&Arc<dynPhysicalExpr>,) -> Result<Option<protobuf::PhysicalExprNode>>{Ok(None)}

I don't like use of protobuf::PhysicalExprNode there as method is not similar to others in the traint

this part is to_proto:

pubfnserialize_physical_expr(value:&Arc<dynPhysicalExpr>,codec:&dynPhysicalExtensionCodec,) -> Result<protobuf::PhysicalExprNode>{match codec.before_encode_expr(value)? {Some(p) => Ok(p),None => serialize_physical_expr(value, codec),}}

maybe tree node like approach may make sense as well

[timsaucer]

This feels like it's more complicated than it needs to be. But I also must be missing something because I still don't see what it was that my original proposal in #18813 didn't meet the needs. I think you have some use cases and context in mind @adriangb that I clearly don't understand.

Regarding this PR, do you also need to cover the other functions in the codec? We now have two methods we need for physical plans, try_decode and deserialize_physical_plan. The first seems very codec-y to me. Given a bunch of bytes, produce a plan. It doesn't even necessarily need to use protobuf (up to the implementer). The second seems to be more about hooking into the process of using the codecs + non-codecs.

I think the problem we're running into with not having an elegant way of hooking these together is demonstrating there's a friction point between what the codec is for and what the protobuf processing part is for. I'm sure I'm biased towards my initial implementation approach, but I still feel that's a cleaner way of doing it.

[adriangb]

@timsaucer thank you for taking a look at this. I have my two key use cases as examples in this PR. When I first looked at your change it seemed like it would not be able to implement these examples, but I can give it another try to see if they can be made to work with your proposal.

[adriangb]

Closing in favor of #19437