Skip to content

Commit b0c0b95

Browse files
a10ya10y
authored
feat: new expression builders (#1829)
For binary expressions, we now have eq, not_eq, lt, lt_eq, gt, gt_eq, and, or (see vortex-expr/src/binary.rs). Add example usage doctests for all For literals we have lit For Identity we have ident
1 parent a81c606 commit b0c0b95

File tree

15 files changed

+380
-330
lines changed

15 files changed

+380
-330
lines changed

‎pyvortex/src/expr.rs‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,9 @@
11
use pyo3::exceptions::PyValueError;
22
use pyo3::prelude::*;
33
use pyo3::types::*;
4-
use vortex::dtype::field::Field;
54
use vortex::dtype::half::f16;
65
use vortex::dtype::{DType, Nullability, PType};
7-
use vortex::expr::{BinaryExpr, Column, ExprRef, Literal, Operator};
6+
use vortex::expr::{col, lit, BinaryExpr, ExprRef, Operator};
87
use vortex::scalar::Scalar;
98

109
use crate::dtype::PyDType;
@@ -243,12 +242,7 @@ impl PyExpr {
243242
pub fn column<'py>(name: &Bound<'py, PyString>) -> PyResult<Bound<'py, PyExpr>> {
244243
let py = name.py();
245244
let name: String = name.extract()?;
246-
Bound::new(
247-
py,
248-
PyExpr {
249-
inner: Column::new_expr(Field::from(name)),
250-
},
251-
)
245+
Bound::new(py, PyExpr { inner: col(name) })
252246
}
253247

254248
#[pyfunction]
@@ -264,7 +258,7 @@ pub fn scalar<'py>(dtype: DType, value: &Bound<'py, PyAny>) -> PyResult<Bound<'p
264258
Bound::new(
265259
py,
266260
PyExpr {
267-
inner: Literal::new_expr(scalar_helper(dtype, value)?),
261+
inner: lit(scalar_helper(dtype, value)?),
268262
},
269263
)
270264
}

