Rename methods

Author: gabotechs

src/channel_manager_ext.rs

@@ -5,7 +5,7 @@ use std::sync::Arc;
 use tonic::body::BoxBody;
 use url::Url;
 
-pub(crate) fn set_channel_resolver(
+pub(crate) fn set_distributed_channel_resolver(
     cfg: &mut SessionConfig,
     channel_resolver: impl ChannelResolver + Send + Sync + 'static,
 ) {
@@ -14,7 +14,7 @@ pub(crate) fn set_channel_resolver(
     ))));
 }
 
-pub(crate) fn get_channel_resolver(
+pub(crate) fn get_distributed_channel_resolver(
     cfg: &SessionConfig,
 ) -> Option<Arc<dyn ChannelResolver + Send + Sync>> {
     cfg.get_extension::<ChannelResolverExtension>()

src/distributed_ext.rs

@@ -1,8 +1,8 @@
-use crate::channel_manager_ext::set_channel_resolver;
+use crate::channel_manager_ext::set_distributed_channel_resolver;
 use crate::config_extension_ext::{
     set_distributed_option_extension, set_distributed_option_extension_from_headers,
 };
-use crate::user_codec_ext::set_user_codec;
+use crate::user_codec_ext::set_distributed_user_codec;
 use crate::ChannelResolver;
 use datafusion::common::DataFusionError;
 use datafusion::config::ConfigExtension;
@@ -151,20 +151,20 @@ pub trait DistributedExt: Sized {
     ///     }
     /// }
     ///
-    /// let config = SessionConfig::new().with_user_codec(CustomExecCodec);
+    /// let config = SessionConfig::new().with_distributed_user_codec(CustomExecCodec);
     ///
     /// async fn build_state(ctx: DistributedSessionBuilderContext) -> Result<SessionState, DataFusionError> {
     ///     // while providing this MyCustomSessionBuilder to an Arrow Flight endpoint, it will
     ///     // know how to deserialize the CustomExtension from the gRPC metadata.
     ///     Ok(SessionStateBuilder::new()
-    ///         .with_user_codec(CustomExecCodec)
+    ///         .with_distributed_user_codec(CustomExecCodec)
     ///         .build())
     /// }
     /// ```
-    fn with_user_codec<T: PhysicalExtensionCodec + 'static>(self, codec: T) -> Self;
+    fn with_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(self, codec: T) -> Self;
 
-    /// Same as [DistributedExt::with_user_codec] but with an in-place mutation
-    fn set_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
+    /// 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);
 
     /// Injects a [ChannelResolver] implementation for Distributed DataFusion to resolve worker
     /// nodes. When running in distributed mode, setting a [ChannelResolver] is required.
@@ -192,21 +192,26 @@ pub trait DistributedExt: Sized {
     ///     }
     /// }
     ///
-    /// let config = SessionConfig::new().with_channel_resolver(CustomChannelResolver);
+    /// let config = SessionConfig::new().with_distributed_channel_resolver(CustomChannelResolver);
     ///
     /// async fn build_state(ctx: DistributedSessionBuilderContext) -> Result<SessionState, DataFusionError> {
     ///     // while providing this MyCustomSessionBuilder to an Arrow Flight endpoint, it will
     ///     // know how to deserialize the CustomExtension from the gRPC metadata.
     ///     Ok(SessionStateBuilder::new()
-    ///         .with_channel_resolver(CustomChannelResolver)
+    ///         .with_distributed_channel_resolver(CustomChannelResolver)
     ///         .build())
     /// }
     /// ```
-    fn with_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(self, resolver: T)
-        -> Self;
+    fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(
+        self,
+        resolver: T,
+    ) -> Self;
 
-    /// Same as [DistributedExt::with_channel_resolver] but with an in-place mutation.
-    fn set_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
+    /// Same as [DistributedExt::with_distributed_channel_resolver] but with an in-place mutation.
+    fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(
+        &mut self,
+        resolver: T,
+    );
 }
 
 impl DistributedExt for SessionConfig {
@@ -224,12 +229,15 @@ impl DistributedExt for SessionConfig {
         set_distributed_option_extension_from_headers::<T>(self, headers)
     }
 
-    fn set_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T) {
-        set_user_codec(self, codec)
+    fn set_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T) {
+        set_distributed_user_codec(self, codec)
     }
 
-    fn set_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T) {
-        set_channel_resolver(self, resolver)
+    fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(
+        &mut self,
+        resolver: T,
+    ) {
+        set_distributed_channel_resolver(self, resolver)
     }
 
     delegate! {
@@ -242,13 +250,13 @@ impl DistributedExt for SessionConfig {
             #[expr($?;Ok(self))]
             fn with_distributed_option_extension_from_headers<T: ConfigExtension + Default>(mut self, headers: &HeaderMap) -> Result<Self, DataFusionError>;
 
-            #[call(set_user_codec)]
+            #[call(set_distributed_user_codec)]
             #[expr($;self)]
-            fn with_user_codec<T: PhysicalExtensionCodec + 'static>(mut self, codec: T) -> Self;
+            fn with_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(mut self, codec: T) -> Self;
 
-            #[call(set_channel_resolver)]
+            #[call(set_distributed_channel_resolver)]
             #[expr($;self)]
-            fn with_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
+            fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
         }
     }
 }
@@ -266,15 +274,15 @@ impl DistributedExt for SessionStateBuilder {
             #[expr($?;Ok(self))]
             fn with_distributed_option_extension_from_headers<T: ConfigExtension + Default>(mut self, headers: &HeaderMap) -> Result<Self, DataFusionError>;
 
-            fn set_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
-            #[call(set_user_codec)]
+            fn set_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
+            #[call(set_distributed_user_codec)]
             #[expr($;self)]
-            fn with_user_codec<T: PhysicalExtensionCodec + 'static>(mut self, codec: T) -> Self;
+            fn with_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(mut self, codec: T) -> Self;
 
-            fn set_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
-            #[call(set_channel_resolver)]
+            fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
+            #[call(set_distributed_channel_resolver)]
             #[expr($;self)]
-            fn with_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
+            fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
         }
     }
 }
