fix: ComposedPhysicalExtensionCodec does not use the same codec as encoding when decoding (#16986)

* fix(example): `ComposedPhysicalExtensionCodec` does not use the same codec as encoding when decoding

Change-Id: I4e92abd9b49ee6f1fcdfd2d915fce8dc85c3b6a6

* naming and move `ComposedPhysicalExtensionCodec` to datafusion

Change-Id: I780e217f061a92cd7fe5fe9b91953e81d42d05c8

* rename

Change-Id: I6f3860053743b81a5598023becb2e18855904292

Author: Thearas

datafusion-examples/examples/composed_extension_codec.rs

@@ -35,13 +35,14 @@ use std::fmt::Debug;
 use std::ops::Deref;
 use std::sync::Arc;
 
+use datafusion::common::internal_err;
 use datafusion::common::Result;
-use datafusion::common::{internal_err, DataFusionError};
 use datafusion::logical_expr::registry::FunctionRegistry;
-use datafusion::logical_expr::{AggregateUDF, ScalarUDF};
 use datafusion::physical_plan::{DisplayAs, ExecutionPlan};
 use datafusion::prelude::SessionContext;
-use datafusion_proto::physical_plan::{AsExecutionPlan, PhysicalExtensionCodec};
+use datafusion_proto::physical_plan::{
+    AsExecutionPlan, ComposedPhysicalExtensionCodec, PhysicalExtensionCodec,
+};
 use datafusion_proto::protobuf;
 
 #[tokio::main]
@@ -54,12 +55,12 @@ async fn main() {
     });
     let ctx = SessionContext::new();
 
-    let composed_codec = ComposedPhysicalExtensionCodec {
-        codecs: vec![
-            Arc::new(ParentPhysicalExtensionCodec {}),
-            Arc::new(ChildPhysicalExtensionCodec {}),
-        ],
-    };
+    // Position in this list is important as it will be used for decoding.
+    // If new codec is added it should go to last position.
+    let composed_codec = ComposedPhysicalExtensionCodec::new(vec![
+        Arc::new(ParentPhysicalExtensionCodec {}),
+        Arc::new(ChildPhysicalExtensionCodec {}),
+    ]);
 
     // serialize execution plan to proto
     let proto: protobuf::PhysicalPlanNode =
@@ -232,60 +233,3 @@ impl PhysicalExtensionCodec for ChildPhysicalExtensionCodec {
         }
     }
 }
-
-/// A PhysicalExtensionCodec that tries one of multiple inner codecs
-/// until one works
-#[derive(Debug)]
-struct ComposedPhysicalExtensionCodec {
-    codecs: Vec<Arc<dyn PhysicalExtensionCodec>>,
-}
-
-impl ComposedPhysicalExtensionCodec {
-    fn try_any<T>(
-        &self,
-        mut f: impl FnMut(&dyn PhysicalExtensionCodec) -> Result<T>,
-    ) -> Result<T> {
-        let mut last_err = None;
-        for codec in &self.codecs {
-            match f(codec.as_ref()) {
-                Ok(node) => return Ok(node),
-                Err(err) => last_err = Some(err),
-            }
-        }
-
-        Err(last_err.unwrap_or_else(|| {
-            DataFusionError::NotImplemented("Empty list of composed codecs".to_owned())
-        }))
-    }
-}
-
-impl PhysicalExtensionCodec for ComposedPhysicalExtensionCodec {
-    fn try_decode(
-        &self,
-        buf: &[u8],
-        inputs: &[Arc<dyn ExecutionPlan>],
-        registry: &dyn FunctionRegistry,
-    ) -> Result<Arc<dyn ExecutionPlan>> {
-        self.try_any(|codec| codec.try_decode(buf, inputs, registry))
-    }
-
-    fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
-        self.try_any(|codec| codec.try_encode(node.clone(), buf))
-    }
-
-    fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> {
-        self.try_any(|codec| codec.try_decode_udf(name, buf))
-    }
-
-    fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> {
-        self.try_any(|codec| codec.try_encode_udf(node, buf))
-    }
-
-    fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result<Arc<AggregateUDF>> {
-        self.try_any(|codec| codec.try_decode_udaf(name, buf))
-    }
-
-    fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec<u8>) -> Result<()> {
-        self.try_any(|codec| codec.try_encode_udaf(node, buf))
-    }
-}

datafusion/proto/src/physical_plan/mod.rs

