format

Author: andygrove

native/core/src/execution/jni_api.rs

@@ -40,7 +40,6 @@ use datafusion::{
     prelude::{SessionConfig, SessionContext},
 };
 use datafusion_comet_proto::spark_operator::{self, Operator};
-use prost::Message;
 use datafusion_spark::function::bitwise::bit_count::SparkBitCount;
 use datafusion_spark::function::bitwise::bit_get::SparkBitGet;
 use datafusion_spark::function::bitwise::bitwise_not::SparkBitwiseNot;
@@ -73,6 +72,7 @@ use jni::{
     sys::{jboolean, jdouble, jint, jlong},
     JNIEnv,
 };
+use prost::Message;
 use std::collections::HashMap;
 use std::path::PathBuf;
 use std::time::{Duration, Instant};
@@ -894,10 +894,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlock(
     try_unwrap_or_throw(&e, |mut env| {
         // Parse the schema from protobuf bytes
         let schema_vec = env.convert_byte_array(&schema_bytes)?;
-        let shuffle_scan =
-            spark_operator::ShuffleScan::decode(schema_vec.as_slice()).map_err(|e| {
-                CometError::Internal(format!("Failed to parse shuffle schema: {e}"))
-            })?;
+        let shuffle_scan = spark_operator::ShuffleScan::decode(schema_vec.as_slice())
+            .map_err(|e| CometError::Internal(format!("Failed to parse shuffle schema: {e}")))?;
         let fields: Vec<Field> = shuffle_scan
             .fields
             .iter()

native/core/src/execution/shuffle/codec.rs

@@ -71,10 +71,7 @@ fn write_array_data<W: Write>(data: &arrow::array::ArrayData, writer: &mut W) ->
     writer.write_all(&num_buffers.to_le_bytes())?;
     for buffer in data.buffers() {
         let len: u32 = buffer.len().try_into().map_err(|_| {
-            DataFusionError::Execution(format!(
-                "Buffer length {} exceeds u32::MAX",
-                buffer.len()
-            ))
+            DataFusionError::Execution(format!("Buffer length {} exceeds u32::MAX", buffer.len()))
         })?;
         writer.write_all(&len.to_le_bytes())?;
         writer.write_all(buffer.as_slice())?;
@@ -100,9 +97,7 @@ fn write_raw_batch<W: Write>(batch: &RecordBatch, writer: &mut W) -> Result<()>
     for col in batch.columns() {
         // Cast dictionary arrays to their value type
         let col = match col.data_type() {
-            DataType::Dictionary(_, value_type) => {
-                cast(col.as_ref(), value_type.as_ref())?
-            }
+            DataType::Dictionary(_, value_type) => cast(col.as_ref(), value_type.as_ref())?,
             _ => Arc::clone(col),
         };
         write_array_data(&col.to_data(), writer)?;
@@ -237,9 +232,7 @@ fn read_bytes<'a>(cursor: &mut &'a [u8], len: usize) -> Result<&'a [u8]> {
 /// Returns child data types for nested Arrow types.
 fn get_child_types(data_type: &DataType) -> Vec<DataType> {
     match data_type {
-        DataType::List(field)
-        | DataType::LargeList(field)
-        | DataType::FixedSizeList(field, _) => {
+        DataType::List(field) | DataType::LargeList(field) | DataType::FixedSizeList(field, _) => {
             vec![field.data_type().clone()]
         }
         DataType::Map(field, _) => {
@@ -709,11 +702,7 @@ mod tests {
 
         let batch = RecordBatch::try_new(
             Arc::clone(&schema),
-            vec![
-                Arc::new(list_arr),
-                Arc::new(struct_arr),
-                Arc::new(map_arr),
-            ],
+            vec![Arc::new(list_arr), Arc::new(struct_arr), Arc::new(map_arr)],
         )
         .unwrap();
 
@@ -732,8 +721,7 @@ mod tests {
 
         let keys = Int32Array::from(vec![Some(0), Some(1), None, Some(0)]);
         let values = StringArray::from(vec!["foo", "bar"]);
-        let dict_arr =
-            DictionaryArray::<Int32Type>::try_new(keys, Arc::new(values)).unwrap();
+        let dict_arr = DictionaryArray::<Int32Type>::try_new(keys, Arc::new(values)).unwrap();
 
         let batch =
             RecordBatch::try_new(Arc::clone(&dict_schema), vec![Arc::new(dict_arr)]).unwrap();

native/core/src/execution/shuffle/shuffle_writer.rs

@@ -296,8 +296,7 @@ mod test {
         ] {
             let mut output = vec![];
             let mut cursor = Cursor::new(&mut output);
-            let writer =
-                ShuffleBlockWriter::try_new(schema.as_ref(), codec.clone()).unwrap();
+            let writer = ShuffleBlockWriter::try_new(schema.as_ref(), codec.clone()).unwrap();
             let length = writer
                 .write_batch(&batch, &mut cursor, &Time::default())
                 .unwrap();