Add debug assertions to ComputeFn results

encodings/fsst/src/compute/compare.rs

@@ -1,10 +1,11 @@
 use fsst::Symbol;
 use vortex_array::array::ConstantArray;
 use vortex_array::compute::{compare, CompareFn, Operator};
-use vortex_array::{ArrayDType, ArrayData, ArrayLen, IntoArrayVariant, ToArrayData};
+use vortex_array::{ArrayDType, ArrayData, ArrayLen, IntoArrayData, IntoArrayVariant};
 use vortex_buffer::Buffer;
-use vortex_dtype::DType;
-use vortex_error::VortexResult;
+use vortex_dtype::{DType, Nullability};
+use vortex_error::{VortexExpect, VortexResult};
+use vortex_scalar::Scalar;
 
 use crate::{FSSTArray, FSSTEncoding};
 
@@ -16,10 +17,16 @@ impl CompareFn<FSSTArray> for FSSTEncoding {
         operator: Operator,
     ) -> VortexResult<Option<ArrayData>> {
         match (rhs.as_constant(), operator) {
-            // TODO(ngates): implement short-circuit comparisons for other operators.
-            (Some(constant_array), Operator::Eq | Operator::NotEq) => compare_fsst_constant(
+            (Some(constant), _) if constant.is_null() => {
+                // All comparisons to null must return null
+                Ok(Some(
+                    ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), lhs.len())
+                        .into_array(),
+                ))
+            }
+            (Some(constant), Operator::Eq | Operator::NotEq) => compare_fsst_constant(
                 lhs,
-                &ConstantArray::new(constant_array, lhs.len()),
+                &ConstantArray::new(constant, lhs.len()),
                 operator == Operator::Eq,
             )
             .map(Some),
@@ -49,34 +56,31 @@ fn compare_fsst_constant(
     let compressor = compressor.build();
 
     let encoded_scalar = match left.dtype() {
-        DType::Utf8(_) => right
-            .scalar()
-            .as_utf8()
-            .value()
-            .map(|scalar| Buffer::from(compressor.compress(scalar.as_bytes()))),
-        DType::Binary(_) => right
-            .scalar()
-            .as_binary()
-            .value()
-            .map(|scalar| Buffer::from(compressor.compress(scalar.as_slice()))),
+        DType::Utf8(_) => {
+            let value = right
+                .scalar()
+                .as_utf8()
+                .value()
+                .vortex_expect("Expected non-null scalar");
+            Buffer::from(compressor.compress(value.as_bytes()))
+        }
+        DType::Binary(_) => {
+            let value = right
+                .scalar()
+                .as_binary()
+                .value()
+                .vortex_expect("Expected non-null scalar");
+            Buffer::from(compressor.compress(value.as_slice()))
+        }
         _ => unreachable!("FSSTArray can only have string or binary data type"),
     };
 
-    match encoded_scalar {
-        None => {
-            // Eq and NotEq on null values yield nulls, per the Arrow behavior.
-            Ok(right.to_array())
-        }
-        Some(encoded_scalar) => {
-            let rhs = ConstantArray::new(encoded_scalar, left.len());
-
-            compare(
-                left.codes(),
-                rhs,
-                if equal { Operator::Eq } else { Operator::NotEq },
-            )
-        }
-    }
+    let rhs = ConstantArray::new(encoded_scalar, left.len());
+    compare(
+        left.codes(),
+        rhs,
+        if equal { Operator::Eq } else { Operator::NotEq },
+    )
 }
 
 #[cfg(test)]

pyvortex/src/expr.rs

@@ -304,7 +304,7 @@ pub fn scalar_helper(dtype: DType, value: &Bound<'_, PyAny>) -> PyResult<Scalar>
                 .iter()
                 .map(|element| scalar_helper(element_type.as_ref().clone(), element))
                 .collect::<PyResult<Vec<_>>>()?;
-            Ok(Scalar::list(element_type, values))
+            Ok(Scalar::list(element_type, values, Nullability::Nullable))
         }
         DType::Extension(..) => todo!(),
     }

vortex-array/src/array/chunked/compute/compare.rs

@@ -1,9 +1,9 @@
-use vortex_dtype::{DType, Nullability};
+use vortex_dtype::DType;
 use vortex_error::VortexResult;
 
 use crate::array::{ChunkedArray, ChunkedEncoding};
 use crate::compute::{compare, slice, CompareFn, Operator};
-use crate::{ArrayData, IntoArrayData};
+use crate::{ArrayDType, ArrayData, IntoArrayData};
 
 impl CompareFn<ChunkedArray> for ChunkedEncoding {
     fn compare(
@@ -24,7 +24,11 @@ impl CompareFn<ChunkedArray> for ChunkedEncoding {
         }
 
         Ok(Some(
-            ChunkedArray::try_new(compare_chunks, DType::Bool(Nullability::Nullable))?.into_array(),
+            ChunkedArray::try_new(
+                compare_chunks,
+                DType::Bool((lhs.dtype().is_nullable() || rhs.dtype().is_nullable()).into()),
+            )?
+            .into_array(),
         ))
     }
 }