@@ -292,15 +300,15 @@ impl DistributedExt for SessionState {
             #[expr($?;Ok(self))]
             fn with_distributed_option_extension_from_headers<T: ConfigExtension + Default>(mut self, headers: &HeaderMap) -> Result<Self, DataFusionError>;
 
-            fn set_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
-            #[call(set_user_codec)]
+            fn set_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
+            #[call(set_distributed_user_codec)]
             #[expr($;self)]
-            fn with_user_codec<T: PhysicalExtensionCodec + 'static>(mut self, codec: T) -> Self;
+            fn with_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(mut self, codec: T) -> Self;
 
-            fn set_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
-            #[call(set_channel_resolver)]
+            fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
+            #[call(set_distributed_channel_resolver)]
             #[expr($;self)]
-            fn with_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
+            fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
         }
     }
 }
@@ -318,15 +326,15 @@ impl DistributedExt for SessionContext {
             #[expr($?;Ok(self))]
             fn with_distributed_option_extension_from_headers<T: ConfigExtension + Default>(self, headers: &HeaderMap) -> Result<Self, DataFusionError>;
 
-            fn set_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
-            #[call(set_user_codec)]
+            fn set_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T);
+            #[call(set_distributed_user_codec)]
             #[expr($;self)]
-            fn with_user_codec<T: PhysicalExtensionCodec + 'static>(self, codec: T) -> Self;
+            fn with_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(self, codec: T) -> Self;
 
-            fn set_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
-            #[call(set_channel_resolver)]
+            fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
+            #[call(set_distributed_channel_resolver)]
             #[expr($;self)]
-            fn with_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(self, resolver: T) -> Self;
+            fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(self, resolver: T) -> Self;
         }
     }
 }

src/flight_service/do_get.rs

@@ -6,7 +6,7 @@ use crate::flight_service::service::ArrowFlightEndpoint;
 use crate::flight_service::session_builder::DistributedSessionBuilderContext;
 use crate::plan::{DistributedCodec, PartitionGroup};
 use crate::stage::{stage_from_proto, ExecutionStage, ExecutionStageProto};
-use crate::user_codec_ext::get_user_codec;
+use crate::user_codec_ext::get_distributed_user_codec;
 use arrow_flight::encode::FlightDataEncoderBuilder;
 use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_server::FlightService;
