vortex_array/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::expr::{Binary, Expression, Operator, VTableExt, and_collect, col, lit, pack};
11
12pub fn projection_expr(u: &mut Unstructured<'_>, dtype: &DType) -> AResult<Option<Expression>> {
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<Expression>> {
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<Expression> {
52    let scalar = random_scalar(u, dtype)?;
53    Ok(Binary.new_expr(
54        arbitrary_comparison_operator(u)?,
55        [col(name.clone()), lit(scalar)],
56    ))
57}
58
59fn arbitrary_comparison_operator(u: &mut Unstructured<'_>) -> AResult<Operator> {
60    Ok(match u.int_in_range(0..=5)? {
61        0 => Operator::Eq,
62        1 => Operator::NotEq,
63        2 => Operator::Gt,
64        3 => Operator::Gte,
65        4 => Operator::Lt,
66        5 => Operator::Lte,
67        _ => unreachable!("range 0..=5"),
68    })
69}
vortex_array/expr/arbitrary.rs

vortex_array/expr/
arbitrary.rs
