Skip to content

Pushdown some expressions to Dict layout reader #8341

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
myrrc wants to merge 1 commit into
base: develop
Choose a base branch
from
+145 −7
Open

Pushdown some expressions to Dict layout reader #8341

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions vortex-array/src/scalar_fn/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,3 +48,18 @@ mod sealed {
/// This can be the **only** implementor for [`super::typed::DynScalarFn`].
impl<V: ScalarFnVTable> Sealed for TypedScalarFnInstance<V> {}
}

/*
* A scalar function has a negative cost if applying it to an array and
* canonicalizing is cheaper than canonicalizing an array and applying it.
*
* Example of negative cost expressions are byte_length() and get_item() since
* they don't depend on input size.
*
* Example of non-negative cost expression is like()
*/
pub fn is_negative_cost(id: ScalarFnId) -> bool {
id == Id::new_static("vortex.byte_length")
|| id == Id::new_static("vortex.get_item")
|| id == Id::new_static("vortex.literal")
}
137 changes: 130 additions & 7 deletions vortex-layout/src/layouts/dict/reader.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,15 @@ use vortex_array::arrays::SharedArray;
use vortex_array::dtype::DType;
use vortex_array::dtype::FieldMask;
use vortex_array::expr::Expression;
use vortex_array::expr::is_root;
use vortex_array::expr::label_is_fallible;
use vortex_array::expr::label_null_sensitive;
use vortex_array::expr::root;
use vortex_array::expr::traversal::NodeExt;
use vortex_array::expr::traversal::Transformed;
use vortex_array::expr::traversal::TraversalOrder;
use vortex_array::optimizer::ArrayOptimizer;
use vortex_array::scalar_fn::is_negative_cost;
use vortex_error::VortexError;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
Expand Down Expand Up @@ -100,10 +107,7 @@ impl DictReader {
)
.vortex_expect("must construct dict values array evaluation")
.map_err(Arc::new)
.map(move |array| {
let array = array?;
Ok(SharedArray::new(array).into_array())
})
.map(move |array| Ok(SharedArray::new(array?).into_array()))
.boxed()
.shared()
})
Expand Down Expand Up @@ -155,6 +159,49 @@ impl DictReader {
}
}

fn references_root(expr: &Expression) -> bool {
is_root(expr) || expr.children().iter().any(references_root)
}

/// Split expression into two parts:
///
/// left is the optional outer part that we want to apply to array after
/// canonicalizing.
/// right is the optional inner part that we want to apply to array before
/// canonicalizing.
///
/// We want to push to array only if expression has a negative cost, is
/// infallible and null-insensitive.
fn split_expression_for_pushdown(expr: Expression) -> (Option<Expression>, Option<Expression>) {
let labelled_expr = expr.clone();
let fallible = label_is_fallible(&labelled_expr);
let null_sensitive = label_null_sensitive(&labelled_expr);
let mut inner: Option<Expression> = None;

let outer = expr
.transform_down(|node| {
if is_negative_cost(node.id())
&& references_root(&node)
&& !fallible.get(&node).copied().unwrap_or(true)
&& !null_sensitive.get(&node).copied().unwrap_or(true)
{
inner = Some(node);
Ok(Transformed {
value: root(),
changed: true,
order: TraversalOrder::Skip,
})
} else {
Ok(Transformed::no(node))
}
})
.vortex_expect("infallible")
.into_inner();

let outer = (!is_root(&outer)).then_some(outer);
(outer, inner)
}

