test: replace parquet files with generated tests

connector_arrow/Cargo.toml

@@ -57,9 +57,7 @@ optional = true
 
 [dev-dependencies]
 env_logger = "0.11"
-arrow = { version = "49", features = ["prettyprint"], default-features = false }
-parquet = { version = "49", features = ["arrow"], default-features = false }
-insta = { version = "1.34.0" }
+arrow = { version = "49", default-features = false }
 similar-asserts = { version = "1.5.0" }
 half = "2.3.1"
 rand = { version = "0.8.5", default-features = false }

connector_arrow/src/util/coerce.rs

@@ -6,23 +6,32 @@ use arrow::record_batch::RecordBatch;
 use itertools::Itertools;
 
 pub fn coerce_batches<F>(
+    schema: SchemaRef,
     batches: &[RecordBatch],
     coerce_fn: F,
-) -> Result<Vec<RecordBatch>, arrow::error::ArrowError>
+    override_nullable: Option<bool>,
+) -> Result<(SchemaRef, Vec<RecordBatch>), arrow::error::ArrowError>
 where
     F: Fn(&DataType) -> Option<DataType> + Copy,
 {
-    batches.iter().map(|b| coerce_batch(b, coerce_fn)).collect()
+    let batches = batches
+        .iter()
+        .map(|b| coerce_batch(b, coerce_fn, override_nullable))
+        .collect::<Result<Vec<_>, _>>()?;
+
+    let schema = coerce_schema(schema, coerce_fn, override_nullable);
+    Ok((schema, batches))
 }
 
 pub fn coerce_batch<F>(
     batch: &RecordBatch,
     coerce_fn: F,
+    override_nullable: Option<bool>,
 ) -> Result<RecordBatch, arrow::error::ArrowError>
 where
     F: Fn(&DataType) -> Option<DataType> + Copy,
 {
-    let new_schema = coerce_schema(batch.schema(), coerce_fn);
+    let new_schema = coerce_schema(batch.schema(), coerce_fn, override_nullable);
 
     let new_columns = batch
         .columns()
@@ -47,17 +56,31 @@ where
     }
 }
 
-pub fn coerce_schema<F>(schema: SchemaRef, coerce_fn: F) -> SchemaRef
+pub fn coerce_schema<F>(
+    schema: SchemaRef,
+    coerce_fn: F,
+    override_nullable: Option<bool>,
+) -> SchemaRef
 where
     F: Fn(&DataType) -> Option<DataType> + Copy,
 {
     Arc::new(Schema::new(
         schema
             .fields()
             .iter()
-            .map(|f| match coerce_fn(f.data_type()) {
-                Some(new_ty) => Field::new(f.name(), new_ty, true),
-                None => Field::clone(f).with_nullable(true),
+            .map(|f| {
+                let field = match coerce_fn(f.data_type()) {
+                    Some(new_ty) => {
+                        let nullable = f.is_nullable() || matches!(f.data_type(), DataType::Null);
+                        Field::new(f.name(), new_ty, nullable)
+                    }
+                    None => Field::clone(f),
+                };
+                if let Some(nullable) = &override_nullable {
+                    field.with_nullable(*nullable)
+                } else {
+                    field
+                }
             })
             .collect_vec(),
     ))