Evolve ChannelResolver trait for requiring a FlightClient instead of a channel

gabotechs · gabotechs · commit 5000e0bbecd3 · 2025-10-01T15:10:54.000+02:00
diff --git a/examples/in_memory_cluster.rs b/examples/in_memory_cluster.rs
@@ -1,13 +1,14 @@
 use arrow::util::pretty::pretty_format_batches;
+use arrow_flight::flight_service_client::FlightServiceClient;
 use arrow_flight::flight_service_server::FlightServiceServer;
 use async_trait::async_trait;
 use datafusion::common::DataFusionError;
 use datafusion::execution::SessionStateBuilder;
 use datafusion::physical_plan::displayable;
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
 use datafusion_distributed::{
-    ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelResolver, DistributedExt,
-    DistributedPhysicalOptimizerRule, DistributedSessionBuilderContext,
+    ArrowFlightEndpoint, ChannelResolver, DistributedExt, DistributedPhysicalOptimizerRule,
+    DistributedSessionBuilderContext,
 };
 use futures::TryStreamExt;
 use hyper_util::rt::TokioIo;
@@ -75,7 +76,7 @@ const DUMMY_URL: &str = "http://localhost:50051";
 /// tokio duplex rather than a TCP connection.
 #[derive(Clone)]
 struct InMemoryChannelResolver {
-    channel: BoxCloneSyncChannel,
+    channel: FlightServiceClient<tonic::transport::Channel>,
 }
 
 impl InMemoryChannelResolver {
@@ -93,7 +94,7 @@ impl InMemoryChannelResolver {
             }));
 
         let this = Self {
-            channel: BoxCloneSyncChannel::new(channel),
+            channel: FlightServiceClient::new(channel),
         };
         let this_clone = this.clone();
 
@@ -127,10 +128,10 @@ impl ChannelResolver for InMemoryChannelResolver {
         Ok(vec![url::Url::parse(DUMMY_URL).unwrap()])
     }
 
-    async fn get_channel_for_url(
+    async fn get_flight_client_for_url(
         &self,
         _: &url::Url,
-    ) -> Result<BoxCloneSyncChannel, DataFusionError> {
+    ) -> Result<FlightServiceClient<tonic::transport::Channel>, DataFusionError> {
         Ok(self.channel.clone())
     }
 }
diff --git a/examples/localhost_run.rs b/examples/localhost_run.rs
@@ -1,13 +1,12 @@
 use arrow::util::pretty::pretty_format_batches;
+use arrow_flight::flight_service_client::FlightServiceClient;
 use async_trait::async_trait;
 use dashmap::{DashMap, Entry};
 use datafusion::common::DataFusionError;
 use datafusion::execution::SessionStateBuilder;
 use datafusion::physical_plan::displayable;
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
-use datafusion_distributed::{
-    BoxCloneSyncChannel, ChannelResolver, DistributedExt, DistributedPhysicalOptimizerRule,
-};
+use datafusion_distributed::{ChannelResolver, DistributedExt, DistributedPhysicalOptimizerRule};
 use futures::TryStreamExt;
 use std::error::Error;
 use std::sync::Arc;
@@ -83,7 +82,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
 #[derive(Clone)]
 struct LocalhostChannelResolver {
     ports: Vec<u16>,
-    cached: DashMap<Url, BoxCloneSyncChannel>,
+    cached: DashMap<Url, FlightServiceClient<Channel>>,
 }
 
 #[async_trait]
@@ -96,14 +95,17 @@ impl ChannelResolver for LocalhostChannelResolver {
             .collect())
     }
 
