adriangb commented on issue #24625:
URL: https://github.com/apache/datafusion/issues/24625#issuecomment-5684836786
Thanks Gabriel and Jay.
An update on what the API looks like now in the PRs:
```rust
impl ExecutionPlan for MyExec {
fn try_to_proto(&self, ctx: &ExecutionPlanEncodeCtx<'_>) ->
Result<Option<PhysicalPlanNode>> {
Ok(Some(PhysicalPlanNode {
physical_plan_type:
Some(PhysicalPlanType::Extension(PhysicalExtensionNode {
// Made up function that knows how to go from MyExec ->
bytes.
// Any encoding, as today; it doesn't have to be protobuf.
node: encode_my_exec(self)?,
inputs: ctx.encode_children(self.children())?,
plan_name: Some(Self::NAME.to_string()),
})),
}))
}
}
impl ExecutionPlanFromProto for MyExec {
const NAME: &'static str = "my-crate.MyExec";
fn try_from_proto(node: &PhysicalPlanNode, ctx:
&ExecutionPlanDecodeCtx<'_>) -> Result<Arc<dyn ExecutionPlan>> {
// The registry hands over the whole node, so the plan matches its
own variant.
let extension = expect_plan_variant!(node,
PhysicalPlanType::Extension, "Extension");
let children = ctx.decode_children(&extension.inputs)?;
let [input] = children.try_into().map_err(|_|
plan_datafusion_err!("{} expects one input", Self::NAME))?;
// Made up function that knows how to go from bytes + children +
session -> MyExec.
// `ctx.task_ctx()` gives the decoding session, for a plan that
rebuilds
// session-scoped state (a worker pool, a channel resolver) at
decode time.
let exec = decode_my_exec(extension.node.as_slice(), input,
ctx.task_ctx())?;
Ok(Arc::new(exec))
}
}
let mut registry = ExecutionPlanRegistry::new();
registry.register::<MyExec>()?;
let config = SessionConfig::new().with_extension(Arc::new(registry));
```
Regarding your three points:
- Protobuf coupling. The payload stays `Vec<u8>`. Only the envelope is
protobuf, and it already is today: a codec's bytes ride inside
`PhysicalExtensionNode.node` either way. The only addition to the envelope is
an optional string. A non-protobuf payload works on both paths.
- Two APIs at once. Agreed, this is worth thinking about. The PR does not
deprecate `PhysicalExtensionCodec`; it is kept around for three reasons, all of
which can be resolved in time:
- Any plans or expressions not migrated yet (gradual rollout / backwards
compatibility).
- UDF/UDAF/UDWF payloads (a function is an instance, not a type). UDFs are
already resolved by name through the session's `FunctionRegistry` before the
codec is consulted, so the codec only matters for UDFs that carry instance
state. The end state is a self-serialization hook on the UDF trait plus a
per-name decoder, keyed through the `FunctionRegistry` we already have. Not
designed yet.
- `FFI_PhysicalExtensionCodec`, which datafusion-python exposes as public
API. `FFI_ExecutionPlan` has no serialization slot today, so foreign plans
always go through the codec. The registry stores decoders privately, so the
storage can later become `Arc<dyn Fn>` with a `register_decoder(name, ..)` for
foreign decoders, and `PhysicalPlanNode` already crosses the ABI as prost
bytes. No breaking change needed to get there.
What the registry makes redundant is `ComposedPhysicalExtensionCodec`, whose
reason to exist is composition. I would rather deprecate that once both plans
and expressions can go through a registry, as a separate decision. Might have
to wait for UDFs to be resolved as well.
A type-keyed successor to `ComposedPhysicalExtensionCodec` is close to what
this PR is, with one difference: the discriminator is on the wire
(`plan_name`), not inside a codec-private wrapper. That is what lets a reader
distinguish "no decoder is registered for X" from "the decoder for X failed",
which resolves the confusion described in
https://github.com/apache/datafusion/issues/24331. It also means no codec type
has to exist per plan at all. We could still do a name-keyed composed codec if
there is an extended migration period, but I hope that is not necessary.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
