feat[array]: list contains scalar (#5713)

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

Author: joseph-isaacs

vortex-array/src/arrays/scalar_fn/kernel.rs

@@ -57,7 +57,10 @@ impl Kernel for ScalarFnKernel {
             return_dtype: self.return_dtype,
         };
 
-        Ok(self.scalar_fn.execute(args)?.ensure_vector(self.row_count))
+        Ok(self
+            .scalar_fn
+            .execute(args)?
+            .unwrap_into_vector(self.row_count))
     }
 
     fn push_down_filter(self: Box<Self>, selection: &Mask) -> VortexResult<PushDownResult> {

vortex-array/src/arrays/scalar_fn/vtable/canonical.rs

@@ -37,7 +37,7 @@ impl CanonicalVTable<ScalarFnVTable> for ScalarFnVTable {
             .scalar_fn
             .execute(ctx)
             .vortex_expect("Canonicalize should be fallible")
-            .ensure_vector(len);
+            .unwrap_into_vector(len);
 
         result_vector.into_array(&array.dtype).to_canonical()
     }

vortex-array/src/executor.rs

@@ -55,6 +55,6 @@ impl VectorExecutor for ArrayRef {
 
     fn execute_vector(&self, session: &VortexSession) -> VortexResult<Vector> {
         let len = self.len();
-        Ok(self.execute_datum(session)?.ensure_vector(len))
+        Ok(self.execute_datum(session)?.unwrap_into_vector(len))
     }
 }

vortex-array/src/expr/exprs/binary.rs

@@ -137,8 +137,14 @@ impl VTable for Binary {
         match op {
             Operator::And => {
                 // FIXME(ngates): implement logical compute over datums
-                let lhs = lhs.ensure_vector(args.row_count).into_bool().into_arrow()?;
-                let rhs = rhs.ensure_vector(args.row_count).into_bool().into_arrow()?;
+                let lhs = lhs
+                    .unwrap_into_vector(args.row_count)
+                    .into_bool()
+                    .into_arrow()?;
+                let rhs = rhs
+                    .unwrap_into_vector(args.row_count)
+                    .into_bool()
+                    .into_arrow()?;
                 return Ok(Datum::Vector(
                     arrow_arith::boolean::and_kleene(&lhs, &rhs)?
                         .into_vector()?
@@ -147,8 +153,14 @@ impl VTable for Binary {
             }
             Operator::Or => {
                 // FIXME(ngates): implement logical compute over datums
-                let lhs = lhs.ensure_vector(args.row_count).into_bool().into_arrow()?;
-                let rhs = rhs.ensure_vector(args.row_count).into_bool().into_arrow()?;
+                let lhs = lhs
+                    .unwrap_into_vector(args.row_count)
+                    .into_bool()
+                    .into_arrow()?;
+                let rhs = rhs
+                    .unwrap_into_vector(args.row_count)
+                    .into_bool()
+                    .into_arrow()?;
                 return Ok(Datum::Vector(
                     arrow_arith::boolean::or_kleene(&lhs, &rhs)?
                         .into_vector()?

vortex-array/src/expr/exprs/list_contains.rs

@@ -20,7 +20,9 @@ use vortex_error::vortex_err;
 use vortex_mask::Mask;
 use vortex_vector::BoolDatum;
 use vortex_vector::Datum;
+use vortex_vector::ScalarOps;
 use vortex_vector::Vector;
+use vortex_vector::VectorMutOps;
 use vortex_vector::VectorOps;
 use vortex_vector::bool::BoolVector;
 use vortex_vector::listview::ListViewScalar;
@@ -128,14 +130,19 @@ impl VTable for ListContains {
 
         let matches = match (lhs.as_scalar().is_some(), rhs.as_scalar().is_some()) {
             (true, true) => {
-                todo!("Implement ListContains for two scalars")
+                let list = lhs.into_scalar().vortex_expect("scalar").into_list();
+                let needle = rhs.into_scalar().vortex_expect("scalar");
+                // Convert the needle scalar to a single-element vector and reuse
+                // constant_list_scalar_contains
+                let needle_vector = needle.repeat(1).freeze();
+                constant_list_scalar_contains(list, needle_vector)
             }
             (true, false) => constant_list_scalar_contains(
                 lhs.into_scalar().vortex_expect("scalar").into_list(),
                 rhs.into_vector().vortex_expect("vector"),
             ),
             (false, true) => list_contains_scalar(
-                lhs.ensure_vector(args.row_count).into_list(),
+                lhs.unwrap_into_vector(args.row_count).into_list(),
                 rhs.into_scalar().vortex_expect("scalar").into_list(),
             ),
             (false, false) => {
@@ -234,7 +241,7 @@ fn list_contains_scalar(list: ListViewVector, value: ListViewScalar) -> VortexRe
             row_count: elems.len(),
             return_dtype: DType::Bool(Nullability::Nullable),
         })?
-        .ensure_vector(elems.len())
+        .unwrap_into_vector(elems.len())
         .into_bool()
         .into_bits();
 
@@ -311,7 +318,7 @@ fn constant_list_scalar_contains(list: ListViewScalar, values: Vector) -> Vortex
             })?
             .into_bool();
         let compared = Datum::from(compared)
-            .ensure_vector(values.len())
+            .unwrap_into_vector(values.len())
             .into_bool();
 
         result = LogicalOr::or(&result, &compared);

vortex-array/src/expr/exprs/pack.rs

@@ -154,7 +154,7 @@ impl VTable for Pack {
             let fields: Box<[_]> = args
                 .datums
                 .into_iter()
-                .map(|v| v.ensure_vector(args.row_count))
+                .map(|v| v.unwrap_into_vector(args.row_count))
                 .collect();
             return Ok(Datum::Vector(
                 StructVector::try_new(Arc::new(fields), Mask::new_true(args.row_count))?.into(),

vortex-vector/src/datum.rs

@@ -57,7 +57,7 @@ impl From<Vector> for Datum {
 
 impl Datum {
     /// Ensure the datum is a vector by repeating the scalar value if necessary.
-    pub fn ensure_vector(self, len: usize) -> Vector {
+    pub fn unwrap_into_vector(self, len: usize) -> Vector {
         match self {
             Datum::Scalar(scalar) => scalar.repeat(len).freeze(),
             Datum::Vector(vector) => {

vortex-vector/src/decimal/generic.rs

@@ -97,11 +97,11 @@ impl<D: NativeDecimalType> DVector<D> {
         }
 
         // We assert that each element is within bounds for the given precision/scale.
-        if !elements.iter().all(|e| ps.is_valid(*e)) {
+        if let Some(invalid) = elements.iter().find(|e| !ps.is_valid(**e)) {
             vortex_bail!(
-                "One or more elements are out of bounds for precision {} and scale {}",
+                "One or more elements (e.g. {invalid}) are out of bounds for precision {} and scale {}",
                 ps.precision(),
-                ps.scale()
+                ps.scale(),
             );
         }