feat: Improve fetch partition performance, support skip validation arrow ipc files

Author: westhide

ballista/client/tests/context_checks.rs

@@ -134,6 +134,42 @@ mod supported {
         Ok(())
     }
 
+    #[rstest]
+    #[case::standalone(standalone_context())]
+    #[case::remote(remote_context())]
+    #[tokio::test]
+    async fn should_execute_sql_show_config_ballista_shuffle_skip_validation(
+        #[future(awt)]
+        #[case]
+        ctx: SessionContext,
+    ) -> datafusion::error::Result<()> {
+        let state = ctx.state();
+        let ballista_config_extension =
+            state.config().options().extensions.get::<BallistaConfig>();
+
+        // ballista configuration should be registered with
+        // session state
+        assert!(ballista_config_extension.is_some());
+
+        let result = ctx
+            .sql("select name, value from information_schema.df_settings where name = 'ballista.shuffle.skip_validation'")
+            .await?
+            .collect()
+            .await?;
+
+        let expected = [
+            "+----------------------------------+-------+",
+            "| name                             | value |",
+            "+----------------------------------+-------+",
+            "| ballista.shuffle.skip_validation | true  |",
+            "+----------------------------------+-------+",
+        ];
+
+        assert_batches_eq!(expected, &result);
+
+        Ok(())
+    }
+
     #[rstest]
     #[case::standalone(standalone_context())]
     #[case::remote(remote_context())]

ballista/core/proto/ballista.proto

