perf: push a prepared list_contains set into Dict and RunEnd needles

encodings/runend/src/compute/list_contains.rs

@@ -0,0 +1,106 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_array::ArrayRef;
+use vortex_array::ArrayView;
+use vortex_array::IntoArray;
+use vortex_array::arrays::ScalarFnArray;
+use vortex_array::scalar_fn::ScalarFnVTableExt;
+use vortex_array::scalar_fn::fns::list_contains::ListContains;
+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_error::VortexResult;
+
+use crate::RunEnd;
+use crate::array::RunEndArrayExt;
+use crate::array::RunEndArraySlotsExt;
+
+/// Pushes `list_contains` of a prepared set into the values of run-end encoded needles.
+///
+/// The values get the prepared set at their own length, sharing the one set, so that the set is
+/// not built again for them.
+impl ListContainsElementReduce for RunEnd {
+    fn list_contains(
+        list: &ArrayRef,
+        needles: ArrayView<'_, RunEnd>,
+        options: &ListContainsOptions,
+    ) -> VortexResult<Option<ArrayRef>> {
+        let Some(set) = PreparedSet::set_of(list) else {
+            return Ok(None);
+        };
+
+        // The values get the shared set at their own length, which can exceed the rows of a
+        // zero-length slice that kept its runs.
+        let values = needles.values();
+        let set = ScalarFnArray::try_new_with_len(PreparedSet.bind(set.clone()), [], values.len())?
+            .into_array();
+        let values = ListContains::try_new_opts(set, values.clone(), *options)?.into_array();
+
+        // SAFETY: the ends are unchanged, and the values have one result for each run.
+        Ok(Some(
+            unsafe {
+                RunEnd::new_unchecked(
+                    needles.ends().clone(),
+                    values,
+                    needles.offset(),
+                    needles.len(),
+                )
+            }
+            .into_array(),
+        ))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_array::IntoArray;
+    use vortex_array::VortexSessionExecute;
+    use vortex_array::array_session;
+    use vortex_array::arrays::BoolArray;
+    use vortex_array::arrays::ConstantArray;
+    use vortex_array::assert_arrays_eq;
+    use vortex_array::dtype::DType;
+    use vortex_array::dtype::Nullability;
+    use vortex_array::dtype::PType;
+    use vortex_array::optimizer::ArrayOptimizer;
+    use vortex_array::scalar::Scalar;
+    use vortex_array::scalar_fn::fns::list_contains::ListContains;
+    use vortex_buffer::buffer;
+    use vortex_error::VortexResult;
+
+    use crate::RunEnd;
+
+    #[test]
+    fn constant_list_is_looked_up_once_per_run() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let needles = RunEnd::try_new(
+            buffer![2u32, 5, 6].into_array(),
+            buffer![1i32, 2, 3].into_array(),
+            &mut ctx,
+        )?
+        .into_array();
+        let element = DType::Primitive(PType::I32, Nullability::NonNullable);
+        let list = Scalar::list(
+            element,
+            vec![
+                Scalar::primitive(2i32, Nullability::NonNullable),
+                Scalar::primitive(3i32, Nullability::NonNullable),
+            ],
+            Nullability::NonNullable,
+        );
+        let list = ConstantArray::new(list, needles.len()).into_array();
+
+        let array = ListContains::try_new(list, needles)?
+            .into_array()
+            .optimize()?;
+        assert!(array.is::<RunEnd>());
+
+        assert_arrays_eq!(
+            array,
+            BoolArray::from_iter([false, false, true, true, true, true]),
+            &mut ctx
+        );
+        Ok(())
+    }
+}

encodings/runend/src/compute/mod.rs

@@ -7,6 +7,7 @@ mod fill_null;
 pub(crate) mod filter;
 pub(crate) mod is_constant;
 pub(crate) mod is_sorted;
+mod list_contains;
 pub(crate) mod min_max;
 pub(crate) mod sum;
 pub(crate) mod take;

encodings/runend/src/rules.rs

@@ -15,6 +15,7 @@ use vortex_array::optimizer::rules::ArrayParentReduceRule;
 use vortex_array::optimizer::rules::ParentRuleSet;
 use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor;
 use vortex_array::scalar_fn::fns::fill_null::FillNullReduceAdaptor;
+use vortex_array::scalar_fn::fns::list_contains::ListContainsElementReduceAdaptor;
 use vortex_error::VortexResult;
 
 use crate::RunEnd;
@@ -26,6 +27,7 @@ pub(super) const RULES: ParentRuleSet<RunEnd> = ParentRuleSet::new(&[
     // eagerly (surfacing out-of-range errors immediately) rather than being pushed lazily into
     // the values array by the generic scalar function push-down rule.
     ParentRuleSet::lift(&CastReduceAdaptor(RunEnd)),
+    ParentRuleSet::lift(&ListContainsElementReduceAdaptor(RunEnd)),
     ParentRuleSet::lift(&RunEndScalarFnRule),
     ParentRuleSet::lift(&FillNullReduceAdaptor(RunEnd)),
 ]);

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

