feat: add config option for skipping arrow ipc read validation (#1374)

* feat: add config option for skipping arrow ipc read validation

* fix: remove unused dependencies

Author: killzoner

Cargo.toml

@@ -33,8 +33,6 @@ rust-version = "1.88.0"
 arrow = { version = "57", features = ["ipc_compression"] }
 arrow-flight = { version = "57", features = ["flight-sql-experimental"] }
 clap = { version = "4.5", features = ["derive", "cargo"] }
-configure_me = { version = "0.4.0" }
-configure_me_codegen = { version = "0.4.4" }
 datafusion = "51.0.0"
 datafusion-cli = "51.0.0"
 datafusion-proto = "51.0.0"

ballista/core/Cargo.toml

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

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -545,9 +545,17 @@ 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: 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(cfg!(feature = "arrow-ipc-optimizations"))
+    };
+
     Ok(reader)
 }
 

ballista/executor/Cargo.toml

@@ -33,8 +33,9 @@ path = "src/bin/main.rs"
 required-features = ["build-binary"]
 
 [features]
+arrow-ipc-optimizations = []
 build-binary = ["clap", "tracing-subscriber", "tracing-appender", "tracing", "ballista-core/build-binary"]
-default = ["build-binary", "mimalloc"]
+default = ["arrow-ipc-optimizations", "build-binary", "mimalloc"]
 
 [dependencies]
 arrow = { workspace = true }

ballista/executor/src/flight_service.rs

@@ -105,8 +105,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: setting `skip_validation` requires `unsafe`, user assures data is valid
+                let reader = unsafe {
+                    StreamReader::try_new(file, None)
+                        .map_err(|e| from_arrow_err(&e))?
+                        .with_skip_validation(cfg!(feature = "arrow-ipc-optimizations"))
+                };
 
                 let (tx, rx) = channel(2);
                 let schema = reader.schema();