@@ -205,6 +205,7 @@ message FetchPartition {
   string path = 4;
   string host = 5;
   uint32 port = 6;
+  bool skip_validation = 7;
 }
 
 message PartitionLocation {
@@ -215,6 +216,7 @@ message PartitionLocation {
   ExecutorMetadata executor_meta = 3;
   PartitionStats partition_stats = 4;
   string path = 5;
+  bool skip_validation = 6;
 }
 
 // Unique identifier for a materialized partition of data
@@ -290,7 +292,7 @@ message ExecutorMetadata {
 }
 
 
-// Used for scheduler-executor 
+// Used for scheduler-executor
 // communication
 message ExecutorRegistration {
   string id = 1;
@@ -523,7 +525,7 @@ message ExecuteQueryParams {
     bytes logical_plan = 1;
     string sql = 2 [deprecated=true]; // I'd suggest to remove this, if SQL needed use `flight-sql`
   }
-  
+
   optional string session_id = 3;
   repeated KeyValuePair settings = 4;
 }

ballista/core/src/client.rs

@@ -89,6 +89,7 @@ impl BallistaClient {
         path: &str,
         host: &str,
         port: u16,
+        skip_validation: bool,
     ) -> Result<SendableRecordBatchStream> {
         let action = Action::FetchPartition {
             job_id: partition_id.job_id.clone(),
@@ -97,6 +98,7 @@ impl BallistaClient {
             path: path.to_owned(),
             host: host.to_owned(),
             port,
+            skip_validation,
         };
         self.execute_action(&action)
             .await

ballista/core/src/config.rs

@@ -34,6 +34,8 @@ pub const BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE: &str =
     "ballista.grpc_client_max_message_size";
 pub const BALLISTA_SHUFFLE_READER_MAX_REQUESTS: &str =
     "ballista.shuffle.max_concurrent_read_requests";
+pub const BALLISTA_SHUFFLE_READER_SKIP_VALIDATION: &str =
+    "ballista.shuffle.skip_validation";
 
 pub type ParseResult<T> = result::Result<T, String>;
 use std::sync::LazyLock;
@@ -54,6 +56,10 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
                          "Maximum concurrent requests shuffle reader can process".to_string(),
                          DataType::UInt64,
                          Some((64).to_string())),
+        ConfigEntry::new(BALLISTA_SHUFFLE_READER_SKIP_VALIDATION.to_string(),
+                         "Enable skip validation for Arrow IPC file when fetch partition".to_string(),
+                         DataType::Boolean,
+                         Some(true.to_string())),
     ];
     entries
         .into_iter()
@@ -175,6 +181,10 @@ impl BallistaConfig {
         self.get_usize_setting(BALLISTA_SHUFFLE_READER_MAX_REQUESTS)
     }
 
+    pub fn shuffle_reader_skip_validation(&self) -> bool {
+        self.get_bool_setting(BALLISTA_SHUFFLE_READER_SKIP_VALIDATION)
+    }
+
     fn get_usize_setting(&self, key: &str) -> usize {
         if let Some(v) = self.settings.get(key) {
             // infallible because we validate all configs in the constructor

ballista/core/src/execution_plans/distributed_query.rs

@@ -337,6 +337,7 @@ async fn fetch_partition(
     let partition_id = location.partition_id.ok_or_else(|| {
         DataFusionError::Internal("Received empty partition id".to_owned())
     })?;
+    let skip_validation = location.skip_validation;
     let host = metadata.host.as_str();
     let port = metadata.port as u16;
     let mut ballista_client = BallistaClient::try_new(host, port, max_message_size)
@@ -349,6 +350,7 @@ async fn fetch_partition(
             &location.path,
             host,
             port,
+            skip_validation,
         )
         .await
         .map_err(|e| DataFusionError::External(Box::new(e)))

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -152,12 +152,14 @@ impl ExecutionPlan for ShuffleReaderExec {
         let max_request_num =
             config.ballista_shuffle_reader_maximum_concurrent_requests();
         let max_message_size = config.ballista_grpc_client_max_message_size();
+        let skip_validation = config.ballista_shuffle_reader_skip_validation();
 
         log::debug!(
-            "ShuffleReaderExec::execute({}) max_request_num: {}, max_message_size: {}",
+            "ShuffleReaderExec::execute({}) max_request_num: {}, max_message_size: {}, skip_validation: {}",
             task_id,
             max_request_num,
-            max_message_size
+            max_message_size,
+            skip_validation
         );
         let mut partition_locations = HashMap::new();
         for p in &self.partition[partition] {
@@ -175,6 +177,9 @@ impl ExecutionPlan for ShuffleReaderExec {
             .collect();
         // Shuffle partitions for evenly send fetching partition requests to avoid hot executors within multiple tasks
         partition_locations.shuffle(&mut rng());
+        partition_locations
+            .iter_mut()
+            .for_each(|p| p.skip_validation = skip_validation);
 
         let response_receiver =
             send_fetch_partitions(partition_locations, max_request_num, max_message_size);
@@ -393,6 +398,7 @@ async fn fetch_partition_remote(
 ) -> result::Result<SendableRecordBatchStream, BallistaError> {
     let metadata = &location.executor_meta;
     let partition_id = &location.partition_id;
+    let skip_validation = location.skip_validation;
     // TODO for shuffle client connections, we should avoid creating new connections again and again.
     // And we should also avoid to keep alive too many connections for long time.
     let host = metadata.host.as_str();
@@ -411,7 +417,14 @@ async fn fetch_partition_remote(
         })?;
 
     ballista_client
-        .fetch_partition(&metadata.id, partition_id, &location.path, host, port)
+        .fetch_partition(
+            &metadata.id,
+            partition_id,
+            &location.path,
+            host,
+            port,
+            skip_validation,
+        )
         .await
 }
 
@@ -421,8 +434,9 @@ async fn fetch_partition_local(
     let path = &location.path;
     let metadata = &location.executor_meta;
     let partition_id = &location.partition_id;
+    let skip_validation = location.skip_validation;
 
-    let reader = fetch_partition_local_inner(path).map_err(|e| {
+    let reader = fetch_partition_local_inner(path, skip_validation).map_err(|e| {
         // return BallistaError::FetchFailed may let scheduler retry this task.
         BallistaError::FetchFailed(
             metadata.id.clone(),
@@ -436,14 +450,22 @@ async fn fetch_partition_local(
 
 fn fetch_partition_local_inner(
     path: &str,
+    skip_validation: bool,
 ) -> result::Result<StreamReader<BufReader<File>>, BallistaError> {
     let file = File::open(path).map_err(|e| {
         BallistaError::General(format!("Failed to open partition file at {path}: {e:?}"))
     })?;
     let file = BufReader::new(file);
-    let reader = StreamReader::try_new(file, None).map_err(|e| {
-        BallistaError::General(format!("Failed to new arrow FileReader at {path}: {e:?}"))
-    })?;
+    // Safety: The caller makes sure ipc file is valid
+    let reader = unsafe {
+        StreamReader::try_new(file, None)
+            .map_err(|e| {
+                BallistaError::General(format!(
+                    "Failed to new arrow FileReader at {path}: {e:?}"
+                ))
+            })?
+            .with_skip_validation(skip_validation)
+    };
     Ok(reader)
 }
 
@@ -567,6 +589,7 @@ mod tests {
                 },
                 partition_stats: Default::default(),
                 path: "test_path".to_string(),
+                skip_validation: true,
             })
         }
 
@@ -601,8 +624,7 @@ mod tests {
         test_send_fetch_partitions(4, 10).await;
     }
 
-    #[tokio::test]
-    async fn test_read_local_shuffle() {
+    async fn test_read_local_shuffle_impl(skip_validation: bool) {
         let session_ctx = SessionContext::new();
         let task_ctx = session_ctx.task_ctx();
         let work_dir = TempDir::new().unwrap();
@@ -629,7 +651,7 @@ mod tests {
 
         // from to input partitions test the first one with two batches
         let file_path = path.value(0);
-        let reader = fetch_partition_local_inner(file_path).unwrap();
+        let reader = fetch_partition_local_inner(file_path, skip_validation).unwrap();
 
         let mut stream: Pin<Box<dyn RecordBatchStream + Send>> =
             async { Box::pin(LocalShuffleStream::new(reader)) }.await;
@@ -645,6 +667,16 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn test_read_local_shuffle_with_validation() {
+        test_read_local_shuffle_impl(false).await;
+    }
+
+    #[tokio::test]
+    async fn test_read_local_shuffle_skip_validation() {
+        test_read_local_shuffle_impl(true).await;
+    }
+
     async fn test_send_fetch_partitions(max_request_num: usize, partition_num: usize) {
         let schema = get_test_partition_schema();
         let data_array = Int32Array::from(vec![1]);
@@ -693,6 +725,7 @@ mod tests {
                 },
                 partition_stats: Default::default(),
                 path: path.clone(),
+                skip_validation: true,
             })
             .collect()
     }

ballista/core/src/extension.rs

@@ -17,7 +17,8 @@
 
 use crate::config::{
     BallistaConfig, BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE, BALLISTA_JOB_NAME,
-    BALLISTA_SHUFFLE_READER_MAX_REQUESTS, BALLISTA_STANDALONE_PARALLELISM,
+    BALLISTA_SHUFFLE_READER_MAX_REQUESTS, BALLISTA_SHUFFLE_READER_SKIP_VALIDATION,
+    BALLISTA_STANDALONE_PARALLELISM,
 };
 use crate::planner::BallistaQueryPlanner;
 use crate::serde::protobuf::KeyValuePair;
@@ -112,6 +113,12 @@ pub trait SessionConfigExt {
         self,
         max_requests: usize,
     ) -> Self;
+
+    /// get skip_validation in flight requests for shuffle reader
+    fn ballista_shuffle_reader_skip_validation(&self) -> bool;
+
+    /// set skip_validation in flight requests for shuffle reader
+    fn with_ballista_shuffle_reader_skip_validation(self, skip_validation: bool) -> Self;
 }
 
 /// [SessionConfigHelperExt] is set of [SessionConfig] extension methods
@@ -308,6 +315,23 @@ impl SessionConfigExt for SessionConfig {
                 .set_usize(BALLISTA_SHUFFLE_READER_MAX_REQUESTS, max_requests)
         }
     }
+
+    fn ballista_shuffle_reader_skip_validation(&self) -> bool {
+        self.options()
+            .extensions
+            .get::<BallistaConfig>()
+            .map(|c| c.shuffle_reader_skip_validation())
+            .unwrap_or_else(|| BallistaConfig::default().shuffle_reader_skip_validation())
+    }
+
+    fn with_ballista_shuffle_reader_skip_validation(self, skip_validation: bool) -> Self {
+        if self.options().extensions.get::<BallistaConfig>().is_some() {
+            self.set_bool(BALLISTA_SHUFFLE_READER_SKIP_VALIDATION, skip_validation)
+        } else {
+            self.with_option_extension(BallistaConfig::default())
+                .set_bool(BALLISTA_SHUFFLE_READER_SKIP_VALIDATION, skip_validation)
+        }
+    }
 }
 
 impl SessionConfigHelperExt for SessionConfig {

ballista/core/src/serde/generated/ballista.rs

@@ -304,6 +304,8 @@ pub struct FetchPartition {
     pub host: ::prost::alloc::string::String,
     #[prost(uint32, tag = "6")]
     pub port: u32,
+    #[prost(bool, tag = "7")]
+    pub skip_validation: bool,
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct PartitionLocation {
@@ -319,6 +321,8 @@ pub struct PartitionLocation {
     pub partition_stats: ::core::option::Option<PartitionStats>,
     #[prost(string, tag = "5")]
     pub path: ::prost::alloc::string::String,
+    #[prost(bool, tag = "6")]
+    pub skip_validation: bool,
 }
 /// Unique identifier for a materialized partition of data
 #[derive(Clone, PartialEq, ::prost::Message)]

ballista/core/src/serde/scheduler/from_proto.rs

@@ -56,6 +56,7 @@ impl TryInto<Action> for protobuf::Action {
                     path: fetch.path,
                     host: fetch.host,
                     port: fetch.port as u16,
+                    skip_validation: fetch.skip_validation,
                 })
             }
             _ => Err(BallistaError::General(
@@ -126,6 +127,7 @@ impl TryInto<PartitionLocation> for protobuf::PartitionLocation {
                 })?
                 .into(),
             path: self.path,
+            skip_validation: self.skip_validation,
         })
     }
 }

ballista/core/src/serde/scheduler/mod.rs

@@ -42,6 +42,8 @@ pub enum Action {
         path: String,
         host: String,
         port: u16,
+        // Skip validation for Arrow IPC file when fetch partition
+        skip_validation: bool,
     },
 }
 
@@ -70,6 +72,7 @@ pub struct PartitionLocation {
     pub executor_meta: ExecutorMetadata,
     pub partition_stats: PartitionStats,
     pub path: String,
+    pub skip_validation: bool,
 }
 
 /// Meta-data for an executor, used when fetching shuffle partitions from other executors