Skip to content

Commit fe47467

Browse files
connortsui20connortsui20
committed
fix expression pushdown
Signed-off-by: Connor Tsui <[email protected]>
1 parent 33571c2 commit fe47467

7 files changed

Lines changed: 320 additions & 41 deletions

File tree

‎vortex-array/src/arrays/scalar_fn/array.rs‎

Lines changed: 13 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,9 @@ use crate::arrays::Extension;
1919
use crate::arrays::ScalarFn;
2020
use crate::arrays::extension::ExtensionArrayExt;
2121
use crate::dtype::DType;
22+
use crate::scalar_fn::RefinementFallbackArg;
2223
use crate::scalar_fn::ScalarFnRef;
24+
use crate::scalar_fn::resolve_return_dtype_with_refinement_fallback;
2325

2426
// ScalarFnArray has a variable number of slots (one per child)
2527

@@ -164,37 +166,21 @@ pub(crate) fn peel_refinements_and_resolve_dtype(
164166
scalar_fn: &ScalarFnRef,
165167
children: &mut [ArrayRef],
166168
) -> VortexResult<DType> {
167-
loop {
168-
let children_dtypes: Vec<_> = children.iter().map(|c| c.dtype().clone()).collect();
169-
170-
match scalar_fn.return_dtype(&children_dtypes) {
171-
Ok(dtype) => return Ok(dtype),
172-
Err(err) => {
173-
let any_peeled = peel_refinement_layers(children);
174-
if !any_peeled {
175-
return Err(err);
176-
}
177-
}
178-
}
179-
}
169+
resolve_return_dtype_with_refinement_fallback(scalar_fn, children)
180170
}
181171

