actually use

Author: adriangb

src/flight_service/do_get.rs

@@ -134,12 +134,18 @@ impl ArrowFlightEndpoint {
             .execute(doget.target_partition as usize, session_state.task_ctx())
             .map_err(|err| Status::internal(format!("Error executing stage plan: {err:#?}")))?;
 
-        let stream = FlightDataEncoderBuilder::new()
-            .with_schema(stream.schema().clone())
-            .with_max_flight_data_size(usize::MAX)
-            .build(stream.map_err(|err| {
-                FlightError::Tonic(Box::new(datafusion_error_to_tonic_status(&err)))
-            }));
+        let mut encoder_builder =
+            FlightDataEncoderBuilder::new().with_schema(stream.schema().clone());
+        if let Some(max_message_size) = self.max_message_size {
+            encoder_builder = encoder_builder.with_max_flight_data_size(max_message_size);
+        }
+
+        let stream =
+            encoder_builder
+                .with_max_flight_data_size(usize::MAX)
+                .build(stream.map_err(|err| {
+                    FlightError::Tonic(Box::new(datafusion_error_to_tonic_status(&err)))
+                }));
 
         let task_data_entries = Arc::clone(&self.task_data_entries);
         let num_partitions_remaining = Arc::clone(&stage_data.num_partitions_remaining);

src/flight_service/service.rs

@@ -28,13 +28,9 @@ pub struct ArrowFlightEndpoint {
     pub(super) task_data_entries: Arc<TTLMap<StageKey, Arc<OnceCell<TaskData>>>>,
     pub(super) session_builder: Arc<dyn DistributedSessionBuilder + Send + Sync>,
     pub(super) hooks: ArrowFlightEndpointHooks,
-    max_message_size: usize,
+    pub(super) max_message_size: Option<usize>,
 }
 
-/// Default maximum message size for FlightData chunks in ArrowFlightEndpoint.
-/// This is the size used for chunking FlightData within the endpoint.
-const DEFAULT_MESSAGE_SIZE: usize = 2 * 1024 * 1024; // 2 MB
-
 impl ArrowFlightEndpoint {
     pub fn try_new(
         session_builder: impl DistributedSessionBuilder + Send + Sync + 'static,
@@ -45,7 +41,7 @@ impl ArrowFlightEndpoint {
             task_data_entries: Arc::new(ttl_map),
             session_builder: Arc::new(session_builder),
             hooks: ArrowFlightEndpointHooks::default(),
-            max_message_size: DEFAULT_MESSAGE_SIZE,
+            max_message_size: None,
         })
     }
 
@@ -62,14 +58,19 @@ impl ArrowFlightEndpoint {
     }
 
     /// Set the maximum message size for FlightData chunks.
-    /// Defaults to 2 MB.
+    ///
+    /// Defaults to None, which uses `arrow-rs` default, curerntly 2MB.
+    /// 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.
+    ///
+    /// [`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 = size;
+        self.max_message_size = Some(size);
         self
     }
 }