From f5dbee1254e314db61f14721556190061b2a9763 Mon Sep 17 00:00:00 2001 From: Sainath Singineedi <44405294+sainad2222@users.noreply.github.com> Date: Sun, 6 Sep 2026 10:36:30 +0000 Subject: [PATCH 1/2] fix: delegate the remaining codec methods in ComposedPhysicalExtensionCodec --- datafusion/proto/src/physical_plan/mod.rs | 213 +++++++++++++++++++++- 1 file changed, 211 insertions(+), 2 deletions(-) diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 26d1a8ed83d36..350715a91fe64 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -724,6 +724,40 @@ mod tests { } } + /// Codec that rejects function serde outright, as a codec which only + /// knows its own plan nodes does. + #[derive(Debug)] + struct RejectsFunctionsCodec; + + impl PhysicalExtensionCodec for RejectsFunctionsCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + internal_err!("not needed for these tests") + } + + fn try_encode( + &self, + _node: Arc, + _buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + internal_err!("not needed for these tests") + } + + fn try_encode_udwf( + &self, + _node: &WindowUDF, + _buf: &mut Vec, + ) -> Result<()> { + internal_err!("not mine") + } + } + /// Codec whose decode hooks only accept an empty payload, to pin the /// by-name decode fallback (registry miss → codec with `&[]`). #[derive(Debug)] @@ -954,6 +988,79 @@ mod tests { "unexpected error: {err}" ); } + + #[test] + fn composed_codec_round_trips_a_child_function_payload() -> Result<()> { + let composed = + ComposedPhysicalExtensionCodec::new(vec![Arc::new(PayloadCodec)]); + + let mut buf = vec![]; + composed.try_encode_udwf(&WindowUDF::from(TestUdwf::new()), &mut buf)?; + assert!(!buf.is_empty(), "a codec with a payload must produce one"); + + let decoded = composed.try_decode_udwf("test_udwf", &buf)?; + assert_eq!(decoded.name(), "test_udwf"); + Ok(()) + } + + #[test] + fn composed_codec_leaves_by_name_functions_payload_free() -> Result<()> { + // A plans-only codec accepts these hooks without writing anything. + // Wrapping that in a `DataEncoderTuple` would make a registry + // function look codec-encoded and strand it at decode time. + let composed = ComposedPhysicalExtensionCodec::new(vec![Arc::new( + DefaultPhysicalExtensionCodec {}, + )]); + + let mut buf = vec![]; + composed.try_encode_udf(&ScalarUDF::from(TestUdf::new()), &mut buf)?; + assert!(buf.is_empty(), "udf: expected no payload, got {buf:?}"); + + let mut buf = vec![]; + composed.try_encode_udaf(&AggregateUDF::from(TestUdaf::new()), &mut buf)?; + assert!(buf.is_empty(), "udaf: expected no payload, got {buf:?}"); + + let mut buf = vec![]; + composed.try_encode_udwf(&WindowUDF::from(TestUdwf::new()), &mut buf)?; + assert!(buf.is_empty(), "udwf: expected no payload, got {buf:?}"); + + Ok(()) + } + + #[test] + fn composed_codec_keeps_by_name_encoding_after_a_rejecting_codec() -> Result<()> { + // The accepting codec sits at position 1, so a `DataEncoderTuple` + // wrapping its empty blob would not encode to zero bytes. Decode + // skips the registry whenever a payload is present, so emitting one + // here would strand every registry-resolvable window function. + let composed = ComposedPhysicalExtensionCodec::new(vec![ + Arc::new(RejectsFunctionsCodec), + Arc::new(DefaultPhysicalExtensionCodec {}), + ]); + + let mut buf = vec![]; + composed.try_encode_udwf(&WindowUDF::from(TestUdwf::new()), &mut buf)?; + assert!(buf.is_empty(), "expected no payload, got {buf:?}"); + Ok(()) + } + + #[test] + fn composed_codec_looks_past_payload_free_codecs() -> Result<()> { + // The payload-free codec is first and must not shadow the codec + // that actually encodes the function. + let composed = ComposedPhysicalExtensionCodec::new(vec![ + Arc::new(DefaultPhysicalExtensionCodec {}), + Arc::new(PayloadCodec), + ]); + + let mut buf = vec![]; + composed.try_encode_udwf(&WindowUDF::from(TestUdwf::new()), &mut buf)?; + assert!(!buf.is_empty(), "expected the second codec's payload"); + + let decoded = composed.try_decode_udwf("test_udwf", &buf)?; + assert_eq!(decoded.name(), "test_udwf"); + Ok(()) + } } } @@ -2015,6 +2122,56 @@ impl ComposedPhysicalExtensionCodec { .encode(buf) .map_err(|e| internal_datafusion_err!("{e}")) } + + /// Like [`Self::encode_protobuf`], but for the function hooks whose trait + /// default is `Ok(())` rather than an error. + /// + /// Those hooks treat an empty buffer as "no custom payload, encode by + /// name", and the decode side only consults the function registry when the + /// payload is absent. Wrapping an empty blob in a [`DataEncoderTuple`] + /// would make a by-name function look codec-encoded and strand it at + /// decode time, so a codec that writes nothing must leave `buf` untouched. + fn encode_protobuf_by_name_aware( + &self, + buf: &mut Vec, + mut encode: impl FnMut(&dyn PhysicalExtensionCodec, &mut Vec) -> Result<()>, + ) -> Result<()> { + let mut data = vec![]; + let mut last_err = None; + let mut encoder_position = None; + let mut any_accepted = false; + + for (position, codec) in self.codecs.iter().enumerate() { + data.clear(); + match encode(codec.as_ref(), &mut data) { + Ok(()) => { + any_accepted = true; + if !data.is_empty() { + encoder_position = Some(position as u32); + break; + } + } + Err(err) => last_err = Some(err), + } + } + + let Some(encoder_position) = encoder_position else { + return match last_err { + // Every codec rejected the function outright. + Some(err) if !any_accepted => Err(err), + // A codec accepted it without a payload: leave it encoded by name. + _ => Ok(()), + }; + }; + + let proto = DataEncoderTuple { + encoder_position, + blob: data, + }; + proto + .encode(buf) + .map_err(|e| internal_datafusion_err!("{e}")) + } } impl PhysicalExtensionCodec for ComposedPhysicalExtensionCodec { @@ -2046,7 +2203,47 @@ impl PhysicalExtensionCodec for ComposedPhysicalExtensionCodec { } fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { - self.encode_protobuf(buf, |codec, data| codec.try_encode_udf(node, data)) + self.encode_protobuf_by_name_aware(buf, |codec, data| { + codec.try_encode_udf(node, data) + }) + } + + fn try_decode_higher_order_function( + &self, + name: &str, + buf: &[u8], + ) -> Result> { + self.decode_protobuf(buf, |codec, data| { + codec.try_decode_higher_order_function(name, data) + }) + } + + fn try_encode_higher_order_function( + &self, + node: &HigherOrderUDF, + buf: &mut Vec, + ) -> Result<()> { + self.encode_protobuf_by_name_aware(buf, |codec, data| { + codec.try_encode_higher_order_function(node, data) + }) + } + + fn try_decode_expr( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &PhysicalExprDecodeCtx<'_>, + ) -> Result> { + self.decode_protobuf(buf, |codec, data| codec.try_decode_expr(data, inputs, ctx)) + } + + fn try_encode_expr( + &self, + node: &Arc, + buf: &mut Vec, + ctx: &PhysicalExprEncodeCtx<'_>, + ) -> Result<()> { + self.encode_protobuf(buf, |codec, data| codec.try_encode_expr(node, data, ctx)) } fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { @@ -2054,7 +2251,19 @@ impl PhysicalExtensionCodec for ComposedPhysicalExtensionCodec { } fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { - self.encode_protobuf(buf, |codec, data| codec.try_encode_udaf(node, data)) + self.encode_protobuf_by_name_aware(buf, |codec, data| { + codec.try_encode_udaf(node, data) + }) + } + + fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { + self.decode_protobuf(buf, |codec, data| codec.try_decode_udwf(name, data)) + } + + fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { + self.encode_protobuf_by_name_aware(buf, |codec, data| { + codec.try_encode_udwf(node, data) + }) } } From e1e35488a3696c72793b276a042b1d33de979313 Mon Sep 17 00:00:00 2001 From: Sainath Singineedi <44405294+sainad2222@users.noreply.github.com> Date: Sun, 6 Sep 2026 16:48:48 +0000 Subject: [PATCH 2/2] fix: clear the encode buffer between composed codec attempts --- datafusion/proto/src/physical_plan/mod.rs | 79 +++++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 350715a91fe64..31425e17b0753 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -758,6 +758,58 @@ mod tests { } } + /// Writes into the buffer before discovering it cannot encode the node. + #[derive(Debug)] + struct WritesThenRejectsCodec; + + impl PhysicalExtensionCodec for WritesThenRejectsCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + internal_err!("not needed for these tests") + } + + fn try_encode( + &self, + _node: Arc, + buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + buf.extend_from_slice(b"partial"); + internal_err!("not mine") + } + } + + /// Encodes any plan as a fixed payload. + #[derive(Debug)] + struct PlanPayloadCodec; + + impl PhysicalExtensionCodec for PlanPayloadCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + internal_err!("not needed for these tests") + } + + fn try_encode( + &self, + _node: Arc, + buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + buf.extend_from_slice(b"plan"); + Ok(()) + } + } + /// Codec whose decode hooks only accept an empty payload, to pin the /// by-name decode fallback (registry miss → codec with `&[]`). #[derive(Debug)] @@ -989,6 +1041,30 @@ mod tests { ); } + #[test] + fn composed_codec_discards_bytes_from_a_failed_attempt() -> Result<()> { + // A codec may write before discovering it cannot encode the node. + // Those bytes must not be committed with the codec that succeeds next. + let composed = ComposedPhysicalExtensionCodec::new(vec![ + Arc::new(WritesThenRejectsCodec), + Arc::new(PlanPayloadCodec), + ]); + let plan: Arc = + Arc::new(EmptyExec::new(Arc::new(Schema::empty()))); + + let mut buf = vec![]; + composed.try_encode(plan, &mut buf, &DefaultPhysicalProtoConverter {})?; + + let proto = DataEncoderTuple::decode(buf.as_slice()) + .map_err(|e| internal_datafusion_err!("{e}"))?; + assert_eq!(proto.encoder_position, 1); + assert_eq!( + proto.blob, b"plan", + "bytes from the failed attempt must not be committed" + ); + Ok(()) + } + #[test] fn composed_codec_round_trips_a_child_function_payload() -> Result<()> { let composed = @@ -2096,6 +2172,9 @@ impl ComposedPhysicalExtensionCodec { // find the encoder for (position, codec) in self.codecs.iter().enumerate() { + // A codec may write before it rejects the node; those bytes must not + // reach whichever codec succeeds next. + data.clear(); match encode(codec.as_ref(), &mut data) { Ok(_) => { encoder_position = Some(position as u32);