Skip to content

Push down some expressions to Dict layout reader's cached values #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
+149 −7
Open

Push down some expressions to Dict layout reader's cached values #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
19 changes: 19 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 @@ -9,6 +9,10 @@

use vortex_session::registry::Id;

use crate::scalar_fn::fns::byte_length::ByteLength;
use crate::scalar_fn::fns::get_item::GetItem;
use crate::scalar_fn::fns::literal::Literal;

mod vtable;
pub use vtable::*;

Expand Down Expand Up @@ -48,3 +52,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 == ScalarFnVTable::id(&ByteLength)
|| id == ScalarFnVTable::id(&GetItem)
|| id == ScalarFnVTable::id(&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,
Comment on lines +175 to +192

@joseph-isaacs joseph-isaacs Jun 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we have a partition function for exactly this!

@myrrc myrrc Jun 12, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can't implement this using partition<> because it pushes children annotations up to the parents. This breaks on

let expr = cast(byte_length(root()), dtype);
// expect inner = byte_length(root), outer = cast(root, dtype)

But instead we get

inner = cast(...)

because cast inherits the annotation from a child. This is also an implementation that's harder to reason about. Let's leave this one.

@joseph-isaacs joseph-isaacs Jun 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you know which part incorrect with this method?

})
} 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