format

duckdb · Apr 24, 2024 · 092a2a9 · 092a2a9
1 parent 2652a54
commit 092a2a9
Showing 1 changed file with 31 additions and 16 deletions.
diff --git a/src/vtab/arrow.rs b/src/vtab/arrow.rs
@@ -555,8 +555,14 @@ mod test {
     use crate::{Connection, Result};
     use arrow::{
         array::{
-            Array, ArrayRef, AsArray, Date32Array, Date64Array, Decimal256Array, Float64Array, GenericListArray, Int32Array, ListArray, OffsetSizeTrait, PrimitiveArray, StringArray, StructArray, Time32SecondArray, Time64MicrosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray
-        }, buffer::{OffsetBuffer, ScalarBuffer}, datatypes::{i256, ArrowNativeType, ArrowPrimitiveType, DataType, Field, Fields, Schema}, record_batch::RecordBatch
+            Array, ArrayRef, AsArray, Date32Array, Date64Array, Decimal256Array, Float64Array, GenericListArray,
+            Int32Array, ListArray, OffsetSizeTrait, PrimitiveArray, StringArray, StructArray, Time32SecondArray,
+            Time64MicrosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
+            TimestampSecondArray,
+        },
+        buffer::{OffsetBuffer, ScalarBuffer},
+        datatypes::{i256, ArrowNativeType, ArrowPrimitiveType, DataType, Field, Fields, Schema},
+        record_batch::RecordBatch,
     };
     use std::{error::Error, sync::Arc};
 
@@ -684,26 +690,28 @@ mod test {
         Ok(())
     }
 
-    fn check_generic_array_roundtrip<T>(
-        arry: GenericListArray<T>,
-    ) -> Result<(), Box<dyn Error>> where T: OffsetSizeTrait{
-
-        let expected_output_array = arry.clone();    
-    
+    fn check_generic_array_roundtrip<T>(arry: GenericListArray<T>) -> Result<(), Box<dyn Error>>
+    where
+        T: OffsetSizeTrait,
+    {
+        let expected_output_array = arry.clone();
+
         let db = Connection::open_in_memory()?;
         db.register_table_function::<ArrowVTab>("arrow")?;
-    
+
         // Roundtrip a record batch from Rust to DuckDB and back to Rust
         let schema = Schema::new(vec![Field::new("a", arry.data_type().clone(), false)]);
-    
+
         let rb = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(arry.clone())])?;
         let param = arrow_recordbatch_to_query_params(rb);
         let mut stmt = db.prepare("select a from arrow(?, ?)")?;
         let rb = stmt.query_arrow(param)?.next().expect("no record batch");
-    
+
         let output_any_array = rb.column(0);
-        assert!(output_any_array.data_type().equals_datatype(expected_output_array.data_type()));
-
+        assert!(output_any_array
+            .data_type()
+            .equals_datatype(expected_output_array.data_type()));
+
         match output_any_array.as_list_opt::<T>() {
             Some(output_array) => {
                 assert_eq!(output_array.len(), expected_output_array.len());
@@ -716,16 +724,23 @@ mod test {
             }
             None => panic!("Expected GenericListArray"),
         }
-    
+
         Ok(())
     }
 
     #[test]
     fn test_array_roundtrip() -> Result<(), Box<dyn Error>> {
         check_generic_array_roundtrip(ListArray::new(
             Arc::new(Field::new("item", DataType::Utf8, true)),
-            OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 4, 5])), 
-            Arc::new(StringArray::from(vec![Some("foo"), Some("baz"), Some("bar"), Some("foo"), Some("baz")])), None
+            OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 4, 5])),
+            Arc::new(StringArray::from(vec![
+                Some("foo"),
+                Some("baz"),
+                Some("bar"),
+                Some("foo"),
+                Some("baz"),
+            ])),
+            None,
         ))?;
 
         Ok(())