‎vortex-expr/src/binary.rs‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,3 +76,183 @@ impl PartialEq<dyn Any> for BinaryExpr {
7676
.unwrap_or(false)
7777
}
7878
}
79+
80+
/// Create a new `BinaryExpr` using the `Eq` operator.
81+
///
82+
/// ## Example usage
83+
///
84+
/// ```
85+
/// use vortex_array::array::{BoolArray, PrimitiveArray };
86+
/// use vortex_array::{IntoArrayData, IntoArrayVariant};
87+
/// use vortex_array::validity::Validity;
88+
/// use vortex_buffer::buffer;
89+
/// use vortex_expr::{eq, ident, lit};
90+
///
91+
/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable).into_array();
92+
/// let result = eq(ident(), lit(3)).evaluate(&xs).unwrap();
93+
///
94+
/// assert_eq!(
95+
/// result.into_bool().unwrap().boolean_buffer(),
96+
/// BoolArray::from_iter(vec![false, false, true]).boolean_buffer(),
97+
/// );
98+
/// ```
99+
pub fn eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
100+
BinaryExpr::new_expr(lhs, Operator::Eq, rhs)
101+
}
102+
103+
/// Create a new `BinaryExpr` using the `NotEq` operator.
104+
///
105+
/// ## Example usage
106+
///
107+
/// ```
108+
/// use vortex_array::array::{BoolArray, PrimitiveArray };
109+
/// use vortex_array::{IntoArrayData, IntoArrayVariant};
110+
/// use vortex_array::validity::Validity;
111+
/// use vortex_buffer::buffer;
112+
/// use vortex_expr::{ident, lit, not_eq};
113+
///
114+
/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable).into_array();
115+
/// let result = not_eq(ident(), lit(3)).evaluate(&xs).unwrap();
116+
///
117+
/// assert_eq!(
118+
/// result.into_bool().unwrap().boolean_buffer(),
119+
/// BoolArray::from_iter(vec![true, true, false]).boolean_buffer(),
120+
/// );
121+
/// ```
122+
pub fn not_eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
123+
BinaryExpr::new_expr(lhs, Operator::NotEq, rhs)
124+
}
125+
126+
/// Create a new `BinaryExpr` using the `Gte` operator.
127+
///
128+
/// ## Example usage
129+
///
130+
/// ```
131+
/// use vortex_array::array::{BoolArray, PrimitiveArray };
132+
/// use vortex_array::{IntoArrayData, IntoArrayVariant};
133+
/// use vortex_array::validity::Validity;
134+
/// use vortex_buffer::buffer;
135+
/// use vortex_expr::{gt_eq, ident, lit};
136+
///
137+
/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable).into_array();
138+
/// let result = gt_eq(ident(), lit(3)).evaluate(&xs).unwrap();
139+
///
140+
/// assert_eq!(
141+
/// result.into_bool().unwrap().boolean_buffer(),
142+
/// BoolArray::from_iter(vec![false, false, true]).boolean_buffer(),
143+
/// );
144+
/// ```
145+
pub fn gt_eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
146+
BinaryExpr::new_expr(lhs, Operator::Gte, rhs)
147+
}
148+
149+
/// Create a new `BinaryExpr` using the `Gt` operator.
150+
///
151+
/// ## Example usage
152+
///
153+
/// ```
154+
/// use vortex_array::array::{BoolArray, PrimitiveArray };
155+
/// use vortex_array::{IntoArrayData, IntoArrayVariant};
156+
/// use vortex_array::validity::Validity;
157+
/// use vortex_buffer::buffer;
158+
/// use vortex_expr::{gt, ident, lit};
159+
///
160+
/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable).into_array();
161+
/// let result = gt(ident(), lit(2)).evaluate(&xs).unwrap();
162+
///
163+
/// assert_eq!(
164+
/// result.into_bool().unwrap().boolean_buffer(),
165+
/// BoolArray::from_iter(vec![false, false, true]).boolean_buffer(),
166+
/// );
167+
/// ```
168+
pub fn gt(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
169+
BinaryExpr::new_expr(lhs, Operator::Gt, rhs)
170+
}
171+
172+
/// Create a new `BinaryExpr` using the `Lte` operator.
173+
///
174+
/// ## Example usage
175+
///
176+
/// ```
177+
/// use vortex_array::array::{BoolArray, PrimitiveArray };
178+
/// use vortex_array::{IntoArrayData, IntoArrayVariant};
179+
/// use vortex_array::validity::Validity;
180+
/// use vortex_buffer::buffer;
181+
/// use vortex_expr::{ident, lit, lt_eq};
182+
///
183+
/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable).into_array();
184+
/// let result = lt_eq(ident(), lit(2)).evaluate(&xs).unwrap();
185+
///
186+
/// assert_eq!(
187+
/// result.into_bool().unwrap().boolean_buffer(),
188+
/// BoolArray::from_iter(vec![true, true, false]).boolean_buffer(),
189+
/// );
190+
/// ```
191+
pub fn lt_eq(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
192+
BinaryExpr::new_expr(lhs, Operator::Lte, rhs)
193+
}
194+
195+
/// Create a new `BinaryExpr` using the `Lt` operator.
196+
///
197+
/// ## Example usage
198+
///
199+
/// ```
200+
/// use vortex_array::array::{BoolArray, PrimitiveArray };
201+
/// use vortex_array::{IntoArrayData, IntoArrayVariant};
202+
/// use vortex_array::validity::Validity;
203+
/// use vortex_buffer::buffer;
204+
/// use vortex_expr::{ident, lit, lt};
205+
///
206+
/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable).into_array();
207+
/// let result = lt(ident(), lit(3)).evaluate(&xs).unwrap();
208+
///
209+
/// assert_eq!(
210+
/// result.into_bool().unwrap().boolean_buffer(),
211+
/// BoolArray::from_iter(vec![true, true, false]).boolean_buffer(),
212+
/// );
213+
/// ```
214+
pub fn lt(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
215+
BinaryExpr::new_expr(lhs, Operator::Lt, rhs)
216+
}
217+
218+
/// Create a new `BinaryExpr` using the `Or` operator.
219+
///
220+
/// ## Example usage
221+
///
222+
/// ```
223+
/// use vortex_array::array::BoolArray;
224+
/// use vortex_array::{IntoArrayData, IntoArrayVariant};
225+
/// use vortex_expr::{ ident, lit, or};
226+
///
227+
/// let xs = BoolArray::from_iter(vec![true, false, true]).into_array();
228+
/// let result = or(ident(), lit(false)).evaluate(&xs).unwrap();
229+
///
230+
/// assert_eq!(
231+
/// result.into_bool().unwrap().boolean_buffer(),
232+
/// BoolArray::from_iter(vec![true, false, true]).boolean_buffer(),
233+
/// );
234+
/// ```
235+
pub fn or(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
236+
BinaryExpr::new_expr(lhs, Operator::Or, rhs)
237+
}
238+
239+
/// Create a new `BinaryExpr` using the `And` operator.
240+
///
241+
/// ## Example usage
242+
///
243+
/// ```
244+
/// use vortex_array::array::BoolArray;
245+
/// use vortex_array::{IntoArrayData, IntoArrayVariant};
246+
/// use vortex_expr::{and, ident, lit};
247+
///
248+
/// let xs = BoolArray::from_iter(vec![true, false, true]).into_array();
249+
/// let result = and(ident(), lit(true)).evaluate(&xs).unwrap();
250+
///
251+
/// assert_eq!(
252+
/// result.into_bool().unwrap().boolean_buffer(),
253+
/// BoolArray::from_iter(vec![true, false, true]).boolean_buffer(),
254+
/// );
255+
/// ```
256+
pub fn and(lhs: ExprRef, rhs: ExprRef) -> ExprRef {
257+
BinaryExpr::new_expr(lhs, Operator::And, rhs)
258+
}

