Use non debug asserts in compute function entry points (#2914)

Author: joseph-isaacs

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

@@ -11,13 +11,22 @@ impl TakeFn<&ConstantArray> for ConstantEncoding {
     fn take(&self, array: &ConstantArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
         match indices.validity_mask()?.boolean_buffer() {
             AllOr::All => {
-                Ok(ConstantArray::new(array.scalar().clone(), indices.len()).into_array())
+                let nullability = array.dtype().nullability() | indices.dtype().nullability();
+                let scalar = Scalar::new(
+                    array.scalar().dtype().with_nullability(nullability),
+                    array.scalar().value().clone(),
+                );
+                Ok(ConstantArray::new(scalar, indices.len()).into_array())
+            }
+            AllOr::None => {
+                Ok(ConstantArray::new(
+                    Scalar::null(array.dtype().with_nullability(
+                        array.dtype().nullability() | indices.dtype().nullability(),
+                    )),
+                    indices.len(),
+                )
+                .into_array())
             }
-            AllOr::None => Ok(ConstantArray::new(
-                Scalar::null(array.dtype().clone()),
-                indices.len(),
-            )
-            .into_array()),
             AllOr::Some(v) => {
                 let arr = ConstantArray::new(array.scalar().clone(), indices.len()).into_array();
 
@@ -38,6 +47,7 @@ impl TakeFn<&ConstantArray> for ConstantEncoding {
 #[cfg(test)]
 mod tests {
     use vortex_buffer::buffer;
+    use vortex_dtype::Nullability;
     use vortex_mask::AllOr;
 
     use crate::arrays::{ConstantArray, PrimitiveArray};
@@ -58,6 +68,10 @@ mod tests {
         )
         .unwrap();
         let valid_indices: &[usize] = &[1usize];
+        assert_eq!(
+            &array.dtype().with_nullability(Nullability::Nullable),
+            taken.dtype()
+        );
         assert_eq!(
             taken.to_primitive().unwrap().as_slice::<i32>(),
             &[42, 42, 42]
@@ -67,4 +81,23 @@ mod tests {
             AllOr::Some(valid_indices)
         );
     }
+
+    #[test]
+    fn take_all_valid_indices() {
+        let array = ConstantArray::new(42, 10).to_array();
+        let taken = take(
+            &array,
+            &PrimitiveArray::new(buffer![0, 5, 7], Validity::AllValid).into_array(),
+        )
+        .unwrap();
+        assert_eq!(
+            &array.dtype().with_nullability(Nullability::Nullable),
+            taken.dtype()
+        );
+        assert_eq!(
+            taken.to_primitive().unwrap().as_slice::<i32>(),
+            &[42, 42, 42]
+        );
+        assert_eq!(taken.validity_mask().unwrap().indices(), AllOr::All);
+    }
 }

vortex-array/src/compute/between.rs

@@ -88,10 +88,10 @@ pub fn between(
     upper: &dyn Array,
     options: &BetweenOptions,
 ) -> VortexResult<ArrayRef> {
-    debug_assert!(arr.dtype().eq_ignore_nullability(lower.dtype()));
-    debug_assert!(arr.dtype().eq_ignore_nullability(upper.dtype()));
-    debug_assert_eq!(arr.len(), lower.len());
-    debug_assert_eq!(arr.len(), upper.len());
+    assert!(arr.dtype().eq_ignore_nullability(lower.dtype()));
+    assert!(arr.dtype().eq_ignore_nullability(upper.dtype()));
+    assert_eq!(arr.len(), lower.len());
+    assert_eq!(arr.len(), upper.len());
 
     // A quick check to see if either array might is a null constant array.
     if lower.is_invalid(0)? || upper.is_invalid(0)? {
@@ -110,8 +110,8 @@ pub fn between(
 
     let result = between_impl(arr, lower, upper, options)?;
 
-    debug_assert_eq!(result.len(), arr.len());
-    debug_assert_eq!(
+    assert_eq!(result.len(), arr.len());
+    assert_eq!(
         result.dtype(),
         &DType::Bool(
             arr.dtype().nullability() | lower.dtype().nullability() | upper.dtype().nullability()

vortex-array/src/compute/binary_numeric.rs

@@ -171,13 +171,13 @@ fn arrow_numeric(
 
 #[inline(always)]
 fn check_numeric_result(result: ArrayRef, lhs: &dyn Array, rhs: &dyn Array) -> ArrayRef {
-    debug_assert_eq!(
+    assert_eq!(
         result.len(),
         lhs.len(),
         "Numeric operation length mismatch {}",
         rhs.encoding()
     );
-    debug_assert_eq!(
+    assert_eq!(
         result.dtype(),
         &DType::Primitive(
             PType::try_from(lhs.dtype())

vortex-array/src/compute/boolean.rs

@@ -115,13 +115,13 @@ pub fn binary_boolean(
         .and_then(|f| f.binary_boolean(lhs, rhs, op).transpose())
         .transpose()?
     {
-        debug_assert_eq!(
+        assert_eq!(
             result.len(),
             lhs.len(),
             "Boolean operation length mismatch {}",
             lhs.encoding()
         );
-        debug_assert_eq!(
+        assert_eq!(
             result.dtype(),
             &DType::Bool((lhs.dtype().is_nullable() || rhs.dtype().is_nullable()).into()),
             "Boolean operation dtype mismatch {}",
@@ -136,13 +136,13 @@ pub fn binary_boolean(
         .and_then(|f| f.binary_boolean(rhs, lhs, op).transpose())
         .transpose()?
     {
-        debug_assert_eq!(
+        assert_eq!(
             result.len(),
             lhs.len(),
             "Boolean operation length mismatch {}",
             rhs.encoding()
         );
-        debug_assert_eq!(
+        assert_eq!(
             result.dtype(),
             &DType::Bool((lhs.dtype().is_nullable() || rhs.dtype().is_nullable()).into()),
             "Boolean operation dtype mismatch {}",

vortex-array/src/compute/cast.rs

@@ -31,13 +31,13 @@ pub fn try_cast(array: &dyn Array, dtype: &DType) -> VortexResult<ArrayRef> {
 
     let casted = try_cast_impl(array, dtype)?;
 
-    debug_assert_eq!(
+    assert_eq!(
         casted.len(),
         array.len(),
         "Cast length mismatch {}",
         array.encoding()
     );
-    debug_assert_eq!(
+    assert_eq!(
         casted.dtype(),
         dtype,
         "Cast dtype mismatch {}",

vortex-array/src/compute/compare.rs

@@ -211,7 +211,7 @@ fn arrow_compare(
 
 #[inline(always)]
 fn check_compare_result(result: &dyn Array, lhs: &dyn Array, rhs: &dyn Array) {
-    debug_assert_eq!(
+    assert_eq!(
         result.len(),
         lhs.len(),
         "CompareFn result length ({}) mismatch for left encoding {}, left len {}, right encoding {}, right len {}",
@@ -221,7 +221,7 @@ fn check_compare_result(result: &dyn Array, lhs: &dyn Array, rhs: &dyn Array) {
         rhs.encoding(),
         rhs.len()
     );
-    debug_assert_eq!(
+    assert_eq!(
         result.dtype(),
         &DType::Bool((lhs.dtype().is_nullable() || rhs.dtype().is_nullable()).into()),
         "CompareFn result dtype ({}) mismatch for left encoding {}, right encoding {}",

vortex-array/src/compute/fill_forward.rs

@@ -41,13 +41,13 @@ pub fn fill_forward(array: &dyn Array) -> VortexResult<ArrayRef> {
             ))
         })?;
 
-    debug_assert_eq!(
+    assert_eq!(
         filled.len(),
         array.len(),
         "FillForward length mismatch {}",
         array.encoding()
     );
-    debug_assert_eq!(
+    assert_eq!(
         filled.dtype(),
         array.dtype(),
         "FillForward dtype mismatch {}",

vortex-array/src/compute/fill_null.rs

@@ -40,13 +40,13 @@ pub fn fill_null(array: &dyn Array, fill_value: Scalar) -> VortexResult<ArrayRef
     let fill_value_nullability = fill_value.dtype().nullability();
     let filled = fill_null_impl(array, fill_value)?;
 
-    debug_assert_eq!(
+    assert_eq!(
         filled.len(),
         array.len(),
         "FillNull length mismatch {}",
         array.encoding()
     );
-    debug_assert_eq!(
+    assert_eq!(
         filled.dtype(),
         &array.dtype().with_nullability(fill_value_nullability),
         "FillNull dtype mismatch {}",

vortex-array/src/compute/filter.rs

@@ -92,13 +92,13 @@ impl ComputeFn for Filter {
 
         let filtered = filter_impl(array, mask)?;
 
-        debug_assert_eq!(
+        assert_eq!(
             filtered.len(),
             true_count,
             "Filter length mismatch {}",
             array.encoding()
         );
-        debug_assert_eq!(
+        assert_eq!(
             filtered.dtype(),
             array.dtype(),
             "Filter dtype mismatch {}",

vortex-array/src/compute/is_constant.rs

@@ -117,7 +117,7 @@ fn is_constant_impl(array: &dyn Array, opts: &IsConstantOpts) -> VortexResult<Op
         }
     }
 
-    debug_assert!(
+    assert!(
         all_valid,
         "All values must be valid as an invariant of the VTable."
     );