vortex_expr/
arbitrary.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::cmp::max;
5
6use arbitrary::{Result as AResult, Unstructured};
7use vortex_dtype::{DType, FieldName};
8use vortex_scalar::arbitrary::random_scalar;
9
10use crate::{BinaryExpr, ExprRef, Operator, and_collect, col, lit, pack};
11
12pub fn projection_expr(u: &mut Unstructured<'_>, dtype: &DType) -> AResult<Option<ExprRef>> {
13    let Some(struct_dtype) = dtype.as_struct_fields_opt() else {
14        return Ok(None);
15    };
16
17    let column_count = u.int_in_range::<usize>(0..=max(struct_dtype.nfields(), 10))?;
18
19    let cols = (0..column_count)
20        .map(|_| {
21            let get_item = u.choose_iter(struct_dtype.names().iter())?;
22            Ok((get_item.clone(), col(get_item.clone())))
23        })
24        .collect::<AResult<Vec<_>>>()?;
25
26    Ok(Some(pack(cols, u.arbitrary()?)))
27}
28
29pub fn filter_expr(u: &mut Unstructured<'_>, dtype: &DType) -> AResult<Option<ExprRef>> {
30    let Some(struct_dtype) = dtype.as_struct_fields_opt() else {
31        return Ok(None);
32    };
33
34    let filter_count = u.int_in_range::<usize>(0..=max(struct_dtype.nfields(), 10))?;
35
36    let filters = (0..filter_count)
37        .map(|_| {
38            let (col, dtype) =
39                u.choose_iter(struct_dtype.names().iter().zip(struct_dtype.fields()))?;
40            random_comparison(u, col, &dtype)
41        })
42        .collect::<AResult<Vec<_>>>()?;
43
44    Ok(and_collect(filters))
45}
46
47fn random_comparison(
48    u: &mut Unstructured<'_>,
49    name: &FieldName,
50    dtype: &DType,
51) -> AResult<ExprRef> {
52    let scalar = random_scalar(u, dtype)?;
53    Ok(BinaryExpr::new_expr(
54        col(name.clone()),
55        arbitrary_comparison_operator(u)?,
56        lit(scalar),
57    ))
58}
59
60fn arbitrary_comparison_operator(u: &mut Unstructured<'_>) -> AResult<Operator> {
61    Ok(match u.int_in_range(0..=5)? {
62        0 => Operator::Eq,
63        1 => Operator::NotEq,
64        2 => Operator::Gt,
65        3 => Operator::Gte,
66        4 => Operator::Lt,
67        5 => Operator::Lte,
68        _ => unreachable!("range 0..=5"),
69    })
70}
vortex_expr/arbitrary.rs

vortex_expr/
arbitrary.rs