182-
// TODO(connor): Is it correct/ok to peel away refinement on all children at once? Do we need
183-
// small-step semantics here?
184-
/// Peels one layer of refinement extensions from all `children`.
185-
///
186-
/// Returns a flag indicating whether any child was actually peeled.
187-
fn peel_refinement_layers(children: &mut [ArrayRef]) -> bool {
188-
let mut any_peeled = false;
189-
190-
for child in children.iter_mut() {
191-
if let Some(peeled) = peel_refinement_child(child) {
192-
*child = peeled;
193-
any_peeled = true;
194-
}
172+
impl RefinementFallbackArg for ArrayRef {
173+
fn current_dtype(&self) -> &DType {
174+
self.dtype()
195175
}
196176

197-
any_peeled
177+
fn peel_one_refinement_layer(&mut self) -> bool {
178+
let Some(peeled) = peel_refinement_child(self) else {
179+
return false;
180+
};
181+
*self = peeled;
182+
true
183+
}
198184
}
199185

200186
/// Peel exactly one refinement wrapper when the array is a real representation wrapper.

‎vortex-array/src/expr/expression.rs‎

Lines changed: 81 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,11 @@ use crate::dtype::DType;
1717
use crate::expr::StatsCatalog;
1818
use crate::expr::display::DisplayTreeExpr;
1919
use crate::expr::stats::Stat;
20+
use crate::scalar_fn::RefinementFallbackArg;
2021
use crate::scalar_fn::ScalarFnRef;
22+
use crate::scalar_fn::fns::literal::Literal;
2123
use crate::scalar_fn::fns::root::Root;
24+
use crate::scalar_fn::resolve_return_dtype_with_refinement_fallback;
2225

2326
/// A node in a Vortex expression tree.
2427
///
@@ -94,16 +97,27 @@ impl Expression {
9497

9598
/// Computes the return dtype of this expression given the input dtype.
9699
pub fn return_dtype(&self, scope: &DType) -> VortexResult<DType> {
100+
Ok(self.return_dtype_info(scope)?.dtype)
101+
}
102+
103+
pub(crate) fn return_dtype_info(&self, scope: &DType) -> VortexResult<ExprReturnInfo> {
97104
if self.is::<Root>() {
98-
return Ok(scope.clone());
105+
return Ok(ExprReturnInfo::storage_chain_repr_wrapper(scope.clone()));
106+
}
107+
108+
if let Some(literal) = self.as_opt::<Literal>() {
109+
return Ok(ExprReturnInfo::storage_chain_repr_wrapper(
110+
literal.dtype().clone(),
111+
));
99112
}
100113

101-
let dtypes: Vec<_> = self
114+
let mut children: Vec<_> = self
102115
.children
103116
.iter()
104-
.map(|c| c.return_dtype(scope))
117+
.map(|c| c.return_dtype_info(scope))
105118
.try_collect()?;
106-
self.scalar_fn.return_dtype(&dtypes)
119+
120+
expr_return_dtype_info(&self.scalar_fn, &mut children)
107121
}
108122

109123
/// Returns a new expression representing the validity mask output of this expression.
@@ -243,3 +257,66 @@ impl Drop for Expression {
243257
}
244258
}
245259
}
260+
261+
#[derive(Clone, Debug)]
262+
pub(crate) struct ExprReturnInfo {
263+
pub(crate) dtype: DType,
264+
peel_strategy: ExprPeelStrategy,
265+
}
266+
267+
#[derive(Clone, Copy, Debug, Default)]
268+
enum ExprPeelStrategy {
269+
#[default]
270+
None,
271+
StorageChainRepresentationWrapper,
272+
}
273+
274+
impl ExprReturnInfo {
275+
pub(crate) fn opaque(dtype: DType) -> Self {
276+
Self {
277+
dtype,
278+
peel_strategy: ExprPeelStrategy::None,
279+
}
280+
}
281+
282+
pub(crate) fn storage_chain_repr_wrapper(dtype: DType) -> Self {
283+
Self {
284+
dtype,
285+
peel_strategy: ExprPeelStrategy::StorageChainRepresentationWrapper,
286+
}
287+
}
288+
}
289+
290+
impl RefinementFallbackArg for ExprReturnInfo {
291+
fn current_dtype(&self) -> &DType {
292+
&self.dtype
293+
}
294+
295+
fn peel_one_refinement_layer(&mut self) -> bool {
296+
if !matches!(
297+
self.peel_strategy,
298+
ExprPeelStrategy::StorageChainRepresentationWrapper
299+
) {
300+
return false;
301+
}
302+
303+
let DType::Extension(ext_dtype) = &self.dtype else {
304+
return false;
305+
};
306+
if !ext_dtype.is_refinement() {
307+
return false;
308+
}
309+
310+
self.dtype = ext_dtype.storage_dtype().clone();
311+
true
312+
}
313+
}
314+
315+
pub(crate) fn expr_return_dtype_info(
316+
scalar_fn: &ScalarFnRef,
317+
children: &mut [ExprReturnInfo],
318+
) -> VortexResult<ExprReturnInfo> {
319+
Ok(ExprReturnInfo::opaque(
320+
resolve_return_dtype_with_refinement_fallback(scalar_fn, children)?,
321+
))
322+
}

‎vortex-array/src/expr/mod.rs‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,9 @@ mod tests {
127127
use crate::dtype::Nullability;
128128
use crate::dtype::PType;
129129
use crate::dtype::StructFields;
130+
use crate::dtype::extension::ExtDType;
130131
use crate::expr::and;
132+
use crate::expr::checked_add;
131133
use crate::expr::col;
132134
use crate::expr::eq;
133135
use crate::expr::get_item;
@@ -142,6 +144,8 @@ mod tests {
142144
use crate::expr::root;
143145
use crate::expr::select;
144146
use crate::expr::select_exclude;
147+
use crate::extension::tests::divisible_int::DivisibleInt;
148+
use crate::extension::tests::divisible_int::Divisor;
145149
use crate::scalar::Scalar;
146150

147151
#[test]
@@ -260,4 +264,34 @@ mod tests {
260264
"{dog: 32u32, cat: \"rufus\"}"
261265
);
262266
}
267+
268+
fn divisible_int_dtype(divisor: u64) -> DType {
269+
DType::Extension(
270+
ExtDType::<DivisibleInt>::try_new(
271+
Divisor(divisor),
272+
DType::Primitive(PType::U64, Nullability::NonNullable),
273+
)
274+
.unwrap()
275+
.erased(),
276+
)
277+
}
278+
279+
#[test]
280+
fn return_dtype_peels_refinement_root_and_literal() {
281+
let scope = divisible_int_dtype(3);
282+
283+
assert_eq!(
284+
checked_add(root(), root()).return_dtype(&scope).unwrap(),
285+
DType::Primitive(PType::U64, Nullability::NonNullable),
286+
);
287+
288+
let ext_dtype = scope.as_extension().clone();
289+
let refined_literal = Scalar::extension_ref(ext_dtype, Scalar::from(3u64));
290+
assert_eq!(
291+
checked_add(root(), lit(refined_literal))
292+
.return_dtype(&scope)
293+
.unwrap(),
294+
DType::Primitive(PType::U64, Nullability::NonNullable),
295+
);
296+
}
263297
}

‎vortex-array/src/expr/optimize.rs‎

Lines changed: 64 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33

44
use std::any::Any;
55
use std::cell::RefCell;
6-
use std::ops::Deref;
76
use std::sync::Arc;
87

98
use itertools::Itertools;
@@ -13,12 +12,15 @@ use vortex_utils::aliases::hash_map::HashMap;
1312

1413
use crate::dtype::DType;
1514
use crate::expr::Expression;
15+
use crate::expr::expression::ExprReturnInfo;
16+
use crate::expr::expression::expr_return_dtype_info;
1617
use crate::expr::transform::match_between::find_between;
1718
use crate::scalar_fn::ReduceCtx;
1819
use crate::scalar_fn::ReduceNode;
1920
use crate::scalar_fn::ReduceNodeRef;
2021
use crate::scalar_fn::ScalarFnRef;
2122
use crate::scalar_fn::SimplifyCtx;
23+
use crate::scalar_fn::fns::literal::Literal;
2224
use crate::scalar_fn::fns::root::Root;
2325

2426
impl Expression {
@@ -206,27 +208,34 @@ impl Expression {
206208

207209
struct SimplifyCache<'a> {
208210
scope: &'a DType,
209-
dtype_cache: RefCell<HashMap<Expression, DType>>,
211+
dtype_cache: RefCell<HashMap<Expression, ExprReturnInfo>>,
210212
}
211213

212-
impl SimplifyCtx for SimplifyCache<'_> {
213-
fn return_dtype(&self, expr: &Expression) -> VortexResult<DType> {
214-
// If the expression is "root", return the scope dtype
214+
impl SimplifyCache<'_> {
215+
fn return_dtype_info(&self, expr: &Expression) -> VortexResult<ExprReturnInfo> {
215216
if expr.is::<Root>() {
216-
return Ok(self.scope.clone());
217+
return Ok(ExprReturnInfo::storage_chain_repr_wrapper(
218+
self.scope.clone(),
219+
));
220+
}
221+
222+
if let Some(literal) = expr.as_opt::<Literal>() {
223+
return Ok(ExprReturnInfo::storage_chain_repr_wrapper(
224+
literal.dtype().clone(),
225+
));
217226
}
218227

219228
if let Some(dtype) = self.dtype_cache.borrow().get(expr) {
220229
return Ok(dtype.clone());
221230
}
222231

223232
// Otherwise, compute dtype from children
224-
let input_dtypes: Vec<_> = expr
233+
let mut input_dtypes: Vec<_> = expr
225234
.children()
226235
.iter()
227-
.map(|c| self.return_dtype(c))
236+
.map(|c| self.return_dtype_info(c))
228237
.try_collect()?;
229-
let dtype = expr.deref().return_dtype(&input_dtypes)?;
238+
let dtype = expr_return_dtype_info(expr.scalar_fn(), &mut input_dtypes)?;
230239
self.dtype_cache
231240
.borrow_mut()
232241
.insert(expr.clone(), dtype.clone());
@@ -235,6 +244,12 @@ impl SimplifyCtx for SimplifyCache<'_> {
235244
}
236245
}
237246

247+
impl SimplifyCtx for SimplifyCache<'_> {
248+
fn return_dtype(&self, expr: &Expression) -> VortexResult<DType> {
249+
Ok(self.return_dtype_info(expr)?.dtype)
250+
}
251+
}
252+
238253
struct ExpressionReduceNode {
239254
expression: Expression,
240255
scope: DType,
@@ -294,3 +309,43 @@ impl ReduceCtx for ExpressionReduceCtx {
294309
}))
295310
}
296311
}
312+
313+
#[cfg(test)]
314+
mod tests {
315+
use crate::dtype::DType;
316+
use crate::dtype::Nullability;
317+
use crate::dtype::PType;
318+
use crate::dtype::extension::ExtDType;
319+
use crate::expr::checked_add;
320+
use crate::expr::fill_null;
321+
use crate::expr::lit;
322+
use crate::expr::root;
323+
use crate::extension::tests::divisible_int::DivisibleInt;
324+
use crate::extension::tests::divisible_int::Divisor;
325+
326+
fn divisible_int_dtype(divisor: u64) -> DType {
327+
DType::Extension(
328+
ExtDType::<DivisibleInt>::try_new(
329+
Divisor(divisor),
330+
DType::Primitive(PType::U64, Nullability::NonNullable),
331+
)
332+
.unwrap()
333+
.erased(),
334+
)
335+
}
336+
337+
#[test]
338+
fn optimize_recursive_uses_refinement_fallback_for_typed_simplify() {
339+
let scope = divisible_int_dtype(3);
340+
let add = checked_add(root(), root());
341+
let expr = fill_null(add.clone(), lit(0u64));
342+
343+
let optimized = expr.optimize_recursive(&scope).unwrap();
344+
345+
assert_eq!(optimized, add);
346+
assert_eq!(
347+
optimized.return_dtype(&scope).unwrap(),
348+
DType::Primitive(PType::U64, Nullability::NonNullable),
349+
);
350+
}
351+
}

