use usize::MAX and add helpers to create clients/servers

Author: adriangb

examples/in_memory_cluster.rs

@@ -1,14 +1,13 @@
 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,
+    DistributedPhysicalOptimizerRule, DistributedSessionBuilderContext, create_flight_client,
 };
 use futures::TryStreamExt;
 use hyper_util::rt::TokioIo;
@@ -76,13 +75,6 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
 const DUMMY_URL: &str = "http://localhost:50051";
 
-/// Maximum message size for FlightData chunks in ArrowFlightEndpoint.
-const ENDPOINT_MESSAGE_SIZE: usize = 128 * 1024 * 1024; // 128 MB
-
-/// Maximum message size for gRPC server encoding and decoding.
-/// This should be 2x the ArrowFlightEndpoint max_message_size to allow for overhead.
-const MAX_MESSAGE_SIZE: usize = 256 * 1024 * 1024; // 256 MB
-
 /// [ChannelResolver] implementation that returns gRPC clients baked by an in-memory
 /// tokio duplex rather than a TCP connection.
 #[derive(Clone)]
@@ -105,7 +97,7 @@ impl InMemoryChannelResolver {
             }));
 
         let this = Self {
-            channel: FlightServiceClient::new(BoxCloneSyncChannel::new(channel)),
+            channel: create_flight_client(BoxCloneSyncChannel::new(channel)),
         };
         let this_clone = this.clone();
 
@@ -120,16 +112,11 @@ impl InMemoryChannelResolver {
                     Ok(builder.build())
                 }
             })
-            .unwrap()
-            .with_max_message_size(ENDPOINT_MESSAGE_SIZE);
+            .unwrap();
 
         tokio::spawn(async move {
             Server::builder()
-                .add_service(
-                    FlightServiceServer::new(endpoint)
-                        .max_decoding_message_size(MAX_MESSAGE_SIZE)
-                        .max_encoding_message_size(MAX_MESSAGE_SIZE),
-                )
+                .add_service(endpoint.into_flight_server())
                 .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server)))
                 .await
         });

examples/localhost_worker.rs

@@ -1,12 +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,
+    DistributedSessionBuilderContext, create_flight_client,
 };
 use std::error::Error;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
@@ -25,13 +24,6 @@ struct Args {
     cluster_ports: Vec<u16>,
 }
 
-/// Maximum message size for FlightData chunks in ArrowFlightEndpoint.
-const ENDPOINT_MESSAGE_SIZE: usize = 128 * 1024 * 1024; // 128 MB
-
-/// Maximum message size for gRPC server encoding and decoding.
-/// This should be 2x the ArrowFlightEndpoint max_message_size to allow for overhead.
-const MAX_MESSAGE_SIZE: usize = 256 * 1024 * 1024; // 256 MB
-
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn Error>> {
     let args = Args::from_args();
@@ -50,15 +42,10 @@ async fn main() -> Result<(), Box<dyn Error>> {
                 .with_default_features()
                 .build())
         }
-    })?
-    .with_max_message_size(ENDPOINT_MESSAGE_SIZE);
+    })?;
 
     Server::builder()
-        .add_service(
-            FlightServiceServer::new(endpoint)
-                .max_decoding_message_size(MAX_MESSAGE_SIZE)
-                .max_encoding_message_size(MAX_MESSAGE_SIZE),
-        )
+        .add_service(endpoint.into_flight_server())
         .serve(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), args.port))
         .await?;
 
@@ -91,7 +78,7 @@ impl ChannelResolver for LocalhostChannelResolver {
                 let channel = Channel::from_shared(url.to_string())
                     .unwrap()
                     .connect_lazy();
-                let channel = FlightServiceClient::new(BoxCloneSyncChannel::new(channel));
+                let channel = create_flight_client(BoxCloneSyncChannel::new(channel));
                 v.insert(channel.clone());
                 Ok(channel)
             }

src/channel_resolver_ext.rs

@@ -35,11 +35,21 @@ pub type BoxCloneSyncChannel = tower::util::BoxCloneSyncService<
 
 /// Abstracts networking details so that users can implement their own network resolution
 /// mechanism.
+///
+/// # Implementation Note
+///
+/// When implementing `get_flight_client_for_url`, it is recommended to use the
+/// [`create_flight_client`] helper function to ensure clients are configured with
+/// appropriate message size limits for internal communication. This helps avoid message
+/// size errors when transferring large datasets.
 #[async_trait]
 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 an Arrow Flight client for communicating to it.
+    ///
+    /// Consider using [`create_flight_client`] to create the client with appropriate
+    /// default message size limits.
     async fn get_flight_client_for_url(
         &self,
         url: &Url,
@@ -59,3 +69,39 @@ impl ChannelResolver for Arc<dyn ChannelResolver + Send + Sync> {
         self.as_ref().get_flight_client_for_url(url).await
     }
 }