‎vortex-expr/src/column.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,12 @@ impl Column {
2626
}
2727
}
2828

29+
pub fn col(field: impl Into<Field>) -> ExprRef {
30+
Arc::new(Column {
31+
field: field.into(),
32+
})
33+
}
34+
2935
impl From<String> for Column {
3036
fn from(value: String) -> Self {
3137
Column {

‎vortex-expr/src/datafusion.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ use datafusion_physical_expr::{expressions, PhysicalExpr};
77
use vortex_error::{vortex_bail, vortex_err, VortexError, VortexResult};
88
use vortex_scalar::Scalar;
99

10-
use crate::{BinaryExpr, Column, ExprRef, Like, Literal, Operator};
10+
use crate::{lit, BinaryExpr, Column, ExprRef, Like, Operator};
1111

1212
pub fn convert_expr_to_vortex(physical_expr: Arc<dyn PhysicalExpr>) -> VortexResult<ExprRef> {
1313
if let Some(binary_expr) = physical_expr
@@ -41,12 +41,12 @@ pub fn convert_expr_to_vortex(physical_expr: Arc<dyn PhysicalExpr>) -> VortexRes
4141
));
4242
}
4343

44-
if let Some(lit) = physical_expr
44+
if let Some(literal) = physical_expr
4545
.as_any()
4646
.downcast_ref::<expressions::Literal>()
4747
{
48-
let value = Scalar::from(lit.value().clone());
49-
return Ok(Literal::new_expr(value));
48+
let value = Scalar::from(literal.value().clone());
49+
return Ok(lit(value));
5050
}
5151

5252
vortex_bail!(

‎vortex-expr/src/identity.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,3 +40,8 @@ impl PartialEq<dyn Any> for Identity {
4040
.unwrap_or(false)
4141
}
4242
}
43+
44+
// Return a global pointer to the identity token.
45+
pub fn ident() -> ExprRef {
46+
Identity::new_expr()
47+
}

‎vortex-expr/src/lib.rs‎

Lines changed: 27 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -96,69 +96,68 @@ mod tests {
9696

9797
#[test]
9898
fn basic_expr_split_test() {
99-
let lhs = Column::new_expr(Field::from("a"));
100-
let rhs = Literal::new_expr(1.into());
101-
let expr = BinaryExpr::new_expr(lhs, Operator::Eq, rhs);
99+
let lhs = col("a");
100+
let rhs = lit(1);
101+
let expr = eq(lhs, rhs);
102102
let conjunction = split_conjunction(&expr);
103103
assert_eq!(conjunction.len(), 1);
104104
}
105105

106106
#[test]
107107
fn basic_conjunction_split_test() {
108-
let lhs = Column::new_expr(Field::from("a"));
109-
let rhs = Literal::new_expr(1.into());
110-
let expr = BinaryExpr::new_expr(lhs, Operator::And, rhs);
108+
let lhs = col("a");
109+
let rhs = lit(1);
110+
let expr = and(lhs, rhs);
111111
let conjunction = split_conjunction(&expr);
112112
assert_eq!(conjunction.len(), 2, "Conjunction is {conjunction:?}");
113113
}
114114

115115
#[test]
116116
fn expr_display() {
117-
assert_eq!(Column::new_expr(Field::from("a")).to_string(), "$a");
118-
assert_eq!(Column::new_expr(Field::Index(1)).to_string(), "[1]");
117+
assert_eq!(col("a").to_string(), "$a");
118+
assert_eq!(col(1).to_string(), "[1]");
119119
assert_eq!(Identity.to_string(), "[]");
120120
assert_eq!(Identity.to_string(), "[]");
121121

122-
let col1: Arc<dyn VortexExpr> = Column::new_expr(Field::from("col1"));
123-
let col2: Arc<dyn VortexExpr> = Column::new_expr(Field::from("col2"));
122+
let col1: Arc<dyn VortexExpr> = col("col1");
123+
let col2: Arc<dyn VortexExpr> = col("col2");
124124
assert_eq!(
125-
BinaryExpr::new_expr(col1.clone(), Operator::And, col2.clone()).to_string(),
125+
and(col1.clone(), col2.clone()).to_string(),
126126
"($col1 and $col2)"
127127
);
128128
assert_eq!(
129-
BinaryExpr::new_expr(col1.clone(), Operator::Or, col2.clone()).to_string(),
129+
or(col1.clone(), col2.clone()).to_string(),
130130
"($col1 or $col2)"
131131
);
132132
assert_eq!(
133-
BinaryExpr::new_expr(col1.clone(), Operator::Eq, col2.clone()).to_string(),
133+
eq(col1.clone(), col2.clone()).to_string(),
134134
"($col1 = $col2)"
135135
);
136136
assert_eq!(
137-
BinaryExpr::new_expr(col1.clone(), Operator::NotEq, col2.clone()).to_string(),
137+
not_eq(col1.clone(), col2.clone()).to_string(),
138138
"($col1 != $col2)"
139139
);
140140
assert_eq!(
141-
BinaryExpr::new_expr(col1.clone(), Operator::Gt, col2.clone()).to_string(),
141+
gt(col1.clone(), col2.clone()).to_string(),
142142
"($col1 > $col2)"
143143
);
144144
assert_eq!(
145-
BinaryExpr::new_expr(col1.clone(), Operator::Gte, col2.clone()).to_string(),
145+
gt_eq(col1.clone(), col2.clone()).to_string(),
146146
"($col1 >= $col2)"
147147
);
148148
assert_eq!(
149-
BinaryExpr::new_expr(col1.clone(), Operator::Lt, col2.clone()).to_string(),
149+
lt(col1.clone(), col2.clone()).to_string(),
150150
"($col1 < $col2)"
151151
);
152152
assert_eq!(
153-
BinaryExpr::new_expr(col1.clone(), Operator::Lte, col2.clone()).to_string(),
153+
lt_eq(col1.clone(), col2.clone()).to_string(),
154154
"($col1 <= $col2)"
155155
);
156156

157157
assert_eq!(
158-
BinaryExpr::new_expr(
159-
BinaryExpr::new_expr(col1.clone(), Operator::Lt, col2.clone()),
160-
Operator::Or,
161-
BinaryExpr::new_expr(col1.clone(), Operator::NotEq, col2.clone())
158+
or(
159+
lt(col1.clone(), col2.clone()),
160+
not_eq(col1.clone(), col2.clone()),
162161
)
163162
.to_string(),
164163
"(($col1 < $col2) or ($col1 != $col2))"
@@ -184,23 +183,20 @@ mod tests {
184183
"Exclude($col1,$col2,[1])"
185184
);
186185

187-
assert_eq!(Literal::new_expr(Scalar::from(0_u8)).to_string(), "0_u8");
186+
assert_eq!(lit(Scalar::from(0_u8)).to_string(), "0_u8");
187+
assert_eq!(lit(Scalar::from(0.0_f32)).to_string(), "0_f32");
188188
assert_eq!(
189-
Literal::new_expr(Scalar::from(0.0_f32)).to_string(),
190-
"0_f32"
191-
);
192-
assert_eq!(
193-
Literal::new_expr(Scalar::from(i64::MAX)).to_string(),
189+
lit(Scalar::from(i64::MAX)).to_string(),
194190
"9223372036854775807_i64"
195191
);
196-
assert_eq!(Literal::new_expr(Scalar::from(true)).to_string(), "true");
192+
assert_eq!(lit(Scalar::from(true)).to_string(), "true");
197193
assert_eq!(
198-
Literal::new_expr(Scalar::null(DType::Bool(Nullability::Nullable))).to_string(),
194+
lit(Scalar::null(DType::Bool(Nullability::Nullable))).to_string(),
199195
"null"
200196
);
201197

202198
assert_eq!(
203-
Literal::new_expr(Scalar::struct_(
199+
lit(Scalar::struct_(
204200
DType::Struct(
205201
StructDType::new(
206202
Arc::from([Arc::from("dog"), Arc::from("cat")]),

0 commit comments

Comments
 (0)