Skip to content

proto: add proto converter reference to PhysicalExtensionCodec trait #21055

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jayshrivastava wants to merge 5 commits into
base: main
Choose a base branch
from
Open

proto: add proto converter reference to PhysicalExtensionCodec trait #21055

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
dc8addd
add proto converter reference to PhysicalExtensionCodec to allow dyna…
jayshrivastava Feb 27, 2026
6dc8faf
wip
jayshrivastava Apr 22, 2026
e0baa89
wip;
jayshrivastava May 6, 2026
c30ee94
more linting
jayshrivastava May 14, 2026
c53347d
typo
jayshrivastava May 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions datafusion-examples/examples/custom_data_source/adapter_serialization.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -275,6 +275,7 @@ impl PhysicalExtensionCodec for AdapterPreservingCodec {
buf: &[u8],
inputs: &[Arc<dyn ExecutionPlan>],
_ctx: &TaskContext,
_proto_converter: &dyn PhysicalProtoConverterExtension,
Copy link
Copy Markdown
Contributor

@adriangb adriangb May 14, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could add new methods with the param that default to the old methods to avoid a breaking change. Did you consider that as an option?

) -> Result<Arc<dyn ExecutionPlan>> {
// Try to parse as our extension payload
if let Ok(payload) = serde_json::from_slice::<ExtensionPayload>(buf)
Expand Down Expand Up @@ -303,6 +304,7 @@ impl PhysicalExtensionCodec for AdapterPreservingCodec {
&self,
_node: Arc<dyn ExecutionPlan>,
_buf: &mut Vec<u8>,
_proto_converter: &dyn PhysicalProtoConverterExtension,
Copy link
Copy Markdown
Contributor

@adriangb adriangb May 14, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Naming these converter would match try_from_physical_plan_with_converter and make the arg name shorter.

) -> Result<()> {
// We don't need this for the example - adapter wrapping happens in
// `execution_plan_to_proto` instead.
Expand Down
17 changes: 15 additions & 2 deletions datafusion-examples/examples/proto/composed_extension_codec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ use datafusion::physical_plan::{DisplayAs, ExecutionPlan};
use datafusion::prelude::SessionContext;
use datafusion_proto::physical_plan::{
AsExecutionPlan, ComposedPhysicalExtensionCodec, PhysicalExtensionCodec,
PhysicalProtoConverterExtension,
};
use datafusion_proto::protobuf;

Expand Down Expand Up @@ -145,6 +146,7 @@ impl PhysicalExtensionCodec for ParentPhysicalExtensionCodec {
buf: &[u8],
inputs: &[Arc<dyn ExecutionPlan>],
_ctx: &TaskContext,
_proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<Arc<dyn ExecutionPlan>> {
if buf == "ParentExec".as_bytes() {
Ok(Arc::new(ParentExec {
Expand All @@ -155,7 +157,12 @@ impl PhysicalExtensionCodec for ParentPhysicalExtensionCodec {
}
}

fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
fn try_encode(
&self,
node: Arc<dyn ExecutionPlan>,
buf: &mut Vec<u8>,
_proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<()> {
if node.is::<ParentExec>() {
buf.extend_from_slice("ParentExec".as_bytes());
Ok(())
Expand Down Expand Up @@ -226,6 +233,7 @@ impl PhysicalExtensionCodec for ChildPhysicalExtensionCodec {
buf: &[u8],
_inputs: &[Arc<dyn ExecutionPlan>],
_ctx: &TaskContext,
_proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<Arc<dyn ExecutionPlan>> {
if buf == "ChildExec".as_bytes() {
Ok(Arc::new(ChildExec {}))
Expand All @@ -234,7 +242,12 @@ impl PhysicalExtensionCodec for ChildPhysicalExtensionCodec {
}
}

fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
fn try_encode(
&self,
node: Arc<dyn ExecutionPlan>,
buf: &mut Vec<u8>,
_proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<()> {
if node.is::<ChildExec>() {
buf.extend_from_slice("ChildExec".as_bytes());
Ok(())
Expand Down
2 changes: 2 additions & 0 deletions datafusion-examples/examples/proto/expression_deduplication.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -187,6 +187,7 @@ impl PhysicalExtensionCodec for CachingCodec {
_buf: &[u8],
_inputs: &[Arc<dyn ExecutionPlan>],
_ctx: &TaskContext,
_proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<Arc<dyn ExecutionPlan>> {
datafusion::common::not_impl_err!("No custom extension nodes")
}
Expand All @@ -196,6 +197,7 @@ impl PhysicalExtensionCodec for CachingCodec {
&self,
_node: Arc<dyn ExecutionPlan>,
_buf: &mut Vec<u8>,
_proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<()> {
datafusion::common::not_impl_err!("No custom extension nodes")
}
Expand Down
50 changes: 40 additions & 10 deletions datafusion/ffi/src/proto/physical_extension_codec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,10 @@ use datafusion_expr::{
AggregateUDF, AggregateUDFImpl, ScalarUDF, ScalarUDFImpl, WindowUDF, WindowUDFImpl,
};
use datafusion_physical_plan::ExecutionPlan;
use datafusion_proto::physical_plan::PhysicalExtensionCodec;
use datafusion_proto::physical_plan::{
DefaultPhysicalProtoConverter, PhysicalExtensionCodec,
PhysicalProtoConverterExtension,
};

use stabby::slice::Slice as SSlice;
use stabby::str::Str as SStr;
Expand Down Expand Up @@ -145,8 +148,12 @@ unsafe extern "C" fn try_decode_fn_wrapper(
.collect::<Result<Vec<_>>>();
let inputs = sresult_return!(inputs);

let plan =
sresult_return!(codec.try_decode(buf.as_ref(), &inputs, task_ctx.as_ref()));
let plan = sresult_return!(codec.try_decode(
buf.as_ref(),
&inputs,
task_ctx.as_ref(),
&DefaultPhysicalProtoConverter {},
Copy link
Copy Markdown
Contributor

@adriangb adriangb May 14, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@timsaucer not sure if you have any suggestions but I think this is fine. Generally we have not made any effort to make dynamic filters / expr deduplication work with plans that go through FFI.

));

FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime))
}
Expand All @@ -160,7 +167,11 @@ unsafe extern "C" fn try_encode_fn_wrapper(
let plan: Arc<dyn ExecutionPlan> = sresult_return!((&node).try_into());

let mut bytes = Vec::new();
sresult_return!(codec.try_encode(plan, &mut bytes));
sresult_return!(codec.try_encode(
plan,
&mut bytes,
&DefaultPhysicalProtoConverter {}
));

FFI_Result::Ok(bytes.into_iter().collect())
}
Expand Down Expand Up @@ -335,6 +346,7 @@ impl PhysicalExtensionCodec for ForeignPhysicalExtensionCodec {
buf: &[u8],
inputs: &[Arc<dyn ExecutionPlan>],
_ctx: &TaskContext,
_proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<Arc<dyn ExecutionPlan>> {
let inputs = inputs
.iter()
Expand All @@ -348,7 +360,12 @@ impl PhysicalExtensionCodec for ForeignPhysicalExtensionCodec {
Ok(plan)
}

fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
fn try_encode(
&self,
node: Arc<dyn ExecutionPlan>,
buf: &mut Vec<u8>,
_proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<()> {
let plan = FFI_ExecutionPlan::new(node, None);
let bytes = df_result!(unsafe { (self.0.try_encode)(&self.0, plan) })?;

Expand Down Expand Up @@ -426,7 +443,10 @@ pub(crate) mod tests {
use datafusion_functions_aggregate::sum::Sum;
use datafusion_functions_window::rank::{Rank, RankType};
use datafusion_physical_plan::ExecutionPlan;
use datafusion_proto::physical_plan::PhysicalExtensionCodec;
use datafusion_proto::physical_plan::{
DefaultPhysicalProtoConverter, PhysicalExtensionCodec,
PhysicalProtoConverterExtension,
};

use crate::execution_plan::tests::EmptyExec;
use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec;
Expand All @@ -449,6 +469,7 @@ pub(crate) mod tests {
buf: &[u8],
_inputs: &[Arc<dyn ExecutionPlan>],
_ctx: &TaskContext,
_proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<Arc<dyn ExecutionPlan>> {
if buf[0] != Self::MAGIC_NUMBER {
return exec_err!(
Expand All @@ -467,6 +488,7 @@ pub(crate) mod tests {
&self,
node: Arc<dyn ExecutionPlan>,
buf: &mut Vec<u8>,
_proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<()> {
buf.push(Self::MAGIC_NUMBER);

Expand Down Expand Up @@ -587,10 +609,18 @@ pub(crate) mod tests {
let exec = create_test_exec();
let input_execs = [create_test_exec()];
let mut bytes = Vec::new();
foreign_codec.try_encode(Arc::clone(&exec), &mut bytes)?;

let returned_exec =
foreign_codec.try_decode(&bytes, &input_execs, ctx.task_ctx().as_ref())?;
foreign_codec.try_encode(
Arc::clone(&exec),
&mut bytes,
&DefaultPhysicalProtoConverter {},
)?;

let returned_exec = foreign_codec.try_decode(
&bytes,
&input_execs,
ctx.task_ctx().as_ref(),
&DefaultPhysicalProtoConverter {},
)?;

assert!(returned_exec.is::<EmptyExec>());

Expand Down
37 changes: 29 additions & 8 deletions datafusion/proto/src/physical_plan/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -616,7 +616,7 @@ impl protobuf::PhysicalPlanNode {
}

let mut buf: Vec<u8> = vec![];
match codec.try_encode(Arc::clone(&plan_clone), &mut buf) {
match codec.try_encode(Arc::clone(&plan_clone), &mut buf, proto_converter) {
Ok(_) => {
let inputs: Vec<protobuf::PhysicalPlanNode> = plan_clone
.children()
Expand Down Expand Up @@ -1722,9 +1722,12 @@ impl protobuf::PhysicalPlanNode {
.map(|i| proto_converter.proto_to_execution_plan(i, ctx))
.collect::<Result<_>>()?;

let extension_node =
ctx.codec()
.try_decode(extension.node.as_slice(), &inputs, ctx.task_ctx())?;
let extension_node = ctx.codec().try_decode(
extension.node.as_slice(),
&inputs,
ctx.task_ctx(),
proto_converter,
)?;

Ok(extension_node)
}
Expand Down Expand Up @@ -3704,9 +3707,15 @@ pub trait PhysicalExtensionCodec: Debug + Send + Sync + Any {
buf: &[u8],
inputs: &[Arc<dyn ExecutionPlan>],
ctx: &TaskContext,
proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<Arc<dyn ExecutionPlan>>;

fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()>;
fn try_encode(
&self,
node: Arc<dyn ExecutionPlan>,
buf: &mut Vec<u8>,
proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<()>;

fn try_decode_udf(&self, name: &str, _buf: &[u8]) -> Result<Arc<ScalarUDF>> {
not_impl_err!("PhysicalExtensionCodec is not provided for scalar function {name}")
Expand Down Expand Up @@ -3760,6 +3769,7 @@ impl PhysicalExtensionCodec for DefaultPhysicalExtensionCodec {
_buf: &[u8],
_inputs: &[Arc<dyn ExecutionPlan>],
_ctx: &TaskContext,
_proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<Arc<dyn ExecutionPlan>> {
not_impl_err!("PhysicalExtensionCodec is not provided")
}
Expand All @@ -3768,6 +3778,7 @@ impl PhysicalExtensionCodec for DefaultPhysicalExtensionCodec {
&self,
_node: Arc<dyn ExecutionPlan>,
_buf: &mut Vec<u8>,
_proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<()> {
not_impl_err!("PhysicalExtensionCodec is not provided")
}
Expand Down Expand Up @@ -4091,12 +4102,22 @@ impl PhysicalExtensionCodec for ComposedPhysicalExtensionCodec {
buf: &[u8],
inputs: &[Arc<dyn ExecutionPlan>],
ctx: &TaskContext,
proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<Arc<dyn ExecutionPlan>> {
self.decode_protobuf(buf, |codec, data| codec.try_decode(data, inputs, ctx))
self.decode_protobuf(buf, |codec, data| {
codec.try_decode(data, inputs, ctx, proto_converter)
})
}

fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
self.encode_protobuf(buf, |codec, data| codec.try_encode(Arc::clone(&node), data))
fn try_encode(
&self,
node: Arc<dyn ExecutionPlan>,
buf: &mut Vec<u8>,
proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<()> {
self.encode_protobuf(buf, |codec, data| {
codec.try_encode(Arc::clone(&node), data, proto_converter)
})
}

fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> {
Expand Down
Loading
Loading