fix[duckdb]: remove validity mask from vector export (#5739)

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

vortex-duckdb/src/exporter/dict.rs

@@ -14,6 +14,8 @@ use vortex::array::arrays::ConstantArray;
 use vortex::array::arrays::ConstantVTable;
 use vortex::array::arrays::DictArray;
 use vortex::array::arrays::PrimitiveArray;
+use vortex::array::builtins::ArrayBuiltins;
+use vortex::array::mask::MaskExecutor;
 use vortex::array::vectors::VectorIntoArray;
 use vortex::compute;
 use vortex::compute2::take::Take;
@@ -186,7 +188,7 @@ pub(crate) fn new_vector_exporter_with_flatten(
     if let Some(constant) = values.as_opt::<ConstantVTable>() {
         return constant::new_exporter_with_mask(
             &ConstantArray::new(constant.scalar().clone(), array.codes().len()),
-            array.codes().validity_mask(),
+            array.codes().is_null()?.not()?.execute_mask(session)?,
             cache,
         );
     }
@@ -299,19 +301,22 @@ impl<I: IntegerPType + AsPrimitive<u32>> ColumnExporter for DictVectorExporter<I
 
 #[cfg(test)]
 mod tests {
+    use vortex::VortexSessionDefault;
     use vortex::array::IntoArray;
     use vortex::array::arrays::ConstantArray;
     use vortex::array::arrays::DictArray;
     use vortex::array::arrays::PrimitiveArray;
     use vortex::buffer::Buffer;
     use vortex::error::VortexResult;
+    use vortex::session::VortexSession;
 
     use crate::cpp;
     use crate::duckdb::DataChunk;
     use crate::duckdb::LogicalType;
     use crate::exporter::ColumnExporter;
     use crate::exporter::ConversionCache;
     use crate::exporter::dict::new_exporter_with_flatten;
+    use crate::exporter::dict::new_vector_exporter_with_flatten;
     use crate::exporter::new_array_exporter;
 
     pub(crate) fn new_exporter(
@@ -344,6 +349,54 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_constant_dict_vector() {
+        let arr = DictArray::new(
+            PrimitiveArray::from_option_iter([None, Some(0u32)]).into_array(),
+            ConstantArray::new(10, 1).into_array(),
+        );
+
+        let mut chunk = DataChunk::new([LogicalType::new(cpp::duckdb_type::DUCKDB_TYPE_INTEGER)]);
+
+        let session = VortexSession::default();
+        new_vector_exporter_with_flatten(&arr, &ConversionCache::default(), &session, false)
+            .unwrap()
+            .export(0, 2, &mut chunk.get_vector(0))
+            .unwrap();
+        chunk.set_len(2);
+
+        assert_eq!(
+            format!("{}", String::try_from(&chunk).unwrap()),
+            r#"Chunk - [1 Columns]
+- FLAT INTEGER: 2 = [ NULL, 10]
+"#
+        );
+    }
+
+    #[test]
+    fn test_constant_dict_vector_null() {
+        let arr = DictArray::new(
+            PrimitiveArray::from_option_iter([None::<u32>, None]).into_array(),
+            ConstantArray::new(10, 1).into_array(),
+        );
+
+        let mut chunk = DataChunk::new([LogicalType::new(cpp::duckdb_type::DUCKDB_TYPE_INTEGER)]);
+
+        let session = VortexSession::default();
+        new_vector_exporter_with_flatten(&arr, &ConversionCache::default(), &session, false)
+            .unwrap()
+            .export(0, 2, &mut chunk.get_vector(0))
+            .unwrap();
+        chunk.set_len(2);
+
+        assert_eq!(
+            format!("{}", String::try_from(&chunk).unwrap()),
+            r#"Chunk - [1 Columns]
+- CONSTANT INTEGER: 2 = [ NULL]
+"#
+        );
+    }
+
     #[test]
     fn test_nullable_dict() {
         let arr = DictArray::new(

vortex-duckdb/src/exporter/fixed_size_list.rs

@@ -11,6 +11,8 @@
 use vortex::array::ArrayRef;
 use vortex::array::ToCanonical;
 use vortex::array::arrays::FixedSizeListArray;
+use vortex::array::validity::Validity;
+use vortex::array::vtable::ValidityHelper;
 use vortex::error::VortexResult;
 use vortex::mask::Mask;
 use vortex::session::VortexSession;
@@ -103,14 +105,14 @@ pub(crate) fn new_vector_exporter(
 
     let ltype: LogicalType = array.dtype().try_into()?;
 
-    let mask = array.validity_mask();
+    let mask = array.validity();
 
-    if let Mask::AllFalse(len) = mask {
-        return Ok(all_invalid::new_exporter(len, &ltype));
+    if matches!(mask, Validity::AllInvalid) {
+        return Ok(all_invalid::new_exporter(array.len(), &ltype));
     }
 
     Ok(Box::new(FixedSizeListExporter {
-        validity: mask,
+        validity: mask.to_mask(array.len()),
         elements_exporter,
         list_size: array.list_size(),
     }))

vortex-duckdb/src/exporter/list.rs

@@ -10,6 +10,7 @@ use vortex::array::ToCanonical;
 use vortex::array::VectorExecutor;
 use vortex::array::arrays::ListViewArray;
 use vortex::array::arrays::PrimitiveArray;
+use vortex::array::vtable::ValidityHelper;
 use vortex::dtype::IntegerPType;
 use vortex::dtype::PTypeDowncastExt;
 use vortex::dtype::match_each_integer_ptype;
@@ -210,7 +211,7 @@ pub(crate) fn new_vector_exporter(
     let boxed = match_each_integer_ptype!(offsets.ptype(), |O| {
         match_each_integer_ptype!(sizes.ptype(), |S| {
             Box::new(ListVectorExporter {
-                validity: array.validity_mask(),
+                validity: array.validity().to_mask(array.len()),
                 duckdb_elements: shared_elements,
                 offsets: offsets.downcast::<O>(),
                 sizes: sizes.downcast::<O>(),

vortex-duckdb/src/exporter/mod.rs

@@ -34,6 +34,7 @@ use vortex::array::arrays::DictVTable;
 use vortex::array::arrays::StructArray;
 use vortex::array::arrays::TemporalArray;
 use vortex::array::iter::ArrayIterator;
+use vortex::array::vtable::ValidityHelper;
 use vortex::dtype::DType;
 use vortex::dtype::datetime::is_temporal_ext_type;
 use vortex::encodings::runend::RunEndVTable;
@@ -109,6 +110,7 @@ impl ArrayExporter {
         cache: &ConversionCache,
         session: &VortexSession,
     ) -> VortexResult<Self> {
+        assert!(array.validity().all_valid(array.len()));
         let fields = array
             .fields()
             .iter()

vortex-duckdb/src/exporter/struct_.rs

@@ -1,15 +1,19 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex::array::IntoArray;
 use vortex::array::arrays::StructArray;
+use vortex::array::optimizer::ArrayOptimizer;
+use vortex::array::vtable::ValidityHelper;
 use vortex::compute::mask;
 use vortex::error::VortexResult;
+use vortex::mask::Mask;
 use vortex::session::VortexSession;
 
+use crate::LogicalType;
 use crate::duckdb::Vector;
 use crate::exporter::ColumnExporter;
 use crate::exporter::ConversionCache;
+use crate::exporter::all_invalid;
 use crate::exporter::new_array_exporter;
 use crate::exporter::new_vector_array_exporter;
 use crate::exporter::validity;
@@ -60,19 +64,23 @@ pub(crate) fn new_vector_exporter(
     cache: &ConversionCache,
     session: &VortexSession,
 ) -> VortexResult<Box<dyn ColumnExporter>> {
-    let validity = array.validity_mask();
-    // DuckDB requires that the validity of the child be a subset of the parent struct so we mask out children with
-    // parents nullability
-    let validity_for_mask = array.dtype().is_nullable().then(|| !&validity);
+    let validity = array.validity().to_mask(array.len());
+
+    if validity.all_false() {
+        return Ok(all_invalid::new_exporter(
+            array.len(),
+            &LogicalType::try_from(array.dtype())?,
+        ));
+    }
 
     let children = array
         .fields()
         .iter()
         .map(|child| {
-            if let Some(mv) = validity_for_mask.as_ref() {
-                new_vector_array_exporter(mask(child, mv)?.into_array(), cache, session)
+            if matches!(validity, Mask::Values(_)) {
+                new_vector_array_exporter(mask(child, &validity)?.optimize()?, cache, session)
             } else {
-                new_vector_array_exporter(child.to_array(), cache, session)
+                new_vector_array_exporter(child.clone(), cache, session)
             }
         })
         .collect::<VortexResult<Vec<_>>>()?;