null fix after copying data from array to duckdb chunk

duckdb · May 24, 2024 · 86b6b44 · 86b6b44
1 parent a1aa55a
commit 86b6b44
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/src/vtab/arrow.rs b/src/vtab/arrow.rs
@@ -257,6 +257,13 @@ pub fn record_batch_to_duckdb_data_chunk(
 fn primitive_array_to_flat_vector<T: ArrowPrimitiveType>(array: &PrimitiveArray<T>, out_vector: &mut FlatVector) {
     // assert!(array.len() <= out_vector.capacity());
     out_vector.copy::<T::Native>(array.values());
+    if let Some(nulls) = array.nulls() {
+        for (i, null) in nulls.into_iter().enumerate() {
+            if !null {
+                out_vector.set_null(i);
+            }
+        }
+    }
 }
 
 fn primitive_array_to_flat_vector_cast<T: ArrowPrimitiveType>(
@@ -267,6 +274,13 @@ fn primitive_array_to_flat_vector_cast<T: ArrowPrimitiveType>(
     let array = arrow::compute::kernels::cast::cast(array, &data_type).unwrap();
     let out_vector: &mut FlatVector = out_vector.as_mut_any().downcast_mut().unwrap();
     out_vector.copy::<T::Native>(array.as_primitive::<T>().values());
+    if let Some(nulls) = array.nulls() {
+        for (i, null) in nulls.iter().enumerate() {
+            if !null {
+                out_vector.set_null(i);
+            }
+        }
+    }
 }
 
 fn primitive_array_to_vector(array: &dyn Array, out: &mut dyn Vector) -> Result<(), Box<dyn std::error::Error>> {