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

Author: westhide

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -441,9 +441,16 @@ fn fetch_partition_local_inner(
         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 ShuffleWriter partition 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(true)
+    };
     Ok(reader)
 }
 

ballista/executor/src/flight_service.rs

@@ -95,9 +95,12 @@ impl FlightService for BallistaFlightService {
                     })
                     .map_err(|e| from_ballista_err(&e))?;
                 let file = BufReader::new(file);
-                let reader =
-                    StreamReader::try_new(file, None).map_err(|e| from_arrow_err(&e))?;
-
+                // Safety: The ShuffleWriter partition file is valid
+                let reader = unsafe {
+                    StreamReader::try_new(file, None)
+                        .map_err(|e| from_arrow_err(&e))?
+                        .with_skip_validation(true)
+                };
                 let (tx, rx) = channel(2);
                 let schema = reader.schema();
                 task::spawn_blocking(move || {