Optimize list_contains with prepared constant sets

encodings/sequence/src/compute/list_contains.rs

@@ -5,10 +5,13 @@ use vortex_array::ArrayRef;
 use vortex_array::ArrayView;
 use vortex_array::IntoArray;
 use vortex_array::arrays::BoolArray;
+use vortex_array::arrays::Constant;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::scalar::Scalar;
 use vortex_array::scalar_fn::fns::list_contains::ListContainsElementReduce;
 use vortex_array::scalar_fn::fns::list_contains::ListContainsOptions;
+use vortex_array::scalar_fn::fns::list_contains::PreparedSet;
+use vortex_array::scalar_fn::fns::list_contains::PreparedSetData;
 use vortex_array::validity::Validity;
 use vortex_buffer::BitBufferMut;
 use vortex_error::VortexResult;
@@ -23,7 +26,14 @@ impl ListContainsElementReduce for Sequence {
         element: ArrayView<'_, Self>,
         options: &ListContainsOptions,
     ) -> VortexResult<Option<ArrayRef>> {
-        let Some(list_scalar) = list.as_constant() else {
+        // Optimization prepares a constant list as a set, so the list is usually a prepared set.
+        let Some(list_scalar) = PreparedSet::set_of(list)
+            .map(PreparedSetData::list)
+            .or_else(|| {
+                list.as_opt::<Constant>()
+                    .map(|constant| constant.data().scalar())
+            })
+        else {
             return Ok(None);
         };
 
@@ -77,9 +87,12 @@ mod tests {
     use std::sync::Arc;
     use std::sync::LazyLock;
 
+    use rstest::rstest;
+    use vortex_array::ArrayRef;
     use vortex_array::IntoArray;
     use vortex_array::VortexSessionExecute;
     use vortex_array::arrays::BoolArray;
+    use vortex_array::arrays::ScalarFn;
     use vortex_array::assert_arrays_eq;
     use vortex_array::dtype::DType;
     use vortex_array::dtype::Nullability;
@@ -90,6 +103,7 @@ mod tests {
     use vortex_array::expr::root;
     use vortex_array::scalar::Scalar;
     use vortex_array::scalar_fn::fns::list_contains::ListContainsOptions;
+    use vortex_error::VortexResult;
     use vortex_session::VortexSession;
 
     use crate::Sequence;
@@ -137,6 +151,37 @@ mod tests {
         }
     }
 
+    /// Applies `list_contains(set, root())` to `array`, as given or optimized first.
+    fn apply_list_contains(array: ArrayRef, set: Scalar, optimize: bool) -> VortexResult<ArrayRef> {
+        let expr = list_contains(lit(set), root());
+        if !optimize {
+            return array.apply(&expr);
+        }
+        let expr = expr.bind(array.dtype())?.optimize()?;
+        array.apply_bound(&expr)
+    }
+
+    #[rstest]
+    #[case::applied(false)]
+    #[case::optimized(true)]
+    fn test_list_contains_reduces_a_prepared_set(#[case] optimize: bool) -> VortexResult<()> {
+        // The constant list becomes a prepared set before the sequence sees it. The sequence
+        // still answers from its bounds rather than leaving the probe to execution.
+        let list_scalar = Scalar::list(
+            Arc::new(I32.into()),
+            vec![1.into(), 3.into()],
+            Nullability::NonNullable,
+        );
+        let array = Sequence::try_new_typed(1, 1, Nullability::NonNullable, 3)?.into_array();
+
+        let result = apply_list_contains(array, list_scalar, optimize)?;
+
+        assert!(!result.is::<ScalarFn>());
+        let expected = BoolArray::from_iter([true, false, true]);
+        assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx());
+        Ok(())
+    }
+
     #[test]
     fn test_list_contains_null_list() {
         // A null list yields null for every row. The reduce rule used to panic on it.

vortex-array/src/arrays/chunked/compute/list_contains.rs

@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use crate::ArrayRef;
+use crate::IntoArray;
+use crate::array::ArrayView;
+use crate::arrays::Chunked;
+use crate::arrays::ChunkedArray;
+use crate::arrays::chunked::ChunkedArrayExt;
+use crate::dtype::DType;
+use crate::optimizer::ArrayOptimizer;
+use crate::scalar_fn::fns::list_contains::ListContains;
+use crate::scalar_fn::fns::list_contains::ListContainsElementReduce;
+use crate::scalar_fn::fns::list_contains::ListContainsOptions;
+use crate::scalar_fn::fns::list_contains::PreparedSet;
+
+/// Pushes `list_contains` of a prepared set into each chunk of the needles.
+///
+/// Each chunk gets a slice of the prepared set, which shares the one set, so that the set is built
+/// once for every chunk.
+impl ListContainsElementReduce for Chunked {
+    fn list_contains(
+        list: &ArrayRef,
+        needles: ArrayView<'_, Chunked>,
+        options: &ListContainsOptions,
+    ) -> VortexResult<Option<ArrayRef>> {
+        if PreparedSet::set_of(list).is_none() {
+            return Ok(None);
+        }
+
+        let mut offset = 0;
+        let chunks = needles
+            .iter_chunks()
+            .map(|chunk| {
+                let set = list.slice(offset..offset + chunk.len())?;
+                offset += chunk.len();
+                ListContains::try_new_opts(set, chunk.clone(), *options)?
+                    .into_array()
+                    .optimize()
+            })
+            .collect::<VortexResult<Vec<_>>>()?;
+
+        let dtype = DType::Bool(options.result_nullability(list.dtype(), needles.dtype()));
+
+        // SAFETY: each chunk is `list_contains` of the prepared set and a needle chunk of one dtype,
+        // so every chunk has the dtype of the whole result.
+        Ok(Some(
+            unsafe { ChunkedArray::new_unchecked(chunks, dtype) }.into_array(),
+        ))
+    }
+}

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

@@ -6,6 +6,7 @@ mod cast;
 mod fill_null;
 mod filter;
 pub(crate) mod kernel;
+mod list_contains;
 mod mask;
 pub(crate) mod rules;
 mod slice;

vortex-array/src/arrays/chunked/compute/rules.rs

@@ -10,7 +10,6 @@ use crate::array::ArrayView;
 use crate::arrays::Chunked;
 use crate::arrays::ChunkedArray;
 use crate::arrays::Constant;
-use crate::arrays::ConstantArray;
 use crate::arrays::ScalarFn;
 use crate::arrays::ScalarFnArray;
 use crate::arrays::chunked::ChunkedArrayExt;
@@ -21,10 +20,12 @@ use crate::optimizer::rules::ArrayParentReduceRule;
 use crate::optimizer::rules::ParentRuleSet;
 use crate::scalar_fn::fns::cast::CastReduceAdaptor;
 use crate::scalar_fn::fns::fill_null::FillNullReduceAdaptor;
+use crate::scalar_fn::fns::list_contains::ListContainsElementReduceAdaptor;
 
 pub(crate) const PARENT_RULES: ParentRuleSet<Chunked> = ParentRuleSet::new(&[
     ParentRuleSet::lift(&CastReduceAdaptor(Chunked)),
     ParentRuleSet::lift(&ChunkedUnaryScalarFnPushDownRule),
+    ParentRuleSet::lift(&ListContainsElementReduceAdaptor(Chunked)),
     ParentRuleSet::lift(&ChunkedConstantScalarFnPushDownRule),
     ParentRuleSet::lift(&FillNullReduceAdaptor(Chunked)),
 ]);
@@ -81,24 +82,25 @@ impl ArrayParentReduceRule<Chunked> for ChunkedConstantScalarFnPushDownRule {
             }
         }
 
+        let mut offset = 0;
         let new_chunks: Vec<_> = array
             .iter_chunks()
             .map(|chunk| {
+                let range = offset..offset + chunk.len();
+                offset += chunk.len();
+                // Each sibling is sliced to the chunk's rows, so that every chunk shares what the
+                // sibling shares rather than getting its own copy of the constant.
                 let new_children: Vec<_> = parent
                     .iter_children()
                     .enumerate()
                     .map(|(idx, child)| {
                         if idx == child_idx {
-                            chunk.clone()
+                            Ok(chunk.clone())
                         } else {
-                            ConstantArray::new(
-                                child.as_::<Constant>().scalar().clone(),
-                                chunk.len(),
-                            )
-                            .into_array()
+                            child.slice(range.clone())
                         }
                     })
-                    .collect();
+                    .try_collect()?;
 
                 ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children)?
                     .into_array()

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

@@ -17,6 +17,7 @@ use crate::optimizer::rules::ParentRuleSet;
 use crate::scalar_fn::fns::between::BetweenReduceAdaptor;
 use crate::scalar_fn::fns::cast::CastReduceAdaptor;
 use crate::scalar_fn::fns::fill_null::FillNullReduceAdaptor;
+use crate::scalar_fn::fns::list_contains::PrepareConstantListRule;
 use crate::scalar_fn::fns::not::NotReduceAdaptor;
 
 pub(crate) const PARENT_RULES: ParentRuleSet<Constant> = ParentRuleSet::new(&[
@@ -26,6 +27,7 @@ pub(crate) const PARENT_RULES: ParentRuleSet<Constant> = ParentRuleSet::new(&[
     ParentRuleSet::lift(&FillNullReduceAdaptor(Constant)),
     ParentRuleSet::lift(&FilterReduceAdaptor(Constant)),
     ParentRuleSet::lift(&NotReduceAdaptor(Constant)),
+    ParentRuleSet::lift(&PrepareConstantListRule),
     ParentRuleSet::lift(&SliceReduceAdaptor(Constant)),
     ParentRuleSet::lift(&TakeReduceAdaptor(Constant)),
 ]);

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

@@ -13,5 +13,5 @@ pub use vtable::ConstantArray;
 pub(crate) mod compute;
 
 mod vtable;
-
 pub use vtable::Constant;
+pub(crate) use vtable::canonical::list_scalar_elements;