@@ -117,7 +117,7 @@ impl ArrowFlightEndpoint {
 
                 let mut combined_codec = ComposedPhysicalExtensionCodec::default();
                 combined_codec.push(DistributedCodec);
-                if let Some(ref user_codec) = get_user_codec(state.config()) {
+                if let Some(ref user_codec) = get_distributed_user_codec(state.config()) {
                     combined_codec.push_arc(Arc::clone(user_codec));
                 }
 

src/flight_service/session_builder.rs

@@ -52,7 +52,7 @@ pub trait DistributedSessionBuilder {
     ///         let mut builder = SessionStateBuilder::new()
     ///             .with_runtime_env(ctx.runtime_env.clone())
     ///             .with_default_features();
-    ///         builder.set_user_codec(CustomExecCodec);
+    ///         builder.set_distributed_user_codec(CustomExecCodec);
     ///         // Add your UDFs, optimization rules, etc...
     ///
     ///         Ok(builder.build())

src/plan/arrow_flight_read.rs

@@ -1,12 +1,12 @@
 use super::combined::CombinedRecordBatchStream;
-use crate::channel_manager_ext::get_channel_resolver;
+use crate::channel_manager_ext::get_distributed_channel_resolver;
 use crate::common::ComposedPhysicalExtensionCodec;
 use crate::config_extension_ext::ContextGrpcMetadata;
 use crate::errors::tonic_status_to_datafusion_error;
 use crate::flight_service::{DoGet, StageKey};
 use crate::plan::DistributedCodec;
 use crate::stage::{proto_from_stage, ExecutionStage};
-use crate::user_codec_ext::get_user_codec;
+use crate::user_codec_ext::get_distributed_user_codec;
 use crate::ChannelResolver;
 use arrow_flight::decode::FlightRecordBatchStream;
 use arrow_flight::error::FlightError;
@@ -159,7 +159,8 @@ impl ExecutionPlan for ArrowFlightReadExec {
         };
 
         // get the channel manager and current stage from our context
-        let Some(channel_resolver) = get_channel_resolver(context.session_config()) else {
+        let Some(channel_resolver) = get_distributed_channel_resolver(context.session_config())
+        else {
             return exec_err!(
                 "ArrowFlightReadExec requires a ChannelResolver in the session config"
             );
@@ -188,7 +189,7 @@ impl ExecutionPlan for ArrowFlightReadExec {
 
         let mut combined_codec = ComposedPhysicalExtensionCodec::default();
         combined_codec.push(DistributedCodec {});
-        if let Some(ref user_codec) = get_user_codec(context.session_config()) {
+        if let Some(ref user_codec) = get_distributed_user_codec(context.session_config()) {
             combined_codec.push_arc(Arc::clone(user_codec));
         }
 

src/stage/execution_stage.rs

@@ -6,7 +6,7 @@ use datafusion::execution::TaskContext;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::SessionContext;
 
-use crate::channel_manager_ext::get_channel_resolver;
+use crate::channel_manager_ext::get_distributed_channel_resolver;
 use crate::task::ExecutionTask;
 use crate::ChannelResolver;
 use itertools::Itertools;
@@ -261,7 +261,8 @@ impl ExecutionPlan for ExecutionStage {
             .downcast_ref::<ExecutionStage>()
             .expect("Unwrapping myself should always work");
 
-        let Some(channel_resolver) = get_channel_resolver(context.session_config()) else {
+        let Some(channel_resolver) = get_distributed_channel_resolver(context.session_config())
+        else {
             return exec_err!("ChannelManager not found in session config");
         };
 

src/test_utils/localhost.rs

@@ -46,7 +46,9 @@ where
     let channel_resolver = LocalHostChannelResolver::new(ports.clone());
     let session_builder = session_builder.map(move |builder: SessionStateBuilder| {
         let channel_resolver = channel_resolver.clone();
-        Ok(builder.with_channel_resolver(channel_resolver).build())
+        Ok(builder
+            .with_distributed_channel_resolver(channel_resolver)
+            .build())
     });
     let mut join_set = JoinSet::new();
     for listener in listeners {

src/user_codec_ext.rs

@@ -4,13 +4,15 @@ use std::sync::Arc;
 
 pub struct UserProvidedCodec(Arc<dyn PhysicalExtensionCodec>);
 
-pub(crate) fn set_user_codec<T: PhysicalExtensionCodec + 'static>(
+pub(crate) fn set_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(
     cfg: &mut SessionConfig,
     codec: T,
 ) {
     cfg.set_extension(Arc::new(UserProvidedCodec(Arc::new(codec))))
 }
 
-pub(crate) fn get_user_codec(cfg: &SessionConfig) -> Option<Arc<dyn PhysicalExtensionCodec>> {
+pub(crate) fn get_distributed_user_codec(
+    cfg: &SessionConfig,
+) -> Option<Arc<dyn PhysicalExtensionCodec>> {
     Some(Arc::clone(&cfg.get_extension::<UserProvidedCodec>()?.0))
 }

tests/custom_config_extension.rs

@@ -32,7 +32,7 @@ mod tests {
                 .with_runtime_env(ctx.runtime_env)
                 .with_default_features()
                 .with_distributed_option_extension_from_headers::<CustomExtension>(&ctx.headers)?
-                .with_user_codec(CustomConfigExtensionRequiredExecCodec)
+                .with_distributed_user_codec(CustomConfigExtensionRequiredExecCodec)
                 .build())
         }
 

tests/custom_extension_codec.rs

@@ -44,7 +44,7 @@ mod tests {
             Ok(SessionStateBuilder::new()
                 .with_runtime_env(ctx.runtime_env)
                 .with_default_features()
-                .with_user_codec(Int64ListExecCodec)
+                .with_distributed_user_codec(Int64ListExecCodec)
                 .build())
         }