+
+/// Creates a [`FlightServiceClient`] with high default message size limits.
+///
+/// This is a convenience function that wraps [`FlightServiceClient::new`] and configures
+/// it with `max_decoding_message_size(usize::MAX)` and `max_encoding_message_size(usize::MAX)`
+/// to avoid message size limitations for internal communication.
+///
+/// Users implementing custom [`ChannelResolver`]s should use this function in their
+/// `get_flight_client_for_url` implementations to ensure consistent behavior with built-in
+/// implementations.
+///
+/// # Example
+///
+/// ```rust,ignore
+/// use datafusion_distributed::{create_flight_client, BoxCloneSyncChannel, ChannelResolver};
+/// use arrow_flight::flight_service_client::FlightServiceClient;
+/// use tonic::transport::Channel;
+///
+/// #[async_trait]
+/// impl ChannelResolver for MyResolver {
+///     async fn get_flight_client_for_url(
+///         &self,
+///         url: &Url,
+///     ) -> Result<FlightServiceClient<BoxCloneSyncChannel>, DataFusionError> {
+///         let channel = Channel::from_shared(url.to_string())?.connect().await?;
+///         Ok(create_flight_client(BoxCloneSyncChannel::new(channel)))
+///     }
+/// }
+/// ```
+pub fn create_flight_client(
+    channel: BoxCloneSyncChannel,
+) -> FlightServiceClient<BoxCloneSyncChannel> {
+    FlightServiceClient::new(channel)
+        .max_decoding_message_size(usize::MAX)
+        .max_encoding_message_size(usize::MAX)
+}

src/flight_service/service.rs

@@ -2,7 +2,7 @@ use crate::common::ttl_map::{TTLMap, TTLMapConfig};
 use crate::flight_service::DistributedSessionBuilder;
 use crate::flight_service::do_get::TaskData;
 use crate::protobuf::StageKey;
-use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::flight_service_server::{FlightService, FlightServiceServer};
 use arrow_flight::{
     Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo,
     HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket,
@@ -41,7 +41,7 @@ impl ArrowFlightEndpoint {
             task_data_entries: Arc::new(ttl_map),
             session_builder: Arc::new(session_builder),
             hooks: ArrowFlightEndpointHooks::default(),
-            max_message_size: None,
+            max_message_size: Some(usize::MAX),
         })
     }
 
@@ -59,20 +59,41 @@ impl ArrowFlightEndpoint {
 
     /// Set the maximum message size for FlightData chunks.
     ///
-    /// Defaults to None, which uses `arrow-rs` default, curerntly 2MB.
+    /// Defaults to `usize::MAX` to minimize chunking overhead for internal communication.
     /// See [`FlightDataEncoderBuilder::with_max_flight_data_size`] for details.
     ///
-    /// If you change this, ensure you configure the server's max_encoding_message_size and
-    /// max_decoding_message_size to at least 2x this value to allow for overhead.
-    /// If your service communication is purely internal and there is no risk of DOS attacks,
-    /// you may want to set this to a considerably larger value to minimize the overhead of chunking
-    /// larger datasets.
+    /// If you change this to a lower value, ensure you configure the server's
+    /// max_encoding_message_size and max_decoding_message_size to at least 2x this value
+    /// to allow for overhead. For most use cases, the default of `usize::MAX` is appropriate.
     ///
     /// [`FlightDataEncoderBuilder::with_max_flight_data_size`]: https://arrow.apache.org/rust/arrow_flight/encode/struct.FlightDataEncoderBuilder.html#structfield.max_flight_data_size
     pub fn with_max_message_size(mut self, size: usize) -> Self {
         self.max_message_size = Some(size);
         self
     }
+
+    /// Converts this endpoint into a [`FlightServiceServer`] with high default message size limits.
+    ///
+    /// This is a convenience method that wraps the endpoint in a [`FlightServiceServer`] and
+    /// configures it with `max_decoding_message_size(usize::MAX)` and
+    /// `max_encoding_message_size(usize::MAX)` to avoid message size limitations for internal
+    /// communication.
+    ///
+    /// You can further customize the returned server by chaining additional tonic methods.
+    ///
+    /// # Example
+    ///
+    /// ```rust,ignore
+    /// let endpoint = ArrowFlightEndpoint::try_new(session_builder)?;
+    /// let server = endpoint.into_flight_server();
+    /// // Can chain additional tonic methods if needed
+    /// // let server = server.some_other_tonic_method(...);
+    /// ```
+    pub fn into_flight_server(self) -> FlightServiceServer<Self> {
+        FlightServiceServer::new(self)
+            .max_decoding_message_size(usize::MAX)
+            .max_encoding_message_size(usize::MAX)
+    }
 }
 
 #[async_trait]

src/lib.rs

