This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 774f0bc34f feat(proto): thread expr encode/decode context into
try_encode_expr / try_decode_expr (#23733)
774f0bc34f is described below
commit 774f0bc34ff3dcd70fff537f828557584b8603c6
Author: Adrian Garcia Badaracco <[email protected]>
AuthorDate: Mon Jul 20 22:44:09 2026 -0500
feat(proto): thread expr encode/decode context into try_encode_expr /
try_decode_expr (#23733)
## Which issue does this PR close?
Closes #22920.
Alternative to #22922 — same goal, but plumbs the per-expr encode/decode
**context** (from the #22418 hook machinery) instead of the raw
`PhysicalProtoConverterExtension`.
## Rationale for this change
#21807 introduced the `DynamicFilterPhysicalExpr` dedup pipeline so
identical references on the wire reconstruct to one shared `Arc<Inner>`
via `expr_id` cache keys, and #22011 hooked it through the `SortExec` /
`AggregateExec` / `HashJoinExec` plan codecs.
The remaining gap is the **expression-level** extension path. In
`serialize_physical_expr_with_converter` /
`parse_physical_expr_with_converter`, the codec's `try_encode_expr` /
`try_decode_expr` is reached only as a fallback after the built-in
`try_to_proto` path and the `ScalarFunctionExpr` downcast — i.e. only
for downstream-defined **custom `PhysicalExpr`** types. When such a
codec serializes nested `PhysicalExprNode` fields *inside its own blob*,
the only helper available today is the free `serialize_physical_expr` /
`parse_physical_expr`, hardwired to `DefaultPhysicalProtoConverter`. So
those nested exprs get `expr_id: None` on the wire and reconstruct as
**distinct** `Inner` allocations — heap-max updates from a `SortExec`
never reach the wrapped reference.
## What changes are included in this PR?
The built-in expressions migrated under #22418 already receive a
`PhysicalExprEncodeCtx` / `PhysicalExprDecodeCtx` in their
`try_to_proto` / `try_from_proto` hooks. Those context objects bundle
the dedup-aware converter **plus** the active schema and task context,
and hide `PhysicalProtoConverterExtension` / `PhysicalExtensionCodec`
from the expression author entirely.
This PR hands the **same context** to the expr-level codec methods:
```rust
fn try_decode_expr(
&self,
buf: &[u8],
inputs: &[Arc<dyn PhysicalExpr>],
ctx: &PhysicalExprDecodeCtx<'_>, // new
) -> Result<Arc<dyn PhysicalExpr>>;
fn try_encode_expr(
&self,
node: &Arc<dyn PhysicalExpr>,
buf: &mut Vec<u8>,
ctx: &PhysicalExprEncodeCtx<'_>, // new
) -> Result<()>;
```
A codec that embeds nested `PhysicalExprNode`s now decodes them with
`ctx.decode(..)` and encodes them with `ctx.encode_child(..)`, which:
- route through any active `DeduplicatingProtoConverter` /
`DeduplicatingDeserializer`, so shared inner expressions cache-hit on
`expr_id`; **and**
- carry the real schema and task context, so nested UDF / column
references resolve against the actual registry.
### Why the context, not the raw converter (cf. #22922)
Threading the bare `PhysicalProtoConverterExtension` still leaves the
codec without the schema/registry, forcing it to fabricate a
`SessionContext::new()` and hard-code a schema to call
`proto_to_physical_expr` — an empty-registry footgun for any nested expr
that references a UDF or column. Passing the existing
`Physical{Encode,Decode}Ctx` avoids that, keeps the extension
escape-hatch consistent with the per-expr proto hooks, and adds no third
converter parameter to the codec API (the concern raised on #22922).
## Are these changes tested?
Yes — `extension_codec_expr_participates_in_deduplication` builds a
`BinaryExpr` whose left operand is a bare `DynamicFilterPhysicalExpr`
and whose right operand is a custom `WrapperExpr` whose codec embeds the
same dynamic filter inside its serialized blob (via `ctx.encode_child`).
After a `DeduplicatingProtoConverter` roundtrip, an `update()` on the
bare-side decoded filter is observed via `current()` on the wrapped-side
filter, proving both refs back the same `Inner`. The codec needs no
fabricated `SessionContext` — it decodes the nested expr through
`ctx.decode`.
## Are there any user-facing changes?
Yes — a **breaking change** for downstream codecs that override
`try_encode_expr` / `try_decode_expr`: they must add the new `ctx`
parameter (name it `_ctx` if the custom expr carries no nested
`PhysicalExprNode`s). Codecs that only override the plan-level
`try_encode` / `try_decode` are unaffected. Wire format is unchanged.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]>
---
datafusion/proto/src/physical_plan/from_proto.rs | 7 +-
datafusion/proto/src/physical_plan/mod.rs | 28 ++++
datafusion/proto/src/physical_plan/to_proto.rs | 2 +-
.../proto/tests/cases/roundtrip_physical_plan.rs | 186 +++++++++++++++++++++
4 files changed, 220 insertions(+), 3 deletions(-)
diff --git a/datafusion/proto/src/physical_plan/from_proto.rs
b/datafusion/proto/src/physical_plan/from_proto.rs
index d003d6db12..f1b324c79d 100644
--- a/datafusion/proto/src/physical_plan/from_proto.rs
+++ b/datafusion/proto/src/physical_plan/from_proto.rs
@@ -383,8 +383,11 @@ pub fn parse_physical_expr_with_converter(
.iter()
.map(|e| proto_converter.proto_to_physical_expr(e,
input_schema, ctx))
.collect::<Result<_>>()?;
- ctx.codec()
- .try_decode_expr(extension.expr.as_slice(), &inputs)? as _
+ ctx.codec().try_decode_expr(
+ extension.expr.as_slice(),
+ &inputs,
+ &decode_ctx,
+ )? as _
}
ExprType::Lambda(_) => LambdaExpr::try_from_proto(proto, &decode_ctx)?,
ExprType::LambdaVariable(_) => {
diff --git a/datafusion/proto/src/physical_plan/mod.rs
b/datafusion/proto/src/physical_plan/mod.rs
index 4f72668813..1ef8237595 100644
--- a/datafusion/proto/src/physical_plan/mod.rs
+++ b/datafusion/proto/src/physical_plan/mod.rs
@@ -62,6 +62,8 @@ use
datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctio
use datafusion_physical_expr::async_scalar_function::AsyncFuncExpr;
use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr;
use datafusion_physical_expr::{LexOrdering, LexRequirement, PhysicalExprRef};
+use
datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx;
+use
datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
use datafusion_physical_plan::aggregates::{
AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy,
};
@@ -4431,18 +4433,44 @@ pub trait PhysicalExtensionCodec: Debug + Send + Sync +
Any {
Ok(())
}
+ /// Decode a custom extension expression from `buf`.
+ ///
+ /// `inputs` holds the already-decoded children carried in the
+ /// `PhysicalExtensionExprNode.inputs` field. If the codec instead embeds
+ /// nested `PhysicalExprNode`s *inside* `buf`, decode them through
+ /// `ctx.decode(..)` (equivalently [`PhysicalExprDecodeCtx::decode`])
rather
+ /// than the free [`parse_physical_expr`] function: `ctx` carries the
active
+ /// schema and task context (so UDF/column references resolve against the
+ /// real registry) and routes through any active
`DeduplicatingDeserializer`,
+ /// so a shared inner expression (e.g. a `DynamicFilterPhysicalExpr`
+ /// referenced both from a `SortExec.filter` and from inside this blob)
+ /// cache-hits on its `expr_id` and re-shares one `Arc<dyn PhysicalExpr>`.
+ ///
+ /// [`parse_physical_expr`]:
crate::physical_plan::from_proto::parse_physical_expr
fn try_decode_expr(
&self,
_buf: &[u8],
_inputs: &[Arc<dyn PhysicalExpr>],
+ _ctx: &PhysicalExprDecodeCtx<'_>,
) -> Result<Arc<dyn PhysicalExpr>> {
not_impl_err!("PhysicalExtensionCodec is not provided")
}
+ /// Encode a custom extension expression into `buf`.
+ ///
+ /// If the codec embeds nested `PhysicalExprNode`s inside `buf`, encode
them
+ /// through `ctx.encode_child(..)` (equivalently
+ /// [`PhysicalExprEncodeCtx::encode_child`]) rather than the free
+ /// [`serialize_physical_expr`] function, so an active
+ /// `DeduplicatingProtoConverter` stamps matching `expr_id`s for shared
+ /// inner expressions. See [`Self::try_decode_expr`].
+ ///
+ /// [`serialize_physical_expr`]:
crate::physical_plan::to_proto::serialize_physical_expr
fn try_encode_expr(
&self,
_node: &Arc<dyn PhysicalExpr>,
_buf: &mut Vec<u8>,
+ _ctx: &PhysicalExprEncodeCtx<'_>,
) -> Result<()> {
not_impl_err!("PhysicalExtensionCodec is not provided")
}
diff --git a/datafusion/proto/src/physical_plan/to_proto.rs
b/datafusion/proto/src/physical_plan/to_proto.rs
index cade397c3b..515a53c087 100644
--- a/datafusion/proto/src/physical_plan/to_proto.rs
+++ b/datafusion/proto/src/physical_plan/to_proto.rs
@@ -332,7 +332,7 @@ pub fn serialize_physical_expr_with_converter(
})
} else {
let mut buf: Vec<u8> = vec![];
- match codec.try_encode_expr(value, &mut buf) {
+ match codec.try_encode_expr(value, &mut buf, &ctx) {
Ok(_) => {
let inputs: Vec<protobuf::PhysicalExprNode> = value
.children()
diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
index d87efcf986..28e70f2ddf 100644
--- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
+++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
@@ -120,6 +120,8 @@ use datafusion_functions_aggregate::min_max::max_udaf;
use datafusion_functions_aggregate::nth_value::nth_value_udaf;
use datafusion_functions_aggregate::string_agg::string_agg_udaf;
use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr;
+use
datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx;
+use
datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
use datafusion_proto::bytes::{
physical_plan_from_bytes_with_proto_converter,
physical_plan_to_bytes_with_proto_converter,
@@ -1369,6 +1371,7 @@ fn roundtrip_parquet_exec_with_custom_predicate_expr() ->
Result<()> {
&self,
buf: &[u8],
inputs: &[Arc<dyn PhysicalExpr>],
+ _ctx: &PhysicalExprDecodeCtx<'_>,
) -> Result<Arc<dyn PhysicalExpr>> {
if buf == "CustomPredicateExpr".as_bytes() {
Ok(Arc::new(CustomPredicateExpr {
@@ -1383,6 +1386,7 @@ fn roundtrip_parquet_exec_with_custom_predicate_expr() ->
Result<()> {
&self,
node: &Arc<dyn PhysicalExpr>,
buf: &mut Vec<u8>,
+ _ctx: &PhysicalExprEncodeCtx<'_>,
) -> Result<()> {
if node.downcast_ref::<CustomPredicateExpr>().is_some() {
buf.extend_from_slice("CustomPredicateExpr".as_bytes());
@@ -4542,3 +4546,185 @@ fn roundtrip_parquet_exec_range_output_partitioning()
-> Result<()> {
Ok(())
}
+
+/// A custom `PhysicalExpr` whose extension codec embeds a nested
+/// `PhysicalExprNode` *inside its own blob* (rather than the standard
+/// `PhysicalExtensionExprNode.inputs` field). This is the case that only
+/// works if the expr-level codec methods receive the encode/decode context.
+#[derive(Debug)]
+struct WrapperExpr {
+ inner: Arc<dyn PhysicalExpr>,
+}
+
+impl Display for WrapperExpr {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(f, "WrapperExpr({})", self.inner)
+ }
+}
+
+impl PartialEq for WrapperExpr {
+ fn eq(&self, other: &Self) -> bool {
+ self.inner.eq(&other.inner)
+ }
+}
+impl Eq for WrapperExpr {}
+
+impl std::hash::Hash for WrapperExpr {
+ fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+ self.inner.hash(state);
+ }
+}
+
+impl PhysicalExpr for WrapperExpr {
+ fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
+ self.inner.data_type(input_schema)
+ }
+ fn nullable(&self, input_schema: &Schema) -> Result<bool> {
+ self.inner.nullable(input_schema)
+ }
+ fn evaluate(&self, _batch: &RecordBatch) -> Result<ColumnarValue> {
+ internal_err!("WrapperExpr is not executable in this test")
+ }
+ fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+ vec![&self.inner]
+ }
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn PhysicalExpr>>,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ Ok(Arc::new(WrapperExpr {
+ inner: Arc::clone(&children[0]),
+ }))
+ }
+ fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ std::fmt::Display::fmt(self, f)
+ }
+}
+
+/// Wire layout for [`WrapperExpr`]: a single nested `PhysicalExprNode`.
+#[derive(Clone, PartialEq, prost::Message)]
+struct WrapperExprProto {
+ #[prost(message, optional, boxed, tag = "1")]
+ inner: Option<Box<datafusion_proto::protobuf::PhysicalExprNode>>,
+}
+
+#[derive(Debug)]
+struct WrapperCodec;
+
+impl PhysicalExtensionCodec for WrapperCodec {
+ fn try_decode(
+ &self,
+ _buf: &[u8],
+ _inputs: &[Arc<dyn ExecutionPlan>],
+ _ctx: &TaskContext,
+ _proto_converter: &dyn PhysicalProtoConverterExtension,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ internal_err!("not used")
+ }
+ fn try_encode(
+ &self,
+ _node: Arc<dyn ExecutionPlan>,
+ _buf: &mut Vec<u8>,
+ _proto_converter: &dyn PhysicalProtoConverterExtension,
+ ) -> Result<()> {
+ internal_err!("not used")
+ }
+ fn try_decode_expr(
+ &self,
+ buf: &[u8],
+ _inputs: &[Arc<dyn PhysicalExpr>],
+ ctx: &PhysicalExprDecodeCtx<'_>,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ let proto = WrapperExprProto::decode(buf)
+ .map_err(|e| internal_datafusion_err!("decode WrapperExprProto:
{e}"))?;
+ let inner_proto = proto
+ .inner
+ .ok_or_else(|| internal_datafusion_err!("missing inner"))?;
+ // Decode the nested expr through the context so it resolves against
+ // the real schema/registry AND participates in dedup — no fabricated
+ // `SessionContext` or hard-coded schema required.
+ let inner = ctx.decode(&inner_proto)?;
+ Ok(Arc::new(WrapperExpr { inner }))
+ }
+ fn try_encode_expr(
+ &self,
+ node: &Arc<dyn PhysicalExpr>,
+ buf: &mut Vec<u8>,
+ ctx: &PhysicalExprEncodeCtx<'_>,
+ ) -> Result<()> {
+ let wrapper = node
+ .downcast_ref::<WrapperExpr>()
+ .ok_or_else(|| internal_datafusion_err!("not WrapperExpr"))?;
+ // Encode the nested expr through the context so an active
+ // `DeduplicatingProtoConverter` stamps a matching `expr_id`.
+ let inner_proto = ctx.encode_child(&wrapper.inner)?;
+ let proto = WrapperExprProto {
+ inner: Some(Box::new(inner_proto)),
+ };
+ proto
+ .encode(buf)
+ .map_err(|e| internal_datafusion_err!("encode WrapperExprProto:
{e}"))?;
+ Ok(())
+ }
+}
+
+/// A `DynamicFilterPhysicalExpr` referenced both as a bare expression and
+/// nested inside a custom expression's codec blob must reconstruct to a
+/// single shared `Inner` after roundtrip.
+///
+/// This exercises the expr-level codec hooks receiving the encode/decode
+/// context: `try_encode_expr` routes its nested `PhysicalExprNode` through
+/// `ctx.encode_child` and `try_decode_expr` through `ctx.decode`, so the
+/// nested filter picks up the same `DeduplicatingProtoConverter` /
+/// `DeduplicatingDeserializer` cache as the bare reference. Without the
+/// context the nested expr would serialize with `expr_id: None` and decode
+/// into a distinct `Inner`, breaking heap-max propagation across the
+/// extension boundary in distributed execution.
+#[test]
+fn extension_codec_expr_participates_in_deduplication() -> Result<()> {
+ use prost::Message;
+
+ // A single composite expression holding TWO references to the same
+ // dynamic filter: bare on the left of an AND, wrapped on the right.
+ let dyn_filter = make_dynamic_filter();
+ let wrapper: Arc<dyn PhysicalExpr> = Arc::new(WrapperExpr {
+ inner: Arc::clone(&dyn_filter),
+ });
+ let composite: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new(
+ Arc::clone(&dyn_filter),
+ Operator::And,
+ Arc::clone(&wrapper),
+ ));
+
+ let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64,
false)]));
+ let codec = WrapperCodec;
+ let converter = DeduplicatingProtoConverter {};
+
+ // Encode, then round-trip through prost bytes to mimic the wire.
+ let proto = converter.physical_expr_to_proto(&composite, &codec)?;
+ let bytes = proto.encode_to_vec();
+ let decoded_proto =
+
datafusion_proto::protobuf::PhysicalExprNode::decode(bytes.as_slice()).unwrap();
+
+ let ctx = SessionContext::new();
+ let task_ctx = ctx.task_ctx();
+ let decode_ctx = PhysicalPlanDecodeContext::new(task_ctx.as_ref(), &codec);
+ let decoded =
+ converter.proto_to_physical_expr(&decoded_proto, &schema,
&decode_ctx)?;
+
+ let binary = decoded
+ .downcast_ref::<BinaryExpr>()
+ .expect("must decode back to BinaryExpr");
+ let decoded_left = Arc::clone(binary.left());
+ let decoded_right = Arc::clone(binary.right());
+ let decoded_wrapper = decoded_right
+ .downcast_ref::<WrapperExpr>()
+ .expect("right side must decode back to WrapperExpr");
+
+ // The load-bearing check: an `update()` on the bare-side filter must be
+ // observable from the wrapped-side filter, proving both refs back the
+ // same `Inner`.
+ assert_dynamic_filter_update_is_visible(&decoded_left,
&decoded_wrapper.inner)?;
+
+ Ok(())
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
