Add Arc variants for the distributed user codec

Author: gabotechs

src/distributed_ext.rs

@@ -3,14 +3,15 @@ use crate::channel_resolver_ext::set_distributed_channel_resolver;
 use crate::config_extension_ext::{
     set_distributed_option_extension, set_distributed_option_extension_from_headers,
 };
-use crate::protobuf::set_distributed_user_codec;
+use crate::protobuf::{set_distributed_user_codec, set_distributed_user_codec_arc};
 use datafusion::common::DataFusionError;
 use datafusion::config::ConfigExtension;
 use datafusion::execution::{SessionState, SessionStateBuilder};
 use datafusion::prelude::{SessionConfig, SessionContext};
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use delegate::delegate;
 use http::HeaderMap;
+use std::sync::Arc;
 
 /// Extends DataFusion with distributed capabilities.
 pub trait DistributedExt: Sized {
@@ -167,6 +168,12 @@ pub trait DistributedExt: Sized {
     /// Same as [DistributedExt::with_distributed_user_codec] but with an in-place mutation
     fn set_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
 
+    /// Same as [DistributedExt::with_distributed_user_codec] but with a dynamic argument.
+    fn with_distributed_user_codec_arc(self, codec: Arc<dyn PhysicalExtensionCodec>) -> Self;
+
+    /// Same as [DistributedExt::set_distributed_user_codec] but with a dynamic argument.
+    fn set_distributed_user_codec_arc(&mut self, codec: Arc<dyn PhysicalExtensionCodec>);
+
     /// Injects a [ChannelResolver] implementation for Distributed DataFusion to resolve worker
     /// nodes. When running in distributed mode, setting a [ChannelResolver] is required.
     ///
@@ -234,6 +241,10 @@ impl DistributedExt for SessionConfig {
         set_distributed_user_codec(self, codec)
     }
 
+    fn set_distributed_user_codec_arc(&mut self, codec: Arc<dyn PhysicalExtensionCodec>) {
+        set_distributed_user_codec_arc(self, codec)
+    }
+
     fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(
         &mut self,
         resolver: T,
@@ -255,6 +266,10 @@ impl DistributedExt for SessionConfig {
             #[expr($;self)]
             fn with_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(mut self, codec: T) -> Self;
 
+            #[call(set_distributed_user_codec_arc)]
+            #[expr($;self)]
+            fn with_distributed_user_codec_arc(mut self, codec: Arc<dyn PhysicalExtensionCodec>) -> Self;
+
             #[call(set_distributed_channel_resolver)]
             #[expr($;self)]
             fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
@@ -280,6 +295,11 @@ impl DistributedExt for SessionStateBuilder {
             #[expr($;self)]
             fn with_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(mut self, codec: T) -> Self;
 
+            fn set_distributed_user_codec_arc(&mut self, codec: Arc<dyn PhysicalExtensionCodec>);
+            #[call(set_distributed_user_codec_arc)]
+            #[expr($;self)]
+            fn with_distributed_user_codec_arc(mut self, codec: Arc<dyn PhysicalExtensionCodec>) -> Self;
+
             fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
             #[call(set_distributed_channel_resolver)]
             #[expr($;self)]
@@ -306,6 +326,11 @@ impl DistributedExt for SessionState {
             #[expr($;self)]
             fn with_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(mut self, codec: T) -> Self;
 
+            fn set_distributed_user_codec_arc(&mut self, codec: Arc<dyn PhysicalExtensionCodec>);
+            #[call(set_distributed_user_codec_arc)]
+            #[expr($;self)]
+            fn with_distributed_user_codec_arc(mut self, codec: Arc<dyn PhysicalExtensionCodec>) -> Self;
+
             fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
             #[call(set_distributed_channel_resolver)]
             #[expr($;self)]
@@ -332,6 +357,11 @@ impl DistributedExt for SessionContext {
             #[expr($;self)]
             fn with_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(self, codec: T) -> Self;
 
+            fn set_distributed_user_codec_arc(&mut self, codec: Arc<dyn PhysicalExtensionCodec>);
+            #[call(set_distributed_user_codec_arc)]
+            #[expr($;self)]
+            fn with_distributed_user_codec_arc(mut self, codec: Arc<dyn PhysicalExtensionCodec>) -> Self;
+
             fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
             #[call(set_distributed_channel_resolver)]
             #[expr($;self)]

src/protobuf/mod.rs

@@ -4,4 +4,6 @@ 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_codecs, set_distributed_user_codec};
+pub(crate) use user_codec::{
+    get_distributed_user_codecs, set_distributed_user_codec, set_distributed_user_codec_arc,
+};

src/protobuf/user_codec.rs

@@ -4,18 +4,25 @@ use std::sync::Arc;
 
 pub struct UserProvidedCodecs(Vec<Arc<dyn PhysicalExtensionCodec>>);
 
-pub(crate) fn set_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(
+pub(crate) fn set_distributed_user_codec_arc(
     cfg: &mut SessionConfig,
-    codec: T,
+    codec: Arc<dyn PhysicalExtensionCodec>,
 ) {
     let mut codecs = match cfg.get_extension::<UserProvidedCodecs>() {
         None => vec![],
         Some(prev) => prev.0.clone(),
     };
-    codecs.push(Arc::new(codec));
+    codecs.push(codec);
     cfg.set_extension(Arc::new(UserProvidedCodecs(codecs)))
 }
 
+pub(crate) fn set_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(
+    cfg: &mut SessionConfig,
+    codec: T,
+) {
+    set_distributed_user_codec_arc(cfg, Arc::new(codec))
+}
+
 pub(crate) fn get_distributed_user_codecs(
     cfg: &SessionConfig,
 ) -> Vec<Arc<dyn PhysicalExtensionCodec>> {