impl LayoutReader for DictReader {
fn name(&self) -> &Arc<str> {
&self.name
Expand Down Expand Up @@ -229,13 +276,18 @@ impl LayoutReader for DictReader {
mask: MaskFuture,
) -> VortexResult<BoxFuture<'static, VortexResult<ArrayRef>>> {
// TODO: fix up expr partitioning with fallible & null sensitive annotations
let values_eval = self.values_array();
let codes_eval = self
.codes
.projection_evaluation(row_range, &root(), mask)
.map_err(|err| err.with_context("While evaluating projection on codes"))?;
let expr = expr.clone();

let (expr_outer, expr_inner) = split_expression_for_pushdown(expr.clone());

let values_eval = if let Some(inner) = expr_inner {
self.values_eval(inner)
} else {
self.values_array()
};
let all_values_referenced = self.layout.has_all_values_referenced();
Ok(async move {
let (values, codes) = try_join!(values_eval.map_err(VortexError::from), codes_eval)?;
Expand All @@ -252,7 +304,11 @@ impl LayoutReader for DictReader {
.into_array()
.optimize()?;

array.apply(&expr)
if let Some(expr) = expr_outer {
array.apply(&expr)
} else {
Ok(array)
}
}
.boxed())
}
Expand Down Expand Up @@ -281,11 +337,20 @@ mod tests {
use vortex_array::dtype::FieldName;
use vortex_array::dtype::FieldNames;
use vortex_array::dtype::Nullability;
use vortex_array::dtype::PType;
use vortex_array::expr::Expression;
use vortex_array::expr::byte_length;
use vortex_array::expr::cast;
use vortex_array::expr::eq;
use vortex_array::expr::is_not_null;
use vortex_array::expr::is_root;
use vortex_array::expr::like;
use vortex_array::expr::lit;
use vortex_array::expr::pack;
use vortex_array::expr::root;
use vortex_array::expr::traversal::NodeExt;
use vortex_array::expr::traversal::Transformed;
use vortex_array::expr::traversal::TraversalOrder;
use vortex_array::scalar_fn::session::ScalarFnSession;
use vortex_array::session::ArraySession;
use vortex_array::validity::Validity;
Expand All @@ -296,6 +361,7 @@ mod tests {
use vortex_io::session::RuntimeSessionExt;
use vortex_session::VortexSession;

use super::split_expression_for_pushdown;
use crate::LayoutId;
use crate::LayoutRef;
use crate::LayoutStrategy;
Expand Down Expand Up @@ -542,4 +608,61 @@ mod tests {
assert_arrays_eq!(actual_canonical, expected);
})
}

fn join_split_expr(initial: &Expression, outer: Option<Expression>, inner: Option<Expression>) {
let outer_expr = outer.unwrap_or_else(root);
let inner_expr = inner.unwrap_or_else(root);
let expected = outer_expr
.transform_down(|node| {
if !is_root(&node) {
return Ok(Transformed::no(node));
}
Ok(Transformed {
value: inner_expr.clone(),
changed: true,
order: TraversalOrder::Skip,
})
})
.vortex_expect("infallible");
assert_eq!(&expected.into_inner(), initial);
}

#[test]
fn split_expr_cast_root() {
let (outer, inner) = split_expression_for_pushdown(root());
assert_eq!(outer, None);
assert_eq!(inner, None); // Applying root to array is useless work
}

#[test]
fn split_expr_partial_pushdown() {
let dtype = DType::Primitive(PType::U64, Nullability::NonNullable);
let expr = cast(byte_length(root()), dtype.clone());
let (outer, inner) = split_expression_for_pushdown(expr.clone());
// [0] = cast([1], dtype)
// [1] = byte_length(root)
assert_eq!(outer, Some(cast(root(), dtype)));
assert_eq!(inner, Some(byte_length(root())));
join_split_expr(&expr, outer, inner);
}

#[test]
fn split_expr_full_pushdown() {
let expr = byte_length(root());
let (outer, inner) = split_expression_for_pushdown(expr.clone());
assert_eq!(outer, None);
assert_eq!(inner, Some(byte_length(root())));
join_split_expr(&expr, outer, inner);
}

#[test]
fn split_expr_no_pushdown() {
// We can push down lit(), but it we replace
// lit() with root(), the semantics change.
let expr = like(root(), lit(1u64));
let (outer, inner) = split_expression_for_pushdown(expr.clone());
assert_eq!(outer, Some(expr.clone()));
assert_eq!(inner, None);
join_split_expr(&expr, outer, inner);
}
}
Loading