@@ -0,0 +1,74 @@
+// 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::ConstantArray;
+use crate::arrays::Dict;
+use crate::arrays::DictArray;
+use crate::arrays::dict::DictArrayExt;
+use crate::arrays::dict::DictArraySlotsExt;
+use crate::builtins::ArrayBuiltins;
+use crate::dtype::DType;
+use crate::optimizer::ArrayOptimizer;
+use crate::scalar::Scalar;
+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 the values of a dictionary of needles.
+///
+/// The values get a slice of the prepared set, which shares the one set, so that the set is not
+/// built again for them.
+impl ListContainsElementReduce for Dict {
+    fn list_contains(
+        list: &ArrayRef,
+        needles: ArrayView<'_, Dict>,
+        options: &ListContainsOptions,
+    ) -> VortexResult<Option<ArrayRef>> {
+        let Some(set) = PreparedSet::set_of(list) else {
+            return Ok(None);
+        };
+        let dtype = DType::Bool(options.result_nullability(list.dtype(), needles.dtype()));
+
+        // Off SQL null semantics an empty list answers `false` for every needle, a null one
+        // included, which the dictionary would answer `null` for from a null code.
+        if set.is_empty() && !options.sql_null_semantics {
+            return Ok(Some(
+                ConstantArray::new(Scalar::bool(false, dtype.nullability()), needles.len())
+                    .into_array(),
+            ));
+        }
+
+        // A sliced dictionary can have more values than code rows. Do not increase the work in
+        // that case.
+        let values = needles.values();
+        let codes = needles.codes();
+        if values.len() > codes.len() {
+            return Ok(None);
+        }
+
+        // A null code answers `null`, as a null needle does against a list with elements.
+        let values =
+            ListContains::try_new_opts(list.slice(0..values.len())?, values.clone(), *options)?
+                .into_array()
+                .optimize()?;
+
+        // SAFETY: the codes are unchanged, and the values have one result for each value.
+        let dict = unsafe {
+            DictArray::new_unchecked(codes.clone(), values)
+                .set_all_values_referenced(needles.has_all_values_referenced())
+        }
+        .into_array();
+
+        if dict.dtype() == &dtype {
+            Ok(Some(dict))
+        } else {
+            dict.cast(dtype).map(Some)
+        }
+    }
+}

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

@@ -7,6 +7,7 @@ mod fill_null;
 pub(crate) mod is_constant;
 pub(crate) mod is_sorted;
 mod like;
+mod list_contains;
 mod mask;
 pub(crate) mod min_max;
 pub(crate) mod rules;

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

@@ -31,6 +31,7 @@ use crate::optimizer::rules::ParentRuleSet;
 use crate::scalar_fn::fns::cast::Cast;
 use crate::scalar_fn::fns::cast::CastReduceAdaptor;
 use crate::scalar_fn::fns::like::LikeReduceAdaptor;
+use crate::scalar_fn::fns::list_contains::ListContainsElementReduceAdaptor;
 use crate::scalar_fn::fns::mask::MaskReduceAdaptor;
 use crate::scalar_fn::fns::pack::Pack;
 use crate::validity::Validity;