@@ -2941,12 +2941,126 @@ impl PhysicalExtensionCodec for DefaultPhysicalExtensionCodec {
     }
 }
 
+/// DataEncoderTuple captures the position of the encoder
+/// in the codec list that was used to encode the data and actual encoded data
+#[derive(Clone, PartialEq, prost::Message)]
+struct DataEncoderTuple {
+    /// The position of encoder used to encode data
+    /// (to be used for decoding)
+    #[prost(uint32, tag = 1)]
+    pub encoder_position: u32,
+
+    #[prost(bytes, tag = 2)]
+    pub blob: Vec<u8>,
+}
+
+/// A PhysicalExtensionCodec that tries one of multiple inner codecs
+/// until one works
+#[derive(Debug)]
+pub struct ComposedPhysicalExtensionCodec {
+    codecs: Vec<Arc<dyn PhysicalExtensionCodec>>,
+}
+
+impl ComposedPhysicalExtensionCodec {
+    // Position in this codesc list is important as it will be used for decoding.
+    // If new codec is added it should go to last position.
+    pub fn new(codecs: Vec<Arc<dyn PhysicalExtensionCodec>>) -> Self {
+        Self { codecs }
+    }
+
+    fn decode_protobuf<R>(
+        &self,
+        buf: &[u8],
+        decode: impl FnOnce(&dyn PhysicalExtensionCodec, &[u8]) -> Result<R>,
+    ) -> Result<R> {
+        let proto = DataEncoderTuple::decode(buf)
+            .map_err(|e| DataFusionError::Internal(e.to_string()))?;
+
+        let codec = self.codecs.get(proto.encoder_position as usize).ok_or(
+            DataFusionError::Internal(
+                "Can't find required codec in codec list".to_owned(),
+            ),
+        )?;
+
+        decode(codec.as_ref(), &proto.blob)
+    }
+
+    fn encode_protobuf(
+        &self,
+        buf: &mut Vec<u8>,
+        mut encode: impl FnMut(&dyn PhysicalExtensionCodec, &mut Vec<u8>) -> Result<()>,
+    ) -> Result<()> {
+        let mut data = vec![];
+        let mut last_err = None;
+        let mut encoder_position = None;
+
+        // find the encoder
+        for (position, codec) in self.codecs.iter().enumerate() {
+            match encode(codec.as_ref(), &mut data) {
+                Ok(_) => {
+                    encoder_position = Some(position as u32);
+                    break;
+                }
+                Err(err) => last_err = Some(err),
+            }
+        }
+
+        let encoder_position = encoder_position.ok_or_else(|| {
+            last_err.unwrap_or_else(|| {
+                DataFusionError::NotImplemented(
+                    "Empty list of composed codecs".to_owned(),
+                )
+            })
+        })?;
+
+        // encode with encoder position
+        let proto = DataEncoderTuple {
+            encoder_position,
+            blob: data,
+        };
+        proto
+            .encode(buf)
+            .map_err(|e| DataFusionError::Internal(e.to_string()))
+    }
+}
+
+impl PhysicalExtensionCodec for ComposedPhysicalExtensionCodec {
+    fn try_decode(
+        &self,
+        buf: &[u8],
+        inputs: &[Arc<dyn ExecutionPlan>],
+        registry: &dyn FunctionRegistry,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        self.decode_protobuf(buf, |codec, data| codec.try_decode(data, inputs, registry))
+    }
+
+    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_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> {
+        self.decode_protobuf(buf, |codec, data| codec.try_decode_udf(name, data))
+    }
+
+    fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> {
+        self.encode_protobuf(buf, |codec, data| codec.try_encode_udf(node, data))
+    }
+
+    fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result<Arc<AggregateUDF>> {
+        self.decode_protobuf(buf, |codec, data| codec.try_decode_udaf(name, data))
+    }
+
+    fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec<u8>) -> Result<()> {
+        self.encode_protobuf(buf, |codec, data| codec.try_encode_udaf(node, data))
+    }
+}
+
 fn into_physical_plan(
     node: &Option<Box<protobuf::PhysicalPlanNode>>,
     registry: &dyn FunctionRegistry,
     runtime: &RuntimeEnv,
     extension_codec: &dyn PhysicalExtensionCodec,
-) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
+) -> Result<Arc<dyn ExecutionPlan>> {
     if let Some(field) = node {
         field.try_into_physical_plan(registry, runtime, extension_codec)
     } else {