This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23733-3a2bc0a0a5673386ab2238b545dbedb82ce6b92b in repository https://gitbox.apache.org/repos/asf/datafusion.git
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]
