-
Notifications
You must be signed in to change notification settings - Fork 202
Expand file tree
/
Copy patharbitrary.rs
More file actions
78 lines (65 loc) · 2.25 KB
/
Copy patharbitrary.rs
File metadata and controls
78 lines (65 loc) · 2.25 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors
use std::cmp::max;
use arbitrary::Result as AResult;
use arbitrary::Unstructured;
use crate::dtype::DType;
use crate::dtype::FieldName;
use crate::expr::Expression;
use crate::expr::and_collect;
use crate::expr::col;
use crate::expr::lit;
use crate::expr::pack;
use crate::scalar::arbitrary::random_scalar;
use crate::scalar_fn::ScalarFnVTableExt;
use crate::scalar_fn::fns::binary::Binary;
use crate::scalar_fn::fns::operators::Operator;
pub fn projection_expr(u: &mut Unstructured<'_>, dtype: &DType) -> AResult<Option<Expression>> {
let Some(struct_dtype) = dtype.as_struct_fields_opt() else {
return Ok(None);
};
let column_count = u.int_in_range::<usize>(0..=max(struct_dtype.nfields(), 10))?;
let cols = (0..column_count)
.map(|_| {
let get_item = u.choose_iter(struct_dtype.names().iter())?;
Ok((get_item.clone(), col(get_item.clone())))
})
.collect::<AResult<Vec<_>>>()?;
Ok(Some(pack(cols, u.arbitrary()?)))
}
pub fn filter_expr(u: &mut Unstructured<'_>, dtype: &DType) -> AResult<Option<Expression>> {
let Some(struct_dtype) = dtype.as_struct_fields_opt() else {
return Ok(None);
};
let filter_count = u.int_in_range::<usize>(0..=max(struct_dtype.nfields(), 10))?;
let filters = (0..filter_count)
.map(|_| {
let (col, dtype) =
u.choose_iter(struct_dtype.names().iter().zip(struct_dtype.fields()))?;
random_comparison(u, col, &dtype)
})
.collect::<AResult<Vec<_>>>()?;
Ok(and_collect(filters))
}
fn random_comparison(
u: &mut Unstructured<'_>,
name: &FieldName,
dtype: &DType,
) -> AResult<Expression> {
let scalar = random_scalar(u, dtype)?;
Ok(Binary.new_expr(
arbitrary_comparison_operator(u)?,
[col(name.clone()), lit(scalar)],
))
}
fn arbitrary_comparison_operator(u: &mut Unstructured<'_>) -> AResult<Operator> {
Ok(match u.int_in_range(0..=5)? {
0 => Operator::Eq,
1 => Operator::NotEq,
2 => Operator::Gt,
3 => Operator::Gte,
4 => Operator::Lt,
5 => Operator::Lte,
_ => unreachable!("range 0..=5"),
})
}