Allow passing a session builder to in_memory_channel_resolver.rs

Author: gabotechs

src/distributed_ext.rs

@@ -360,7 +360,7 @@ impl DistributedExt for SessionContext {
             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 with_distributed_user_codec_arc(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)]

src/execution_plans/metrics.rs

@@ -310,7 +310,7 @@ mod tests {
         let state = SessionStateBuilder::new()
             .with_default_features()
             .with_config(config)
-            .with_distributed_channel_resolver(InMemoryChannelResolver::new())
+            .with_distributed_channel_resolver(InMemoryChannelResolver::default())
             .with_physical_optimizer_rule(Arc::new(
                 DistributedPhysicalOptimizerRule::default()
                     .with_network_coalesce_tasks(2)

src/test_utils/in_memory_channel_resolver.rs

@@ -1,6 +1,7 @@
 use crate::{
     ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelResolver, DistributedExt,
-    DistributedSessionBuilderContext,
+    DistributedSessionBuilder, DistributedSessionBuilderContext,
+    MappedDistributedSessionBuilderExt,
 };
 use arrow_flight::flight_service_server::FlightServiceServer;
 use async_trait::async_trait;
@@ -20,12 +21,17 @@ pub struct InMemoryChannelResolver {
 
 impl Default for InMemoryChannelResolver {
     fn default() -> Self {
-        Self::new()
+        Self::new(|ctx: DistributedSessionBuilderContext| async move {
+            let builder = SessionStateBuilder::new()
+                .with_default_features()
+                .with_runtime_env(ctx.runtime_env.clone());
+            Ok(builder.build())
+        })
     }
 }
 
 impl InMemoryChannelResolver {
-    pub fn new() -> Self {
+    pub fn new(builder: impl DistributedSessionBuilder + Send + Sync + 'static) -> Self {
         let (client, server) = tokio::io::duplex(1024 * 1024);
 
         let mut client = Some(client);
@@ -43,18 +49,10 @@ impl InMemoryChannelResolver {
         };
         let this_clone = this.clone();
 
-        let endpoint =
-            ArrowFlightEndpoint::try_new(move |ctx: DistributedSessionBuilderContext| {
-                let this = this.clone();
-                async move {
-                    let builder = SessionStateBuilder::new()
-                        .with_default_features()
-                        .with_distributed_channel_resolver(this)
-                        .with_runtime_env(ctx.runtime_env.clone());
-                    Ok(builder.build())
-                }
-            })
-            .unwrap();
+        let builder =
+            builder.map(move |b| Ok(b.with_distributed_channel_resolver(this.clone()).build()));
+
+        let endpoint = ArrowFlightEndpoint::try_new(builder).unwrap();
 
         tokio::spawn(async move {
             Server::builder()