-    async fn get_channel_for_url(&self, url: &Url) -> Result<BoxCloneSyncChannel, DataFusionError> {
+    async fn get_flight_client_for_url(
+        &self,
+        url: &Url,
+    ) -> Result<FlightServiceClient<Channel>, DataFusionError> {
         match self.cached.entry(url.clone()) {
             Entry::Occupied(v) => Ok(v.get().clone()),
             Entry::Vacant(v) => {
                 let channel = Channel::from_shared(url.to_string())
                     .unwrap()
                     .connect_lazy();
-                let channel = BoxCloneSyncChannel::new(channel);
+                let channel = FlightServiceClient::new(channel);
                 v.insert(channel.clone());
                 Ok(channel)
             }
diff --git a/examples/localhost_worker.rs b/examples/localhost_worker.rs
@@ -1,11 +1,11 @@
+use arrow_flight::flight_service_client::FlightServiceClient;
 use arrow_flight::flight_service_server::FlightServiceServer;
 use async_trait::async_trait;
 use dashmap::{DashMap, Entry};
 use datafusion::common::DataFusionError;
 use datafusion::execution::SessionStateBuilder;
 use datafusion_distributed::{
-    ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelResolver, DistributedExt,
-    DistributedSessionBuilderContext,
+    ArrowFlightEndpoint, ChannelResolver, DistributedExt, DistributedSessionBuilderContext,
 };
 use std::error::Error;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
@@ -55,7 +55,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
 #[derive(Clone)]
 struct LocalhostChannelResolver {
     ports: Vec<u16>,
-    cached: DashMap<Url, BoxCloneSyncChannel>,
+    cached: DashMap<Url, FlightServiceClient<Channel>>,
 }
 
 #[async_trait]
@@ -68,14 +68,17 @@ impl ChannelResolver for LocalhostChannelResolver {
             .collect())
     }
 
-    async fn get_channel_for_url(&self, url: &Url) -> Result<BoxCloneSyncChannel, DataFusionError> {
+    async fn get_flight_client_for_url(
+        &self,
+        url: &Url,
+    ) -> Result<FlightServiceClient<Channel>, DataFusionError> {
         match self.cached.entry(url.clone()) {
             Entry::Occupied(v) => Ok(v.get().clone()),
             Entry::Vacant(v) => {
                 let channel = Channel::from_shared(url.to_string())
                     .unwrap()
                     .connect_lazy();
-                let channel = BoxCloneSyncChannel::new(channel);
+                let channel = FlightServiceClient::new(channel);
                 v.insert(channel.clone());
                 Ok(channel)
             }
diff --git a/src/channel_resolver_ext.rs b/src/channel_resolver_ext.rs
@@ -1,9 +1,11 @@
+use arrow_flight::flight_service_client::FlightServiceClient;
 use async_trait::async_trait;
 use datafusion::common::exec_datafusion_err;
 use datafusion::error::DataFusionError;
 use datafusion::prelude::SessionConfig;
 use std::sync::Arc;
 use tonic::body::BoxBody;
+use tonic::transport::Channel;
 use url::Url;
 
 pub(crate) fn set_distributed_channel_resolver(
@@ -38,8 +40,11 @@ pub type BoxCloneSyncChannel = tower::util::BoxCloneSyncService<
 pub trait ChannelResolver {
     /// Gets all available worker URLs. Used during stage assignment.
     fn get_urls(&self) -> Result<Vec<Url>, DataFusionError>;
-    /// For a given URL, get a channel for communicating to it.
-    async fn get_channel_for_url(&self, url: &Url) -> Result<BoxCloneSyncChannel, DataFusionError>;
+    /// For a given URL, get an Arrow Flight client for communicating to it.
+    async fn get_flight_client_for_url(
+        &self,
+        url: &Url,
+    ) -> Result<FlightServiceClient<Channel>, DataFusionError>;
 }
 
 #[async_trait]
@@ -48,7 +53,10 @@ impl ChannelResolver for Arc<dyn ChannelResolver + Send + Sync> {
         self.as_ref().get_urls()
     }
 
-    async fn get_channel_for_url(&self, url: &Url) -> Result<BoxCloneSyncChannel, DataFusionError> {
-        self.as_ref().get_channel_for_url(url).await
+    async fn get_flight_client_for_url(
+        &self,
+        url: &Url,
+    ) -> Result<FlightServiceClient<Channel>, DataFusionError> {
+        self.as_ref().get_flight_client_for_url(url).await
     }
 }
diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs
@@ -180,10 +180,12 @@ pub trait DistributedExt: Sized {
     /// Example:
     ///
     /// ```
+    /// # use arrow_flight::flight_service_client::FlightServiceClient;
     /// # use async_trait::async_trait;
     /// # use datafusion::common::DataFusionError;
     /// # use datafusion::execution::{SessionState, SessionStateBuilder};
     /// # use datafusion::prelude::SessionConfig;
+    /// # use tonic::transport::Channel;
     /// # use url::Url;
     /// # use datafusion_distributed::{BoxCloneSyncChannel, ChannelResolver, DistributedExt, DistributedSessionBuilderContext};
     ///
@@ -195,7 +197,7 @@ pub trait DistributedExt: Sized {
     ///         todo!()
     ///     }
     ///
-    ///     async fn get_channel_for_url(&self, url: &Url) -> Result<BoxCloneSyncChannel, DataFusionError> {
+    ///     async fn get_flight_client_for_url(&self, url: &Url) -> Result<FlightServiceClient<Channel>, DataFusionError> {
     ///         todo!()
     ///     }
     /// }
diff --git a/src/execution_plans/network_coalesce.rs b/src/execution_plans/network_coalesce.rs
@@ -12,7 +12,6 @@ use crate::protobuf::{map_flight_to_datafusion_error, map_status_to_datafusion_e
 use arrow_flight::Ticket;
 use arrow_flight::decode::FlightRecordBatchStream;
 use arrow_flight::error::FlightError;
-use arrow_flight::flight_service_client::FlightServiceClient;
 use dashmap::DashMap;
 use datafusion::common::{exec_err, internal_datafusion_err, internal_err, plan_err};
 use datafusion::error::DataFusionError;
@@ -285,8 +284,8 @@ impl ExecutionPlan for NetworkCoalesceExec {
 
         let metrics_collection_capture = self_ready.metrics_collection.clone();
         let stream = async move {
-            let channel = channel_resolver.get_channel_for_url(&url).await?;
-            let stream = FlightServiceClient::new(channel)
+            let mut client = channel_resolver.get_flight_client_for_url(&url).await?;
+            let stream = client
                 .do_get(ticket)
                 .await
                 .map_err(map_status_to_datafusion_error)?
diff --git a/src/execution_plans/network_shuffle.rs b/src/execution_plans/network_shuffle.rs
@@ -12,7 +12,6 @@ use crate::protobuf::{map_flight_to_datafusion_error, map_status_to_datafusion_e
 use arrow_flight::Ticket;
 use arrow_flight::decode::FlightRecordBatchStream;
 use arrow_flight::error::FlightError;
-use arrow_flight::flight_service_client::FlightServiceClient;
 use dashmap::DashMap;
 use datafusion::common::{exec_err, internal_datafusion_err, internal_err, plan_err};
 use datafusion::error::DataFusionError;
@@ -337,8 +336,8 @@ impl ExecutionPlan for NetworkShuffleExec {
                     "NetworkShuffleExec: task is unassigned, cannot proceed"
                 ))?;
 
-                let channel = channel_resolver.get_channel_for_url(&url).await?;
-                let stream = FlightServiceClient::new(channel)
+                let mut client = channel_resolver.get_flight_client_for_url(&url).await?;
+                let stream = client
                     .do_get(ticket)
                     .await
                     .map_err(map_status_to_datafusion_error)?
diff --git a/src/test_utils/in_memory_channel_resolver.rs b/src/test_utils/in_memory_channel_resolver.rs
@@ -1,21 +1,21 @@
 use crate::{
-    ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelResolver, DistributedExt,
-    DistributedSessionBuilderContext,
+    ArrowFlightEndpoint, ChannelResolver, DistributedExt, DistributedSessionBuilderContext,
 };
+use arrow_flight::flight_service_client::FlightServiceClient;
 use arrow_flight::flight_service_server::FlightServiceServer;
 use async_trait::async_trait;
 use datafusion::common::DataFusionError;
 use datafusion::execution::SessionStateBuilder;
 use hyper_util::rt::TokioIo;
-use tonic::transport::{Endpoint, Server};
+use tonic::transport::{Channel, Endpoint, Server};
 
 const DUMMY_URL: &str = "http://localhost:50051";
 
 /// [ChannelResolver] implementation that returns gRPC clients backed by an in-memory
 /// tokio duplex rather than a TCP connection.
 #[derive(Clone)]
 pub struct InMemoryChannelResolver {
-    channel: BoxCloneSyncChannel,
+    channel: FlightServiceClient<Channel>,
 }
 
 impl Default for InMemoryChannelResolver {
@@ -39,7 +39,7 @@ impl InMemoryChannelResolver {
             }));
 
         let this = Self {
-            channel: BoxCloneSyncChannel::new(channel),
+            channel: FlightServiceClient::new(channel),
         };
         let this_clone = this.clone();
 
@@ -73,10 +73,10 @@ impl ChannelResolver for InMemoryChannelResolver {
         Ok(vec![url::Url::parse(DUMMY_URL).unwrap()])
     }
 
-    async fn get_channel_for_url(
+    async fn get_flight_client_for_url(
         &self,
         _: &url::Url,
-    ) -> Result<BoxCloneSyncChannel, DataFusionError> {
+    ) -> Result<FlightServiceClient<tonic::transport::Channel>, DataFusionError> {
         Ok(self.channel.clone())
     }
 }
diff --git a/src/test_utils/localhost.rs b/src/test_utils/localhost.rs
@@ -1,8 +1,8 @@
 use crate::{
-    ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelResolver, DistributedExt,
-    DistributedSessionBuilder, DistributedSessionBuilderContext,
-    MappedDistributedSessionBuilderExt,
+    ArrowFlightEndpoint, ChannelResolver, DistributedExt, DistributedSessionBuilder,
+    DistributedSessionBuilderContext, MappedDistributedSessionBuilderExt,
 };
+use arrow_flight::flight_service_client::FlightServiceClient;
 use arrow_flight::flight_service_server::FlightServiceServer;
 use async_trait::async_trait;
 use datafusion::common::DataFusionError;
@@ -98,10 +98,13 @@ impl ChannelResolver for LocalHostChannelResolver {
             .map(|url| Url::parse(&url).map_err(external_err))
             .collect::<Result<Vec<Url>, _>>()
     }
-    async fn get_channel_for_url(&self, url: &Url) -> Result<BoxCloneSyncChannel, DataFusionError> {
+    async fn get_flight_client_for_url(
+        &self,
+        url: &Url,
+    ) -> Result<FlightServiceClient<Channel>, DataFusionError> {
         let endpoint = Channel::from_shared(url.to_string()).map_err(external_err)?;
         let channel = endpoint.connect().await.map_err(external_err)?;
-        Ok(BoxCloneSyncChannel::new(channel))
+        Ok(FlightServiceClient::new(channel))
     }
 }
 
