Fix null take on structs

vortex-array/src/arrays/dict/take.rs

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_error::VortexResult;
+use vortex_scalar::Scalar;
 
 use super::DictArray;
 use super::DictVTable;
@@ -10,6 +11,7 @@ use crate::ArrayRef;
 use crate::Canonical;
 use crate::ExecutionCtx;
 use crate::IntoArray;
+use crate::arrays::ConstantArray;
 use crate::expr::stats::Precision;
 use crate::expr::stats::Stat;
 use crate::expr::stats::StatsProvider;
@@ -63,14 +65,13 @@ fn precondition<V: VTable>(array: &V::Array, indices: &dyn Array) -> Option<Arra
         return Some(Canonical::empty(&result_dtype).into_array());
     }
 
-    // TODO(joe): shall we enable this seems expensive.
-    // if indices.all_invalid()? {
-    //     return Ok(
-    //         ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len())
-    //             .into_array()
-    //             .into(),
-    //     );
-    // }
+    // Fast-path for empty arrays: all indices must be null, return all-invalid result.
+    if array.is_empty() {
+        return Some(
+            ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len())
+                .into_array(),
+        );
+    }
 
     None
 }

vortex-array/src/arrays/struct_/compute/mod.rs

@@ -57,6 +57,27 @@ mod tests {
         );
     }
 
+    #[test]
+    fn take_empty_struct_with_nullable_indices() {
+        let struct_arr = StructArray::try_from_iter_with_validity(
+            [("a", BoolArray::from_iter(Vec::<bool>::new()).into_array())],
+            Validity::AllValid,
+        )
+        .unwrap();
+        let indices = PrimitiveArray::from_option_iter([Option::<u64>::None]);
+        let taken = take(struct_arr.as_ref(), indices.as_ref()).unwrap();
+        assert_eq!(taken.len(), 1);
+        assert!(taken.all_invalid().unwrap());
+    }
+
+    #[test]
+    fn take_empty_primitive_with_nullable_indices() {
+        let arr = PrimitiveArray::from_iter(Vec::<u64>::new());
+        let indices = PrimitiveArray::from_option_iter([Option::<u64>::None]);
+        let taken = take(arr.as_ref(), indices.as_ref()).unwrap();
+        assert_eq!(taken.len(), 1);
+    }
+
     #[test]
     fn take_field_struct() {
         let struct_arr =

vortex-array/src/arrays/struct_/compute/take.rs

@@ -13,7 +13,6 @@ use crate::arrays::StructArray;
 use crate::arrays::StructVTable;
 use crate::arrays::TakeExecute;
 use crate::compute;
-use crate::validity::Validity;
 use crate::vtable::ValidityHelper;
 
 impl TakeExecute for StructVTable {
@@ -22,18 +21,6 @@ impl TakeExecute for StructVTable {
         indices: &dyn Array,
         _ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
-        // If the struct array is empty then the indices must be all null, otherwise it will access
-        // an out of bounds element
-        if array.is_empty() {
-            return StructArray::try_new_with_dtype(
-                array.unmasked_fields().clone(),
-                array.struct_fields().clone(),
-                indices.len(),
-                Validity::AllInvalid,
-            )
-            .map(StructArray::into_array)
-            .map(Some);
-        }
         // The validity is applied to the struct validity,
         let inner_indices = &compute::fill_null(
             indices,