Allow passing multiple user codecs

Author: gabotechs

src/distributed_ext.rs

@@ -125,7 +125,8 @@ pub trait DistributedExt: Sized {
     ) -> Result<(), DataFusionError>;
 
     /// Injects a user-defined [PhysicalExtensionCodec] that is capable of encoding/decoding
-    /// custom execution nodes.
+    /// custom execution nodes. Multiple user-defined [PhysicalExtensionCodec] can be added
+    /// by calling this method several times.
     ///
     /// Example:
     ///

src/protobuf/distributed_codec.rs

@@ -1,4 +1,4 @@
-use super::get_distributed_user_codec;
+use super::get_distributed_user_codecs;
 use crate::common::ComposedPhysicalExtensionCodec;
 use crate::execution_plans::{NetworkCoalesceExec, NetworkCoalesceReady, NetworkShuffleReadyExec};
 use crate::{NetworkShuffleExec, PartitionIsolatorExec};
@@ -25,9 +25,7 @@ pub struct DistributedCodec;
 impl DistributedCodec {
     pub fn new_combined_with_user(cfg: &SessionConfig) -> impl PhysicalExtensionCodec + use<> {
         let mut codecs: Vec<Arc<dyn PhysicalExtensionCodec>> = vec![Arc::new(DistributedCodec {})];
-        if let Some(ref user_codec) = get_distributed_user_codec(cfg) {
-            codecs.push(Arc::clone(user_codec));
-        }
+        codecs.extend(get_distributed_user_codecs(cfg));
         ComposedPhysicalExtensionCodec::new(codecs)
     }
 }

src/protobuf/mod.rs

@@ -4,4 +4,4 @@ mod user_codec;
 
 pub(crate) use distributed_codec::DistributedCodec;
 pub(crate) use stage_proto::{StageExecProto, StageKey, proto_from_stage, stage_from_proto};
-pub(crate) use user_codec::{get_distributed_user_codec, set_distributed_user_codec};
+pub(crate) use user_codec::{get_distributed_user_codecs, set_distributed_user_codec};

src/protobuf/user_codec.rs

@@ -2,17 +2,25 @@ use datafusion::prelude::SessionConfig;
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use std::sync::Arc;
 
-pub struct UserProvidedCodec(Arc<dyn PhysicalExtensionCodec>);
+pub struct UserProvidedCodecs(Vec<Arc<dyn PhysicalExtensionCodec>>);
 
 pub(crate) fn set_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(
     cfg: &mut SessionConfig,
     codec: T,
 ) {
-    cfg.set_extension(Arc::new(UserProvidedCodec(Arc::new(codec))))
+    let mut codecs = match cfg.get_extension::<UserProvidedCodecs>() {
+        None => vec![],
+        Some(prev) => prev.0.clone(),
+    };
+    codecs.push(Arc::new(codec));
+    cfg.set_extension(Arc::new(UserProvidedCodecs(codecs)))
 }
 
-pub(crate) fn get_distributed_user_codec(
+pub(crate) fn get_distributed_user_codecs(
     cfg: &SessionConfig,
-) -> Option<Arc<dyn PhysicalExtensionCodec>> {
-    Some(Arc::clone(&cfg.get_extension::<UserProvidedCodec>()?.0))
+) -> Vec<Arc<dyn PhysicalExtensionCodec>> {
+    match cfg.get_extension::<UserProvidedCodecs>() {
+        None => vec![],
+        Some(v) => v.0.clone(),
+    }
 }