vortex-array/src/array/list/compute/mod.rs

@@ -23,7 +23,11 @@ impl ScalarAtFn<ListArray> for ListEncoding {
         let elem = array.elements_at(index)?;
         let scalars: Vec<Scalar> = (0..elem.len()).map(|i| scalar_at(&elem, i)).try_collect()?;
 
-        Ok(Scalar::list(Arc::new(elem.dtype().clone()), scalars))
+        Ok(Scalar::list(
+            Arc::new(elem.dtype().clone()),
+            scalars,
+            array.dtype().nullability(),
+        ))
     }
 }
 

vortex-array/src/array/list/mod.rs

@@ -197,7 +197,7 @@ impl ValidityVTable<ListArray> for ListEncoding {
 mod test {
     use std::sync::Arc;
 
-    use vortex_dtype::PType;
+    use vortex_dtype::{Nullability, PType};
     use vortex_scalar::Scalar;
 
     use crate::array::list::ListArray;
@@ -228,15 +228,27 @@ mod test {
             ListArray::try_new(elements.into_array(), offsets.into_array(), validity).unwrap();
 
         assert_eq!(
-            Scalar::list(Arc::new(PType::I32.into()), vec![1.into(), 2.into()]),
+            Scalar::list(
+                Arc::new(PType::I32.into()),
+                vec![1.into(), 2.into()],
+                Nullability::Nullable
+            ),
             scalar_at(&list, 0).unwrap()
         );
         assert_eq!(
-            Scalar::list(Arc::new(PType::I32.into()), vec![3.into(), 4.into()]),
+            Scalar::list(
+                Arc::new(PType::I32.into()),
+                vec![3.into(), 4.into()],
+                Nullability::Nullable
+            ),
             scalar_at(&list, 1).unwrap()
         );
         assert_eq!(
-            Scalar::list(Arc::new(PType::I32.into()), vec![5.into()]),
+            Scalar::list(
+                Arc::new(PType::I32.into()),
+                vec![5.into()],
+                Nullability::Nullable
+            ),
             scalar_at(&list, 2).unwrap()
         );
     }

vortex-array/src/array/varbin/compute/compare.rs

@@ -6,7 +6,7 @@ use vortex_error::{vortex_bail, VortexResult};
 use crate::array::{VarBinArray, VarBinEncoding};
 use crate::arrow::{Datum, FromArrowArray};
 use crate::compute::{CompareFn, Operator};
-use crate::{ArrayData, IntoArrayData};
+use crate::{ArrayDType, ArrayData, IntoArrayData};
 
 // This implementation exists so we can have custom translation of RHS to arrow that's not the same as IntoCanonical
 impl CompareFn<VarBinArray> for VarBinEncoding {
@@ -17,6 +17,8 @@ impl CompareFn<VarBinArray> for VarBinEncoding {
         operator: Operator,
     ) -> VortexResult<Option<ArrayData>> {
         if let Some(rhs_const) = rhs.as_constant() {
+            let nullable = lhs.dtype().is_nullable() || rhs_const.dtype().is_nullable();
+
             let lhs = Datum::try_from(lhs.clone().into_array())?;
 
             // TODO(robert): Handle LargeString/Binary arrays
@@ -46,7 +48,7 @@ impl CompareFn<VarBinArray> for VarBinEncoding {
                 Operator::Lte => cmp::lt_eq(&lhs, arrow_rhs)?,
             };
 
-            Ok(Some(ArrayData::from_arrow(&array, true)))
+            Ok(Some(ArrayData::from_arrow(&array, nullable)))
         } else {
             Ok(None)
         }

vortex-array/src/builders/list.rs

@@ -156,17 +156,27 @@ mod tests {
 
         builder
             .append_value(
-                Scalar::list(dtype.clone(), vec![1i32.into(), 2i32.into(), 3i32.into()]).as_list(),
+                Scalar::list(
+                    dtype.clone(),
+                    vec![1i32.into(), 2i32.into(), 3i32.into()],
+                    Nullability::NonNullable,
+                )
+                .as_list(),
             )
             .unwrap();
 
         builder
-            .append_value(Scalar::empty(dtype.clone()).as_list())
+            .append_value(Scalar::list_empty(dtype.clone(), Nullability::NonNullable).as_list())
             .unwrap();
 
         builder
             .append_value(
-                Scalar::list(dtype, vec![4i32.into(), 5i32.into(), 6i32.into()]).as_list(),
+                Scalar::list(
+                    dtype,
+                    vec![4i32.into(), 5i32.into(), 6i32.into()],
+                    Nullability::NonNullable,
+                )
+                .as_list(),
             )
             .unwrap();
 

vortex-array/src/compute/binary_numeric.rs

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use arrow_array::ArrayRef;
-use vortex_dtype::DType;
+use vortex_dtype::{DType, PType};
 use vortex_error::{vortex_bail, VortexError, VortexResult};
 use vortex_scalar::{BinaryNumericOperator, Scalar};
 
@@ -117,13 +117,43 @@ pub fn binary_numeric(
     // Check if LHS supports the operation directly.
     if let Some(fun) = lhs.encoding().binary_numeric_fn() {
         if let Some(result) = fun.binary_numeric(lhs, rhs, op)? {
+            debug_assert_eq!(
+                result.len(),
+                lhs.len(),
+                "Numeric operation length mismatch {}",
+                lhs.encoding().id()
+            );
+            debug_assert_eq!(
+                result.dtype(),
+                &DType::Primitive(
+                    PType::try_from(lhs.dtype())?,
+                    (lhs.dtype().is_nullable() || rhs.dtype().is_nullable()).into()
+                ),
+                "Numeric operation dtype mismatch {}",
+                lhs.encoding().id()
+            );
             return Ok(result);
         }
     }
 
     // Check if RHS supports the operation directly.
     if let Some(fun) = rhs.encoding().binary_numeric_fn() {
         if let Some(result) = fun.binary_numeric(rhs, lhs, op)? {
+            debug_assert_eq!(
+                result.len(),
+                lhs.len(),
+                "Numeric operation length mismatch {}",
+                rhs.encoding().id()
+            );
+            debug_assert_eq!(
+                result.dtype(),
+                &DType::Primitive(
+                    PType::try_from(lhs.dtype())?,
+                    (lhs.dtype().is_nullable() || rhs.dtype().is_nullable()).into()
+                ),
+                "Numeric operation dtype mismatch {}",
+                rhs.encoding().id()
+            );
             return Ok(result);
         }
     }

vortex-array/src/compute/boolean.rs

@@ -2,6 +2,7 @@ use std::sync::Arc;
 
 use arrow_array::cast::AsArray;
 use arrow_array::ArrayRef;
+use vortex_dtype::DType;
 use vortex_error::{vortex_bail, vortex_err, VortexError, VortexResult};
 
 use crate::arrow::FromArrowArray;
@@ -106,16 +107,42 @@ pub fn binary_boolean(
         .encoding()
         .binary_boolean_fn()
         .and_then(|f| f.binary_boolean(lhs, rhs, op).transpose())
+        .transpose()?
     {
-        return result;
+        debug_assert_eq!(
+            result.len(),
+            lhs.len(),
+            "Boolean operation length mismatch {}",
+            lhs.encoding().id()
+        );
+        debug_assert_eq!(
+            result.dtype(),
+            &DType::Bool((lhs.dtype().is_nullable() || rhs.dtype().is_nullable()).into()),
+            "Boolean operation dtype mismatch {}",
+            lhs.encoding().id()
+        );
+        return Ok(result);
     }
 
     if let Some(result) = rhs
         .encoding()
         .binary_boolean_fn()
         .and_then(|f| f.binary_boolean(rhs, lhs, op).transpose())
+        .transpose()?
     {
-        return result;
+        debug_assert_eq!(
+            result.len(),
+            lhs.len(),
+            "Boolean operation length mismatch {}",
+            rhs.encoding().id()
+        );
+        debug_assert_eq!(
+            result.dtype(),
+            &DType::Bool((lhs.dtype().is_nullable() || rhs.dtype().is_nullable()).into()),
+            "Boolean operation dtype mismatch {}",
+            rhs.encoding().id()
+        );
+        return Ok(result);
     }
 
     log::debug!(

vortex-array/src/compute/cast.rs

@@ -33,6 +33,25 @@ pub fn try_cast(array: impl AsRef<ArrayData>, dtype: &DType) -> VortexResult<Arr
         return Ok(array.clone());
     }
 
+    let casted = try_cast_impl(array, dtype)?;
+
+    debug_assert_eq!(
+        casted.len(),
+        array.len(),
+        "Cast length mismatch {}",
+        array.encoding().id()
+    );
+    debug_assert_eq!(
+        casted.dtype(),
+        dtype,
+        "Cast dtype mismatch {}",
+        array.encoding().id()
+    );
+
+    Ok(casted)
+}
+
+fn try_cast_impl(array: &ArrayData, dtype: &DType) -> VortexResult<ArrayData> {
     // TODO(ngates): check for null_count if dtype is non-nullable
     if let Some(f) = array.encoding().cast_fn() {
         return f.cast(array, dtype);