Evolve ChannelResolver trait for requiring a FlightClient instead of a channel (#172)

Author: gabotechs

examples/in_memory_cluster.rs

@@ -1,4 +1,5 @@
 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;
@@ -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<BoxCloneSyncChannel>,
 }
 
 impl InMemoryChannelResolver {
@@ -93,7 +94,7 @@ impl InMemoryChannelResolver {
             }));
 
         let this = Self {
-            channel: BoxCloneSyncChannel::new(channel),
+            channel: FlightServiceClient::new(BoxCloneSyncChannel::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<BoxCloneSyncChannel>, DataFusionError> {
         Ok(self.channel.clone())
     }
 }

examples/localhost_run.rs

@@ -1,4 +1,5 @@
 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;
@@ -83,7 +84,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
 #[derive(Clone)]
 struct LocalhostChannelResolver {
     ports: Vec<u16>,
-    cached: DashMap<Url, BoxCloneSyncChannel>,
+    cached: DashMap<Url, FlightServiceClient<BoxCloneSyncChannel>>,
 }
 
 #[async_trait]
@@ -96,14 +97,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<BoxCloneSyncChannel>, 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(BoxCloneSyncChannel::new(channel));
                 v.insert(channel.clone());
                 Ok(channel)
             }

examples/localhost_worker.rs

@@ -1,3 +1,4 @@
+use arrow_flight::flight_service_client::FlightServiceClient;
 use arrow_flight::flight_service_server::FlightServiceServer;
 use async_trait::async_trait;
 use dashmap::{DashMap, Entry};
@@ -55,7 +56,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
 #[derive(Clone)]
 struct LocalhostChannelResolver {
     ports: Vec<u16>,
-    cached: DashMap<Url, BoxCloneSyncChannel>,
+    cached: DashMap<Url, FlightServiceClient<BoxCloneSyncChannel>>,
 }
 
 #[async_trait]
@@ -68,14 +69,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<BoxCloneSyncChannel>, 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(BoxCloneSyncChannel::new(channel));
                 v.insert(channel.clone());
                 Ok(channel)
             }

src/channel_resolver_ext.rs

@@ -1,3 +1,4 @@
+use arrow_flight::flight_service_client::FlightServiceClient;
 use async_trait::async_trait;
 use datafusion::common::exec_datafusion_err;
 use datafusion::error::DataFusionError;
@@ -38,8 +39,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<BoxCloneSyncChannel>, DataFusionError>;
 }
 
 #[async_trait]
@@ -48,7 +52,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<BoxCloneSyncChannel>, DataFusionError> {
+        self.as_ref().get_flight_client_for_url(url).await
     }
 }

src/distributed_ext.rs

@@ -180,6 +180,7 @@ 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};
@@ -195,7 +196,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<BoxCloneSyncChannel>, DataFusionError> {
     ///         todo!()
     ///     }
     /// }

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)?

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)?

src/flight_service/do_get.rs

@@ -10,15 +10,15 @@ use crate::protobuf::{
     AppMetadata, DistributedCodec, FlightAppMetadata, MetricsCollection, StageKey, TaskMetrics,
     datafusion_error_to_tonic_status, stage_from_proto,
 };
-use arrow::array::RecordBatch;
-use arrow::datatypes::SchemaRef;
-use arrow::ipc::writer::{DictionaryTracker, IpcDataGenerator, IpcWriteOptions};
 use arrow_flight::FlightData;
 use arrow_flight::Ticket;
 use arrow_flight::encode::FlightDataEncoderBuilder;
 use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_server::FlightService;
 use bytes::Bytes;
+use datafusion::arrow::array::RecordBatch;
+use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::arrow::ipc::writer::{DictionaryTracker, IpcDataGenerator, IpcWriteOptions};
 use datafusion::common::exec_datafusion_err;
 use datafusion::execution::SendableRecordBatchStream;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;

src/test_utils/in_memory_channel_resolver.rs

@@ -2,6 +2,7 @@ use crate::{
     ArrowFlightEndpoint, BoxCloneSyncChannel, 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;
@@ -15,7 +16,7 @@ const DUMMY_URL: &str = "http://localhost:50051";
 /// tokio duplex rather than a TCP connection.
 #[derive(Clone)]
 pub struct InMemoryChannelResolver {
-    channel: BoxCloneSyncChannel,
+    channel: FlightServiceClient<BoxCloneSyncChannel>,
 }
 
 impl Default for InMemoryChannelResolver {
@@ -39,7 +40,7 @@ impl InMemoryChannelResolver {
             }));
 
         let this = Self {
-            channel: BoxCloneSyncChannel::new(channel),
+            channel: FlightServiceClient::new(BoxCloneSyncChannel::new(channel)),
         };
         let this_clone = this.clone();
 
@@ -73,10 +74,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<BoxCloneSyncChannel>, DataFusionError> {
         Ok(self.channel.clone())
     }
 }

src/test_utils/localhost.rs

@@ -3,6 +3,7 @@ use crate::{
     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 +99,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<BoxCloneSyncChannel>, 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(BoxCloneSyncChannel::new(channel)))
     }
 }