‎vortex-array/src/scalar_fn/mod.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@ pub use typed::*;
2424
mod erased;
2525
pub use erased::*;
2626

27+
mod refinement;
28+
pub(crate) use refinement::*;
29+
2730
mod options;
2831
pub use options::*;
2932

‎vortex-array/src/scalar_fn/refinement.rs‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
use vortex_error::VortexResult;
5+
6+
use crate::dtype::DType;
7+
use crate::scalar_fn::ScalarFnRef;
8+
9+
/// An argument that may be able to peel one refinement layer and retry scalar-fn dtype
10+
/// resolution on its storage dtype.
11+
pub(crate) trait RefinementFallbackArg {
12+
fn current_dtype(&self) -> &DType;
13+
14+
fn peel_one_refinement_layer(&mut self) -> bool;
15+
}
16+
17+
/// Resolve a scalar function's return dtype, retrying after peeling one refinement layer from all
18+
/// currently peelable children when the original argument dtypes are rejected.
19+
pub(crate) fn resolve_return_dtype_with_refinement_fallback<A: RefinementFallbackArg>(
20+
scalar_fn: &ScalarFnRef,
21+
args: &mut [A],
22+
) -> VortexResult<DType> {
23+
loop {
24+
let arg_dtypes: Vec<_> = args.iter().map(|arg| arg.current_dtype().clone()).collect();
25+
26+
match scalar_fn.return_dtype(&arg_dtypes) {
27+
Ok(dtype) => return Ok(dtype),
28+
Err(err) => {
29+
let mut any_peeled = false;
30+
for arg in args.iter_mut() {
31+
any_peeled |= arg.peel_one_refinement_layer();
32+
}
33+
if !any_peeled {
34+
return Err(err);
35+
}
36+
}
37+
}
38+
}
39+
}

0 commit comments

Comments
 (0)