@@ -41,6 +42,7 @@ pub(crate) const PARENT_RULES: ParentRuleSet<Dict> = ParentRuleSet::new(&[
     ParentRuleSet::lift(&MaskReduceAdaptor(Dict)),
     ParentRuleSet::lift(&LikeReduceAdaptor(Dict)),
     ParentRuleSet::lift(&DictionaryChunkedValuesPullUpRule),
+    ParentRuleSet::lift(&ListContainsElementReduceAdaptor(Dict)),
     ParentRuleSet::lift(&DictionaryScalarFnValuesPushDownRule),
     ParentRuleSet::lift(&DictionaryScalarFnCodesPullUpRule),
     ParentRuleSet::lift(&SliceReduceAdaptor(Dict)),

vortex-array/src/scalar_fn/fns/list_contains/mod.rs

@@ -1899,6 +1899,28 @@ mod tests {
         )
     }
 
+    #[rstest]
+    // Non-strict, so a null code keeps the dictionary rewrite from pushing the function into the
+    // values: the needle is executed into its canonical encoding and looked up in the set.
+    #[case::executed(ListContainsOptions::default(), [Some(false), Some(true), None, Some(true)])]
+    // Strict, so the rewrite looks up the dictionary's two values instead of its four rows.
+    #[case::pushed_into_values(SQL, [Some(false), Some(true), None, Some(true)])]
+    fn constant_set_looks_up_a_dictionary_needle(
+        #[case] options: ListContainsOptions,
+        #[case] expected: [Option<bool>; 4],
+    ) -> VortexResult<()> {
+        let needles = DictArray::try_new(
+            PrimitiveArray::from_option_iter([Some(0u32), Some(1), None, Some(1)]).into_array(),
+            PrimitiveArray::from_iter([1i32, 2]).into_array(),
+        )?
+        .into_array();
+        let set = i32_set(vec![Some(2), Some(4)]);
+        assert_rows_agree(
+            needles.apply(&list_contains_opts(lit(set), root(), options))?,
+            BoolArray::from_iter(expected),
+        )
+    }
+
     #[test]
     fn constant_set_of_strings_looks_up_an_encoded_needle() -> VortexResult<()> {
         let needles = VarBinArray::from_iter(

vortex-array/src/scalar_fn/fns/list_contains/prepared/function.rs

@@ -40,8 +40,9 @@ use crate::scalar_fn::ScalarFnVTableExt;
 /// Every row of such an array holds the list. A slice or a filter of it keeps the options, so it
 /// looks up in the same set. A [`ListContainsElementReduce`] rule of the needle encoding can push
 /// the function into the needle's parts with slices of the array, so that the chunks of a chunked
-/// needle all look up in one set, and a [`ListContainsElementKernel`] of the needle encoding gets
-/// the set through [`PreparedSet::set_of`].
+/// needle or the values of a dictionary all look up in one set, and a
+/// [`ListContainsElementKernel`] of the needle encoding gets the set through
+/// [`PreparedSet::set_of`].
 ///
 /// It serializes as the list, and deserializes into a set prepared again.
 ///

vortex-array/src/scalar_fn/fns/list_contains/prepared/tests.rs

@@ -19,13 +19,16 @@ use crate::arrays::Chunked;
 use crate::arrays::ChunkedArray;
 use crate::arrays::ConstantArray;
 use crate::arrays::DecimalArray;
+use crate::arrays::Dict;
+use crate::arrays::DictArray;
 use crate::arrays::FixedSizeListArray;
 use crate::arrays::ListArray;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::ScalarFn;
 use crate::arrays::StructArray;
 use crate::arrays::VarBinViewArray;
 use crate::arrays::chunked::ChunkedArrayExt;
+use crate::arrays::dict::DictArraySlotsExt;
 use crate::arrays::scalar_fn::ScalarFnArrayExt;
 use crate::assert_arrays_eq;
 use crate::builders::builder_with_capacity_in;
@@ -554,3 +557,67 @@ fn test_optimized_expression_over_chunked_needles() -> VortexResult<()> {
     );
     Ok(())
 }
+
+/// Dictionary needles with a null code, against a set that is empty or not.
+#[rstest]
+#[case::default(ListContainsOptions::default(), set_with_null(), [Some(true), Some(false), None, Some(true)])]
+#[case::sql(ListContainsOptions { sql_null_semantics: true }, set_with_null(), [Some(true), None, None, Some(true)])]
+// Off SQL null semantics an empty list answers `false` even for the null code.
+#[case::empty(ListContainsOptions::default(), empty_i32_set(), [Some(false); 4])]
+#[case::empty_sql(ListContainsOptions { sql_null_semantics: true }, empty_i32_set(), [Some(false), Some(false), None, Some(false)])]
+fn test_dictionary_needles_with_null_codes(
+    #[case] options: ListContainsOptions,
+    #[case] set: Scalar,
+    #[case] expected: [Option<bool>; 4],
+) -> VortexResult<()> {
+    let mut ctx = array_session().create_execution_ctx();
+    let needles = DictArray::try_new(
+        PrimitiveArray::from_option_iter([Some(0u8), Some(1), None, Some(0)]).into_array(),
+        buffer![2i32, 3].into_array(),
+    )?
+    .into_array();
+    let list = ConstantArray::new(set, needles.len()).into_array();
+
+    let array = ListContains::try_new_opts(list, needles, options)?
+        .into_array()
+        .optimize()?;
+    // The rule took the node: it is a dictionary of answers, or a constant for the empty list.
+    assert!(!array.is::<ScalarFn>());
+
+    assert_arrays_eq!(array, BoolArray::from_iter(expected), &mut ctx);
+    Ok(())
+}
+
+/// The empty set of nullable `i32`.
+fn empty_i32_set() -> Scalar {
+    Scalar::list(
+        DType::Primitive(PType::I32, Nullability::Nullable),
+        vec![],
+        Nullability::NonNullable,
+    )
+}
+
+#[test]
+fn test_constant_list_over_dictionary_needles() -> VortexResult<()> {
+    let mut ctx = array_session().create_execution_ctx();
+    let needles = DictArray::try_new(
+        buffer![0u8, 1, 0, 2, 1].into_array(),
+        PrimitiveArray::from_option_iter([Some(2i32), Some(3), None]).into_array(),
+    )?
+    .into_array();
+    let list = ConstantArray::new(set_with_null(), needles.len()).into_array();
+
+    // The prepared list still lets the function into the values of the dictionary.
+    let array = ListContains::try_new(list, needles)?
+        .into_array()
+        .optimize()?;
+    let values = array.as_::<Dict>().values().clone();
+    assert!(PreparedSet::set_of(values.as_::<ScalarFn>().get_child(0)).is_some());
+
+    assert_arrays_eq!(
+        array,
+        BoolArray::from_iter([Some(true), Some(false), Some(true), None, Some(false)]),
+        &mut ctx
+    );
+    Ok(())
+}