vortex_expr/exprs/
list_contains.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::{Debug, Display, Formatter};
5use std::hash::Hash;
6
7use vortex_array::compute::list_contains as compute_list_contains;
8use vortex_array::{ArrayRef, DeserializeMetadata, EmptyMetadata};
9use vortex_dtype::DType;
10use vortex_error::{VortexResult, vortex_bail};
11
12use crate::{
13    AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, LiteralVTable, Scope, StatsCatalog,
14    VTable, and, gt, lit, lt, or, vtable,
15};
16
17vtable!(ListContains);
18
19#[allow(clippy::derived_hash_with_manual_eq)]
20#[derive(Debug, Clone, Hash, Eq)]
21pub struct ListContainsExpr {
22    list: ExprRef,
23    value: ExprRef,
24}
25
26impl PartialEq for ListContainsExpr {
27    fn eq(&self, other: &Self) -> bool {
28        self.list.eq(&other.list) && self.value.eq(&other.value)
29    }
30}
31
32pub struct ListContainsExprEncoding;
33
34impl VTable for ListContainsVTable {
35    type Expr = ListContainsExpr;
36    type Encoding = ListContainsExprEncoding;
37    type Metadata = EmptyMetadata;
38
39    fn id(_encoding: &Self::Encoding) -> ExprId {
40        ExprId::new_ref("list_contains")
41    }
42
43    fn encoding(_expr: &Self::Expr) -> ExprEncodingRef {
44        ExprEncodingRef::new_ref(ListContainsExprEncoding.as_ref())
45    }
46
47    fn metadata(_expr: &Self::Expr) -> Option<Self::Metadata> {
48        Some(EmptyMetadata)
49    }
50
51    fn children(expr: &Self::Expr) -> Vec<&ExprRef> {
52        vec![&expr.list, &expr.value]
53    }
54
55    fn with_children(_expr: &Self::Expr, children: Vec<ExprRef>) -> VortexResult<Self::Expr> {
56        Ok(ListContainsExpr::new(
57            children[0].clone(),
58            children[1].clone(),
59        ))
60    }
61
62    fn build(
63        _encoding: &Self::Encoding,
64        _metadata: &<Self::Metadata as DeserializeMetadata>::Output,
65        children: Vec<ExprRef>,
66    ) -> VortexResult<Self::Expr> {
67        if children.len() != 2 {
68            vortex_bail!(
69                "ListContains expression must have exactly 2 children, got {}",
70                children.len()
71            );
72        }
73        Ok(ListContainsExpr::new(
74            children[0].clone(),
75            children[1].clone(),
76        ))
77    }
78
79    fn evaluate(expr: &Self::Expr, scope: &Scope) -> VortexResult<ArrayRef> {
80        compute_list_contains(
81            expr.list.evaluate(scope)?.as_ref(),
82            expr.value.evaluate(scope)?.as_ref(),
83        )
84    }
85
86    fn return_dtype(expr: &Self::Expr, scope: &DType) -> VortexResult<DType> {
87        Ok(DType::Bool(
88            expr.list.return_dtype(scope)?.nullability()
89                | expr.value.return_dtype(scope)?.nullability(),
90        ))
91    }
92}
93
94impl ListContainsExpr {
95    pub fn new(list: ExprRef, value: ExprRef) -> Self {
96        Self { list, value }
97    }
98
99    pub fn new_expr(list: ExprRef, value: ExprRef) -> ExprRef {
100        Self::new(list, value).into_expr()
101    }
102
103    pub fn value(&self) -> &ExprRef {
104        &self.value
105    }
106}
107
108/// Creates an expression that checks if a value is contained in a list.
109///
110/// Returns a boolean array indicating whether the value appears in each list.
111///
112/// ```rust
113/// # use vortex_expr::{list_contains, lit, root};
114/// let expr = list_contains(root(), lit(42));
115/// ```
116pub fn list_contains(list: ExprRef, value: ExprRef) -> ExprRef {
117    ListContainsExpr::new(list, value).into_expr()
118}
119
120impl Display for ListContainsExpr {
121    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
122        write!(f, "contains({}, {})", &self.list, &self.value)
123    }
124}
125
126impl AnalysisExpr for ListContainsExpr {
127    // falsification(contains([1,2,5], x)) =>
128    //   falsification(x != 1) and falsification(x != 2) and falsification(x != 5)
129
130    fn stat_falsification(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
131        let min = self.list.min(catalog)?;
132        let max = self.list.max(catalog)?;
133        // If the list is constant when we can compare each element to the value
134        if min == max {
135            let list_ = min
136                .as_opt::<LiteralVTable>()
137                .and_then(|l| l.value().as_list_opt())
138                .and_then(|l| l.elements())?;
139            if list_.is_empty() {
140                // contains([], x) is always false.
141                return Some(lit(true));
142            }
143            let value_max = self.value.max(catalog)?;
144            let value_min = self.value.min(catalog)?;
145
146            return list_
147                .iter()
148                .map(move |v| {
149                    or(
150                        lt(value_max.clone(), lit(v.clone())),
151                        gt(value_min.clone(), lit(v.clone())),
152                    )
153                })
154                .reduce(and);
155        }
156
157        None
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use vortex_array::arrays::{BoolArray, BooleanBuffer, ListArray, PrimitiveArray};
164    use vortex_array::stats::Stat;
165    use vortex_array::validity::Validity;
166    use vortex_array::{Array, ArrayRef, IntoArray};
167    use vortex_dtype::PType::I32;
168    use vortex_dtype::{DType, Field, FieldPath, FieldPathSet, Nullability, StructFields};
169    use vortex_scalar::Scalar;
170    use vortex_utils::aliases::hash_map::HashMap;
171
172    use crate::list_contains::list_contains;
173    use crate::pruning::checked_pruning_expr;
174    use crate::{Arc, HashSet, Scope, and, col, get_item, gt, lit, lt, or, root};
175
176    fn test_array() -> ArrayRef {
177        ListArray::try_new(
178            PrimitiveArray::from_iter(vec![1, 1, 2, 2, 2, 2, 2, 3, 3, 3]).into_array(),
179            PrimitiveArray::from_iter(vec![0, 5, 10]).into_array(),
180            Validity::AllValid,
181        )
182        .unwrap()
183        .into_array()
184    }
185
186    #[test]
187    pub fn test_one() {
188        let arr = test_array();
189
190        let expr = list_contains(root(), lit(1));
191        let item = expr.evaluate(&Scope::new(arr)).unwrap();
192
193        assert_eq!(item.scalar_at(0), Scalar::bool(true, Nullability::Nullable));
194        assert_eq!(
195            item.scalar_at(1),
196            Scalar::bool(false, Nullability::Nullable)
197        );
198    }
199
200    #[test]
201    pub fn test_all() {
202        let arr = test_array();
203
204        let expr = list_contains(root(), lit(2));
205        let item = expr.evaluate(&Scope::new(arr)).unwrap();
206
207        assert_eq!(item.scalar_at(0), Scalar::bool(true, Nullability::Nullable));
208        assert_eq!(item.scalar_at(1), Scalar::bool(true, Nullability::Nullable));
209    }
210
211    #[test]
212    pub fn test_none() {
213        let arr = test_array();
214
215        let expr = list_contains(root(), lit(4));
216        let item = expr.evaluate(&Scope::new(arr)).unwrap();
217
218        assert_eq!(
219            item.scalar_at(0),
220            Scalar::bool(false, Nullability::Nullable)
221        );
222        assert_eq!(
223            item.scalar_at(1),
224            Scalar::bool(false, Nullability::Nullable)
225        );
226    }
227
228    #[test]
229    pub fn test_empty() {
230        let arr = ListArray::try_new(
231            PrimitiveArray::from_iter(vec![1, 1, 2, 2, 2]).into_array(),
232            PrimitiveArray::from_iter(vec![0, 5, 5]).into_array(),
233            Validity::AllValid,
234        )
235        .unwrap()
236        .into_array();
237
238        let expr = list_contains(root(), lit(2));
239        let item = expr.evaluate(&Scope::new(arr)).unwrap();
240
241        assert_eq!(item.scalar_at(0), Scalar::bool(true, Nullability::Nullable));
242        assert_eq!(
243            item.scalar_at(1),
244            Scalar::bool(false, Nullability::Nullable)
245        );
246    }
247
248    #[test]
249    pub fn test_nullable() {
250        let arr = ListArray::try_new(
251            PrimitiveArray::from_iter(vec![1, 1, 2, 2, 2]).into_array(),
252            PrimitiveArray::from_iter(vec![0, 5, 5]).into_array(),
253            Validity::Array(BoolArray::from(BooleanBuffer::from(vec![true, false])).into_array()),
254        )
255        .unwrap()
256        .into_array();
257
258        let expr = list_contains(root(), lit(2));
259        let item = expr.evaluate(&Scope::new(arr)).unwrap();
260
261        assert_eq!(item.scalar_at(0), Scalar::bool(true, Nullability::Nullable));
262        assert!(!item.is_valid(1).unwrap());
263    }
264
265    #[test]
266    pub fn test_return_type() {
267        let scope = DType::Struct(
268            StructFields::new(
269                ["array"].into(),
270                vec![DType::List(
271                    Arc::new(DType::Primitive(I32, Nullability::NonNullable)),
272                    Nullability::Nullable,
273                )],
274            ),
275            Nullability::NonNullable,
276        );
277
278        let expr = list_contains(get_item("array", root()), lit(2));
279
280        // Expect nullable, although scope is non-nullable
281        assert_eq!(
282            expr.return_dtype(&scope).unwrap(),
283            DType::Bool(Nullability::Nullable)
284        );
285    }
286
287    #[test]
288    pub fn list_falsification() {
289        let expr = list_contains(
290            lit(Scalar::list(
291                Arc::new(DType::Primitive(I32, Nullability::NonNullable)),
292                vec![1.into(), 2.into(), 3.into()],
293                Nullability::NonNullable,
294            )),
295            col("a"),
296        );
297
298        let (expr, st) = checked_pruning_expr(
299            &expr,
300            &FieldPathSet::from_iter([
301                FieldPath::from_iter([Field::Name("a".into()), Field::Name("max".into())]),
302                FieldPath::from_iter([Field::Name("a".into()), Field::Name("min".into())]),
303            ]),
304        )
305        .unwrap();
306
307        assert_eq!(
308            &expr,
309            &and(
310                and(
311                    or(lt(col("a_max"), lit(1i32)), gt(col("a_min"), lit(1i32)),),
312                    or(lt(col("a_max"), lit(2i32)), gt(col("a_min"), lit(2i32)),)
313                ),
314                or(lt(col("a_max"), lit(3i32)), gt(col("a_min"), lit(3i32)),)
315            )
316        );
317
318        assert_eq!(
319            st.map(),
320            &HashMap::from_iter([(
321                FieldPath::from_name("a"),
322                HashSet::from([Stat::Min, Stat::Max])
323            )])
324        );
325    }
326}
vortex_expr/exprs/list_contains.rs

vortex_expr/exprs/
list_contains.rs
