feat: add config option for skipping arrow ipc read validation

Author: killzoner

ballista/client/src/extension.rs

@@ -93,7 +93,7 @@ impl SessionContextExt for SessionContext {
     async fn remote_with_state(
         url: &str,
         state: SessionState,
-    ) -> datafusion::error::Result<SessionContext> {
+    ) -> datafusion::error::Result<Self> {
         let scheduler_url = Extension::parse_url(url)?;
         log::info!(
             "Connecting to Ballista scheduler at {}",
@@ -110,7 +110,7 @@ impl SessionContextExt for SessionContext {
         Ok(SessionContext::new_with_state(session_state))
     }
 
-    async fn remote(url: &str) -> datafusion::error::Result<SessionContext> {
+    async fn remote(url: &str) -> datafusion::error::Result<Self> {
         let scheduler_url = Extension::parse_url(url)?;
         log::info!(
             "Connecting to Ballista scheduler at: {}",
@@ -129,7 +129,9 @@ impl SessionContextExt for SessionContext {
     #[cfg(feature = "standalone")]
     async fn standalone_with_state(
         state: SessionState,
-    ) -> datafusion::error::Result<SessionContext> {
+    ) -> datafusion::error::Result<Self> {
+        log::info!("Running in local mode. Scheduler will be run in-proc");
+
         let scheduler_url = Extension::setup_standalone(Some(&state)).await?;
 
         let session_state = state.upgrade_for_ballista(scheduler_url)?;

ballista/core/Cargo.toml

@@ -35,6 +35,8 @@ exclude = ["*.proto"]
 rustc-args = ["--cfg", "docsrs"]
 
 [features]
+default = ["arrow-ipc-optimizations"]
+arrow-ipc-optimizations = []
 build-binary = ["aws-config", "aws-credential-types", "clap", "object_store"]
 docsrs = []
 # Used for testing ONLY: causes all values to hash to the same value (test for collisions)

ballista/core/src/config.rs

@@ -31,32 +31,37 @@ use datafusion::{
 pub const BALLISTA_JOB_NAME: &str = "ballista.job.name";
 /// Configuration key for standalone processing parallelism.
 pub const BALLISTA_STANDALONE_PARALLELISM: &str = "ballista.standalone.parallelism";
+
+// gRPC configuration
+
+/// Configuration key for gRPC client connection timeout in seconds.
+pub const BALLISTA_GRPC_CLIENT_CONNECT_TIMEOUT_SECONDS: &str =
+    "ballista.grpc.client.connect_timeout_seconds";
+/// Configuration key for HTTP/2 keep-alive interval for gRPC clients in seconds.
+pub const BALLISTA_GRPC_CLIENT_HTTP2_KEEPALIVE_INTERVAL_SECONDS: &str =
+    "ballista.grpc.client.http2_keepalive_interval_seconds";
 /// max message size for gRPC clients
 pub const BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE: &str =
     "ballista.grpc_client_max_message_size";
-/// Configuration key for maximum concurrent shuffle read requests.
-pub const BALLISTA_SHUFFLE_READER_MAX_REQUESTS: &str =
-    "ballista.shuffle.max_concurrent_read_requests";
+/// Configuration key for TCP keep-alive interval for gRPC clients in seconds.
+pub const BALLISTA_GRPC_CLIENT_TCP_KEEPALIVE_SECONDS: &str =
+    "ballista.grpc.client.tcp_keepalive_seconds";
+/// Configuration key for gRPC client request timeout in seconds.
+pub const BALLISTA_GRPC_CLIENT_TIMEOUT_SECONDS: &str =
+    "ballista.grpc.client.timeout_seconds";
+
+// Shuffle reader configuration
+
 /// Configuration key to force remote reads even for local partitions.
 pub const BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ: &str =
     "ballista.shuffle.force_remote_read";
+/// Configuration key for maximum concurrent shuffle read requests.
+pub const BALLISTA_SHUFFLE_READER_MAX_REQUESTS: &str =
+    "ballista.shuffle.max_concurrent_read_requests";
 /// Configuration key to prefer Flight protocol for remote shuffle reads.
 pub const BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT: &str =
     "ballista.shuffle.remote_read_prefer_flight";
 
-/// Configuration key for gRPC client connection timeout in seconds.
-pub const BALLISTA_GRPC_CLIENT_CONNECT_TIMEOUT_SECONDS: &str =
-    "ballista.grpc.client.connect_timeout_seconds";
-/// Configuration key for gRPC client request timeout in seconds.
-pub const BALLISTA_GRPC_CLIENT_TIMEOUT_SECONDS: &str =
-    "ballista.grpc.client.timeout_seconds";
-/// Configuration key for TCP keep-alive interval for gRPC clients in seconds.
-pub const BALLISTA_GRPC_CLIENT_TCP_KEEPALIVE_SECONDS: &str =
-    "ballista.grpc.client.tcp_keepalive_seconds";
-/// Configuration key for HTTP/2 keep-alive interval for gRPC clients in seconds.
-pub const BALLISTA_GRPC_CLIENT_HTTP2_KEEPALIVE_INTERVAL_SECONDS: &str =
-    "ballista.grpc.client.http2_keepalive_interval_seconds";
-
 /// Result type for configuration parsing operations.
 pub type ParseResult<T> = result::Result<T, String>;
 use std::sync::LazyLock;
@@ -69,38 +74,43 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
         ConfigEntry::new(BALLISTA_STANDALONE_PARALLELISM.to_string(),
                          "Standalone processing parallelism ".to_string(),
                          DataType::UInt16, Some(std::thread::available_parallelism().map(|v| v.get()).unwrap_or(1).to_string())),
+
+        // gRPC configuration
+        ConfigEntry::new(BALLISTA_GRPC_CLIENT_CONNECT_TIMEOUT_SECONDS.to_string(),
+                         "Connection timeout for gRPC client in seconds".to_string(),
+                         DataType::UInt64,
+                         Some((20).to_string())),
+        ConfigEntry::new(BALLISTA_GRPC_CLIENT_HTTP2_KEEPALIVE_INTERVAL_SECONDS.to_string(),
+                         "HTTP/2 keep-alive interval for gRPC client in seconds".to_string(),
+                         DataType::UInt64,
+                         Some((300).to_string())),
         ConfigEntry::new(BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE.to_string(),
                          "Configuration for max message size in gRPC clients".to_string(),
                          DataType::UInt64,
                          Some((16 * 1024 * 1024).to_string())),
-        ConfigEntry::new(BALLISTA_SHUFFLE_READER_MAX_REQUESTS.to_string(),
-                         "Maximum concurrent requests shuffle reader can process".to_string(),
+        ConfigEntry::new(BALLISTA_GRPC_CLIENT_TCP_KEEPALIVE_SECONDS.to_string(),
+                         "TCP keep-alive interval for gRPC client in seconds".to_string(),
                          DataType::UInt64,
-                         Some((64).to_string())),
+                         Some((3600).to_string())),
+        ConfigEntry::new(BALLISTA_GRPC_CLIENT_TIMEOUT_SECONDS.to_string(),
+                         "Request timeout for gRPC client in seconds".to_string(),
+                         DataType::UInt64,
+                         Some((20).to_string())),
+
+
+        // Shuffle reader configuration
         ConfigEntry::new(BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ.to_string(),
                          "Forces the shuffle reader to always read partitions via the Arrow Flight client, even when partitions are local to the node.".to_string(),
                          DataType::Boolean,
                          Some((false).to_string())),
+        ConfigEntry::new(BALLISTA_SHUFFLE_READER_MAX_REQUESTS.to_string(),
+                         "Maximum concurrent requests shuffle reader can process".to_string(),
+                         DataType::UInt64,
+                         Some((64).to_string())),
         ConfigEntry::new(BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT.to_string(),
                          "Forces the shuffle reader to use flight reader instead of block reader for remote read. Block reader usually has better performance and resource utilization".to_string(),
                          DataType::Boolean,
                          Some((false).to_string())),
-        ConfigEntry::new(BALLISTA_GRPC_CLIENT_CONNECT_TIMEOUT_SECONDS.to_string(),
-                         "Connection timeout for gRPC client in seconds".to_string(),
-                         DataType::UInt64,
-                         Some((20).to_string())),
-        ConfigEntry::new(BALLISTA_GRPC_CLIENT_TIMEOUT_SECONDS.to_string(),
-                         "Request timeout for gRPC client in seconds".to_string(),
-                         DataType::UInt64,
-                         Some((20).to_string())),
-        ConfigEntry::new(BALLISTA_GRPC_CLIENT_TCP_KEEPALIVE_SECONDS.to_string(),
-                         "TCP keep-alive interval for gRPC client in seconds".to_string(),
-                         DataType::UInt64,
-                         Some((3600).to_string())),
-        ConfigEntry::new(BALLISTA_GRPC_CLIENT_HTTP2_KEEPALIVE_INTERVAL_SECONDS.to_string(),
-                         "HTTP/2 keep-alive interval for gRPC client in seconds".to_string(),
-                         DataType::UInt64,
-                         Some((300).to_string()))
     ];
     entries
         .into_iter()
@@ -212,6 +222,19 @@ impl BallistaConfig {
         &self.settings
     }
 
+    /// Allows skipping redundant validation of arrow IPC valid data.
+    pub fn arrow_ipc_reader_skip_validation(&self) -> bool {
+        #[cfg(feature = "arrow-ipc-optimizations")]
+        {
+            true
+        }
+
+        #[cfg(not(feature = "arrow-ipc-optimizations"))]
+        {
+            false
+        }
+    }
+
     /// Returns the maximum message size for gRPC clients in bytes.
     pub fn default_grpc_client_max_message_size(&self) -> usize {
         self.get_usize_setting(BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE)

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -162,6 +162,8 @@ impl ExecutionPlan for ShuffleReaderExec {
         let max_message_size = config.ballista_grpc_client_max_message_size();
         let force_remote_read = config.ballista_shuffle_reader_force_remote_read();
         let prefer_flight = config.ballista_shuffle_reader_remote_prefer_flight();
+        let arrow_ipc_reader_skip_validation =
+            config.ballista_arrow_ipc_reader_skip_validation();
 
         if force_remote_read {
             debug!(
@@ -195,6 +197,7 @@ impl ExecutionPlan for ShuffleReaderExec {
             max_message_size,
             force_remote_read,
             prefer_flight,
+            arrow_ipc_reader_skip_validation,
         );
 
         let result = RecordBatchStreamAdapter::new(
@@ -390,6 +393,7 @@ fn send_fetch_partitions(
     max_message_size: usize,
     force_remote_read: bool,
     flight_transport: bool,
+    arrow_ipc_reader_skip_validation: bool,
 ) -> AbortableReceiverStream {
     let (response_sender, response_receiver) = mpsc::channel(max_request_num);
     let semaphore = Arc::new(Semaphore::new(max_request_num));
@@ -409,7 +413,12 @@ fn send_fetch_partitions(
     spawned_tasks.push(SpawnedTask::spawn(async move {
         for p in local_locations {
             let r = PartitionReaderEnum::Local
-                .fetch_partition(&p, max_message_size, flight_transport)
+                .fetch_partition(
+                    &p,
+                    max_message_size,
+                    flight_transport,
+                    arrow_ipc_reader_skip_validation,
+                )
                 .await;
             if let Err(e) = response_sender_c.send(r).await {
                 error!("Fail to send response event to the channel due to {e}");
@@ -424,7 +433,12 @@ fn send_fetch_partitions(
             // Block if exceeds max request number.
             let permit = semaphore.acquire_owned().await.unwrap();
             let r = PartitionReaderEnum::FlightRemote
-                .fetch_partition(&p, max_message_size, flight_transport)
+                .fetch_partition(
+                    &p,
+                    max_message_size,
+                    flight_transport,
+                    arrow_ipc_reader_skip_validation,
+                )
                 .await;
             // Block if the channel buffer is full.
             if let Err(e) = response_sender.send(r).await {
@@ -451,6 +465,7 @@ trait PartitionReader: Send + Sync + Clone {
         location: &PartitionLocation,
         max_message_size: usize,
         flight_transport: bool,
+        arrow_ipc_reader_skip_validation: bool,
     ) -> result::Result<SendableRecordBatchStream, BallistaError>;
 }
 
@@ -470,12 +485,15 @@ impl PartitionReader for PartitionReaderEnum {
         location: &PartitionLocation,
         max_message_size: usize,
         flight_transport: bool,
+        arrow_ipc_reader_skip_validation: bool,
     ) -> result::Result<SendableRecordBatchStream, BallistaError> {
         match self {
             PartitionReaderEnum::FlightRemote => {
                 fetch_partition_remote(location, max_message_size, flight_transport).await
             }
-            PartitionReaderEnum::Local => fetch_partition_local(location).await,
+            PartitionReaderEnum::Local => {
+                fetch_partition_local(location, arrow_ipc_reader_skip_validation).await
+            }
             PartitionReaderEnum::ObjectStoreRemote => {
                 fetch_partition_object_store(location).await
             }
@@ -521,33 +539,44 @@ async fn fetch_partition_remote(
 
 async fn fetch_partition_local(
     location: &PartitionLocation,
+    arrow_ipc_reader_skip_validation: bool,
 ) -> result::Result<SendableRecordBatchStream, BallistaError> {
     let path = &location.path;
     let metadata = &location.executor_meta;
     let partition_id = &location.partition_id;
 
-    let reader = fetch_partition_local_inner(path).map_err(|e| {
-        // return BallistaError::FetchFailed may let scheduler retry this task.
-        BallistaError::FetchFailed(
-            metadata.id.clone(),
-            partition_id.stage_id,
-            partition_id.partition_id,
-            e.to_string(),
-        )
-    })?;
+    let reader = fetch_partition_local_inner(path, arrow_ipc_reader_skip_validation)
+        .map_err(|e| {
+            // return BallistaError::FetchFailed may let scheduler retry this task.
+            BallistaError::FetchFailed(
+                metadata.id.clone(),
+                partition_id.stage_id,
+                partition_id.partition_id,
+                e.to_string(),
+            )
+        })?;
     Ok(Box::pin(LocalShuffleStream::new(reader)))
 }
 
 fn fetch_partition_local_inner(
     path: &str,
+    arrow_ipc_reader_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: setting `skip_validation` requires `unsafe`, user assures data is valid
+    let reader = unsafe {
+        StreamReader::try_new(file, None)
+            .map_err(|e| {
+                BallistaError::General(format!(
+                    "Failed to create new arrow StreamReader at {path}: {e:?}"
+                ))
+            })?
+            .with_skip_validation(arrow_ipc_reader_skip_validation)
+    };
+
     Ok(reader)
 }
 
@@ -881,7 +910,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, true).unwrap();
 
         let mut stream: Pin<Box<dyn RecordBatchStream + Send>> =
             async { Box::pin(LocalShuffleStream::new(reader)) }.await;
@@ -951,6 +980,7 @@ mod tests {
             4 * 1024 * 1024,
             false,
             true,
+            true,
         );
 
         let stream = RecordBatchStreamAdapter::new(