@@ -14,7 +14,7 @@ mod protobuf;
 #[cfg(any(feature = "integration", test))]
 pub mod test_utils;
 
-pub use channel_resolver_ext::{BoxCloneSyncChannel, ChannelResolver};
+pub use channel_resolver_ext::{BoxCloneSyncChannel, ChannelResolver, create_flight_client};
 pub use distributed_ext::DistributedExt;
 pub use distributed_planner::{
     DistributedConfig, DistributedPhysicalOptimizerRule, InputStageInfo, NetworkBoundary,

src/test_utils/in_memory_channel_resolver.rs

@@ -1,9 +1,8 @@
 use crate::{
     ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelResolver, DistributedExt,
-    DistributedSessionBuilderContext,
+    DistributedSessionBuilderContext, create_flight_client,
 };
 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;
@@ -12,13 +11,6 @@ use tonic::transport::{Endpoint, Server};
 
 const DUMMY_URL: &str = "http://localhost:50051";
 
-/// Maximum message size for FlightData chunks in ArrowFlightEndpoint.
-const ENDPOINT_MESSAGE_SIZE: usize = 128 * 1024 * 1024; // 128 MB
-
-/// Maximum message size for gRPC server encoding and decoding.
-/// This should be 2x the ArrowFlightEndpoint max_message_size to allow for overhead.
-const MAX_MESSAGE_SIZE: usize = 256 * 1024 * 1024; // 256 MB
-
 /// [ChannelResolver] implementation that returns gRPC clients backed by an in-memory
 /// tokio duplex rather than a TCP connection.
 #[derive(Clone)]
@@ -47,7 +39,7 @@ impl InMemoryChannelResolver {
             }));
 
         let this = Self {
-            channel: FlightServiceClient::new(BoxCloneSyncChannel::new(channel)),
+            channel: create_flight_client(BoxCloneSyncChannel::new(channel)),
         };
         let this_clone = this.clone();
 
@@ -62,16 +54,11 @@ impl InMemoryChannelResolver {
                     Ok(builder.build())
                 }
             })
-            .unwrap()
-            .with_max_message_size(ENDPOINT_MESSAGE_SIZE);
+            .unwrap();
 
         tokio::spawn(async move {
             Server::builder()
-                .add_service(
-                    FlightServiceServer::new(endpoint)
-                        .max_decoding_message_size(MAX_MESSAGE_SIZE)
-                        .max_encoding_message_size(MAX_MESSAGE_SIZE),
-                )
+                .add_service(endpoint.into_flight_server())
                 .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server)))
                 .await
         });

src/test_utils/localhost.rs

@@ -1,10 +1,9 @@
 use crate::{
     ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelResolver, DistributedExt,
     DistributedSessionBuilder, DistributedSessionBuilderContext,
-    MappedDistributedSessionBuilderExt,
+    MappedDistributedSessionBuilderExt, create_flight_client,
 };
 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::common::runtime::JoinSet;
@@ -18,13 +17,6 @@ use tokio::net::TcpListener;
 use tonic::transport::{Channel, Server};
 use url::Url;
 
-/// Maximum message size for FlightData chunks in ArrowFlightEndpoint.
-const ENDPOINT_MESSAGE_SIZE: usize = 128 * 1024 * 1024; // 128 MB
-
-/// Maximum message size for gRPC server encoding and decoding.
-/// This should be 2x the ArrowFlightEndpoint max_message_size to allow for overhead.
-const MAX_MESSAGE_SIZE: usize = 256 * 1024 * 1024; // 256 MB
-
 pub async fn start_localhost_context<B>(
     num_workers: usize,
     session_builder: B,
@@ -112,25 +104,20 @@ impl ChannelResolver for LocalHostChannelResolver {
     ) -> 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(FlightServiceClient::new(BoxCloneSyncChannel::new(channel)))
+        Ok(create_flight_client(BoxCloneSyncChannel::new(channel)))
     }
 }
 
 pub async fn spawn_flight_service(
     session_builder: impl DistributedSessionBuilder + Send + Sync + 'static,
     incoming: TcpListener,
 ) -> Result<(), Box<dyn Error + Send + Sync>> {
-    let endpoint =
-        ArrowFlightEndpoint::try_new(session_builder)?.with_max_message_size(ENDPOINT_MESSAGE_SIZE);
+    let endpoint = ArrowFlightEndpoint::try_new(session_builder)?;
 
     let incoming = tokio_stream::wrappers::TcpListenerStream::new(incoming);
 
     Ok(Server::builder()
-        .add_service(
-            FlightServiceServer::new(endpoint)
-                .max_decoding_message_size(MAX_MESSAGE_SIZE)
-                .max_encoding_message_size(MAX_MESSAGE_SIZE),
-        )
+        .add_service(endpoint.into_flight_server())
         .serve_with_incoming(incoming)
         .await?)
 }