Repository navigation
Redesign protobuf encode/decode hooks & state #18477
Description
Workin on #18192 it became clear to me that there are both shortcomings in the current design and competing interests in terms of what is supported / what state is passed around. I thought it would be worth sitting down for a bit and trying to come up with a new design that satisfies these concerns.
My design goals were:
- More flexible pre/post processing. The current
Codecsystem only operates as a "fallback" and doesn't support hooking into the moment before trying to deserialize or the moment after a successful deserialization. I have a lot of example use cases for this that I can share. - Optionality of features. Some use cases demand more information / capabilities for serialization / deserialization (e.g. distributed systems want a function registry, a cache similar to Deduplicate PhysicalExpr on proto ser/de using Arc pointer addresses #18192, etc. while something that is just looking to cross the FFI boundary might be okay with no registry/state and just failing if it encounters a UDF).
The design I landed on is something like this:
// DataFusion code
trait Deserializer {
fn deserialize_physical_expr(&mut self, proto: PhysicalExprNode) -> Result<Arc<dyn PhysicalExpr>> {
// most methods have default implementations so you only have to implement those you want to override
deserialize_physical_expr(self, proto)
}
fn deserialize_udf(&mut self, proto: ScalarUdfExprNode) -> Result<()> {
// default implementation errors, we don't have a registry to call. you must override if you do.
internal_datafusion_err!("you must implement deserialize_udf yourself and put the function into a function registry")
}
}
/// Both the top level entry point (if what you have is a PhysicalExprNode)
/// and the default implementation that will get delegated back to recursively
fn deserialize_physical_expr<D: Deserializer + ?Sized>(deserializer: &mut D, proto: PhysicalExprNode) -> Result<Arc<dyn PhysicalExpr>> {
let children: Vec<PhysicalExprNode> = vec![]; // Placeholder for child nodes, from proto.children()?
let deserialized_children = children
.into_iter()
// it's important here that we call back into the trait!
.map(|child_proto| deserializer.deserialize_physical_expr(child_proto))
.collect::<Result<Vec<_>>>()?;
// Reconstruct the node with it's children
todo!()
}
fn deserialize_udf<D: Deserializer + ?Sized>(deserializer: &mut D, registry: &dyn FunctionRegistry, proto: ScalarUdfExprNode) -> Result<()> {
todo!();
}
// Users code: wants to override deserialization for PhysicalExpr
struct MySerializer {
field: String,
cache: HashMap<String, Arc<dyn PhysicalExpr>>,
}
impl Deserializer for MySerializer {
fn deserialize_physical_expr(&mut self, proto: PhysicalExprNode) -> Result<Arc<dyn PhysicalExpr>> {
// can run arbitrary stuff before, short circuit, etc.
println!("Deserializing using MySerializer {}", self.field);
let res = deserialize_physical_expr(&mut *self, proto)?;
// can run arbitrary stuff after and replace the result if necessary
println!("Deserialized result: {:?}", res);
Ok(res)
}
}
// Users code: wants to support deserialization of UDFs
struct MyDeserializerWithUDF<'a> {
registry: &'a dyn FunctionRegistry,
}
impl Deserializer for MyDeserializerWithUDF<'_> {
fn deserialize_physical_expr(&mut self, proto: PhysicalExprNode) -> Result<Arc<dyn PhysicalExpr>> {
deserialize_physical_expr(self, proto)
}
fn deserialize_udf(&mut self, proto: ScalarUdfExprNode) -> Result<()> {
deserialize_udf(self, self.registry, proto)
}
}Before continuing down this path I want go gauge appetite for this change, and understand if it satisfies the use cases currently out there or falls short.
@timsaucer @milenkovicm @Jefffrey any thoughts?
Activity
milenkovicm commentedon Nov 3, 2025
If I'm not mistaken extension codec will be invoked only for non datafusion types, would it make sense to add methods to it like you mentioned here with default implementations, which we already have. This way they can be overridden by users and no additional trait would be introduced
adriangb commentedon Nov 3, 2025
Right that's one of the limitations of the current codec system. I considered doing that but I think it would be hard to maintain both things and would get messy. I'd rather design the desired end state and then work backwards to see if there's a good backwards compatibility store. If there is, great. If there isn't then we just move to a better state and accept the churn.
milenkovicm commentedon Nov 3, 2025
Make Deserializer part of extension codec, and let the codec handle user defined types and datafusion core types? Would it make sense, and practical to implement?
milenkovicm commentedon Nov 3, 2025
It would help understanding if you could demonstrate how Deserializer would interact with current design
adriangb commentedon Nov 3, 2025
It doesn't. I am proposing we replace the current design entirely. Codec goes away, etc. Maybe we can make some backwards compatible version where the old code calls into the new code and we keep the old public functions around as deprecated for a bit, but my primary goal is to design a good end state with a transition plan or intermediate state being a secondary concern.
milenkovicm commentedon Nov 4, 2025
all right, that would be a big change of the API which was constant for quite some time
adriangb commentedon Nov 4, 2025
Didn't we just change some of the parameters to the functions in the past couple weeks? That's almost just as much of a breaking change from a users perspective.
milenkovicm commentedon Nov 4, 2025
I apologise, I don't get your idea fully, trying to understand.
If we go back to last week proposal, comment #18192 (comment) in particular
We went a bit tangental after that, unfortunately.
Would making
DecodeContexttrait provide similar flexibility likeDeserializer? Users could provide their specific implementations, and if they do they take risk of caching or whatever they do in it?timsaucer commentedon Nov 4, 2025
Would both of you be available tomorrow for the community meeting? This sounds like it could benefit from real time discussion.
adriangb commentedon Nov 4, 2025
Yep unless I have a second appendix that needs removing I plan on joining 😄
milenkovicm commentedon Nov 4, 2025
cant promise 100% but will try to join
alamb commentedon Nov 5, 2025
I think the idea of adding a trait makes a lot of sense to me
If we implemented this new trait for
TaskContextmost other users would not have to change their codetimsaucer commentedon Nov 19, 2025
In #18813 I took a first stab at what this might look like. This is impactful to the work I've been doing for the FFI crates and I think it might make my life vastly simpler if we have it or something like it implemented.
I'm not crazy about how a lot of the methods just mirror those in the
PhysicalExtensionCodecbut what this does do is give us a mutable wrapper around those so we can do tracking and other things. I don't think we can just update thePhysicalExtensionCodecwith&mut selfbecause it appears there is a use case where we pass aroundArc<dyn PhysicalExtensionCodec>and alsoArc<TaskContext>. In my experimentation I found that to implement mutable version ofPhysicalExtensionCodecI had to end up cloningTaskContext, which isn't ideal.milenkovicm commentedon Nov 19, 2025
I think your proposal makes sense @timsaucer .
One proposal, not related to your PR directly.
ExtensionCodecsare provided when there is a new physical/logical extension. When you need to deliver two artefacts for it, the actual extension andExtensionCodecswhich have to be merged into the current extension codec or added to some kind of composite extension codec.Would it make sense to add
ToSerializedandFromSerialzedtraits which extensions could implement directly, and make them usable without altering current extension codec?We would have to register
FromSerialzedsomewhere, like session context or similar so we could do reverse lookup for deserialisation. We would probably have specialised version ofTo/FromSerializedfor each type we currently cover by extension codec, but that needs investigation.We could follow similar approach to substrait extensions
where extensions would declare
extension_urlwhich would identify them at deserialiation point.feat: implement protobuf converter trait to allow control over serial…
feat: implement protobuf converter trait to allow control over serial…
feat: implement protobuf converter trait to allow control over serial…
feat: implement protobuf converter trait to allow control over serial…