Flat Layout Execution

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-array/src/arrays/constant/mod.rs

@@ -6,7 +6,6 @@ pub use array::ConstantArray;
 
 mod compute;
 
-mod vector;
 mod vtable;
 
 pub use vtable::ConstantVTable;

vortex-array/src/arrays/constant/vector.rs

@@ -7,12 +7,12 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_scalar::Scalar;
 use vortex_scalar::ScalarValue;
+use vortex_vector::ScalarOps;
 use vortex_vector::Vector;
 use vortex_vector::VectorMutOps;
 
 use crate::EmptyMetadata;
 use crate::arrays::ConstantArray;
-use crate::arrays::constant::vector::to_vector;
 use crate::execution::ExecutionCtx;
 use crate::serde::ArrayChildren;
 use crate::vtable;
@@ -86,6 +86,10 @@ impl VTable for ConstantVTable {
     }
 
     fn batch_execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult<Vector> {
-        Ok(to_vector(array.scalar().clone(), array.len()).freeze())
+        Ok(array
+            .scalar()
+            .to_vector_scalar()
+            .repeat(array.len())
+            .freeze())
     }
 }

vortex-array/src/arrays/constant/vtable/mod.rs

@@ -7,15 +7,22 @@ use std::ops::Deref;
 
 use arrow_buffer::bit_iterator::BitIndexIterator;
 use vortex_buffer::BitBuffer;
+use vortex_compute::comparison::Compare;
+use vortex_compute::comparison::Equal;
+use vortex_compute::logical::LogicalOr;
 use vortex_dtype::DType;
 use vortex_dtype::IntegerPType;
 use vortex_dtype::Nullability;
 use vortex_dtype::PTypeDowncastExt;
 use vortex_dtype::match_each_integer_ptype;
+use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_err;
+use vortex_mask::Mask;
+use vortex_vector::BoolDatum;
 use vortex_vector::Datum;
+use vortex_vector::Vector;
 use vortex_vector::VectorOps;
 use vortex_vector::bool::BoolVector;
 use vortex_vector::listview::ListViewScalar;
@@ -121,15 +128,26 @@ impl VTable for ListContains {
             .try_into()
             .map_err(|_| vortex_err!("Wrong number of arguments for ListContains expression"))?;
 
-        let rhs = rhs
-            .into_scalar()
-            .ok_or_else(|| vortex_err!("Only supports constant RHS"))?;
+        let matches = match (lhs.as_scalar().is_some(), rhs.as_scalar().is_some()) {
+            (true, true) => {
+                todo!("Implement ListContains for two scalars")
+            }
+            (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(),
+                rhs.into_scalar().vortex_expect("scalar").into_list(),
+            ),
+            (false, false) => {
+                vortex_bail!(
+                    "ListContains currently only supports constant needle (RHS) or constant list (LHS)"
+                )
+            }
+        }?;
 
-        let result = list_contains_scalar(
-            lhs.ensure_vector(args.row_count).into_list(),
-            rhs.into_list(),
-        )?;
-        Ok(Datum::Vector(result.into()))
+        Ok(Datum::Vector(matches.into()))
     }
 
     fn stat_falsification(
@@ -190,7 +208,8 @@ pub fn list_contains(list: Expression, value: Expression) -> Expression {
     ListContains.new_expr(EmptyOptions, [list, value])
 }
 
-/// Returns a [`BoolArray`] where each bit represents if a list contains the scalar.
+/// Returns a [`BoolVector`] where each bit represents if a list contains the scalar.
+// FIXME(ngates): test implementation and move to vortex-compute
 fn list_contains_scalar(list: ListViewVector, value: ListViewScalar) -> VortexResult<BoolVector> {
     // If the list array is constant, we perform a single comparison.
     // if list.len() > 1 && list.is_constant() {
@@ -269,6 +288,30 @@ fn list_contains_scalar(list: ListViewVector, value: ListViewScalar) -> VortexRe
     Ok(BoolVector::new(list_matches, list.validity().clone()))
 }
 
+// Then there is a constant list scalar (haystack) being compared to an array of needles.
+// FIXME(ngates): test implementation and move to vortex-compute
+fn constant_list_scalar_contains(list: ListViewScalar, values: Vector) -> VortexResult<BoolVector> {
+    let elements = list.value().elements();
+
+    // For each element in the list, we perform a full comparison over the values and OR
+    // the results together.
+    let mut result: BoolVector = BoolVector::new(
+        BitBuffer::new_unset(values.len()),
+        Mask::new(values.len(), false),
+    );
+    for i in 0..elements.len() {
+        let element = Datum::Scalar(elements.scalar_at(i));
+        let compared: BoolDatum = Compare::<Equal>::compare(Datum::Vector(values.clone()), element);
+        let compared = Datum::from(compared)
+            .ensure_vector(values.len())
+            .into_bool();
+
+        result = LogicalOr::or(result, &compared);
+    }
+
+    Ok(result)
+}
+
 /// Returns a [`BitBuffer`] where each bit represents if a list contains the scalar, derived from a
 /// [`BoolArray`] of matches on the child elements array.
 ///

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

@@ -7,18 +7,25 @@ use vortex_error::VortexResult;
 use crate::Array;
 use crate::ArrayRef;
 use crate::IntoArray;
+use crate::arrays::ConstantArray;
 use crate::arrays::ScalarFnArray;
 use crate::expr::Expression;
+use crate::expr::Literal;
 use crate::expr::Root;
 
 impl dyn Array + '_ {
     /// Apply the expression to this array, producing a new array in constant time.
     pub fn apply(&self, expr: &Expression) -> VortexResult<ArrayRef> {
+        // If the expression is a root, return self.
         if expr.is::<Root>() {
-            // If the expression is a root, return self.
             return Ok(self.to_array());
         }
 
+        // Manually convert literals to ConstantArray.
+        if let Some(scalar) = expr.as_opt::<Literal>() {
+            return Ok(ConstantArray::new(scalar.clone(), self.len()).into_array());
+        }
+
         // Otherwise, collect the child arrays.
         let children: Vec<_> = expr
             .children()