Skip to content

Flatten BoundKind into BoundExpression #9256

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
mhk197 wants to merge 1 commit into from
Open

Flatten BoundKind into BoundExpression #9256

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
109 changes: 55 additions & 54 deletions vortex-array/src/expr/bound_expression.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,17 +27,11 @@ use crate::scalar_fn::ScalarFnRef;
/// Binding is purely logical: it deals only in [`DType`]s and never sees an array, a length, or an
/// encoding.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct BoundExpression {
kind: BoundKind,
dtype: DType,
}

/// The per-variant contents of a [`BoundExpression`], mirroring the logical variants of
/// [`Expression`].
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum BoundKind {
pub enum BoundExpression {
/// A scalar function applied to bound children.
Scalar {
/// The dtype this node evaluates to.
dtype: DType,
/// The scalar function for this node.
scalar_fn: ScalarFnRef,
/// The bound children, in argument order.
Expand All @@ -47,7 +41,10 @@ pub enum BoundKind {
children: Arc<Vec<BoundExpression>>,
},
/// The scope itself. Its dtype is the scope's root dtype.
Root,
Root {
/// The dtype this node evaluates to.
dtype: DType,
},
}

/// A bound-expression wrapper that compares shared tree identity instead of structure.
Expand All @@ -56,23 +53,31 @@ pub struct ExactBoundExpr(pub BoundExpression);

impl PartialEq for ExactBoundExpr {
fn eq(&self, other: &Self) -> bool {
match (&self.0.kind, &other.0.kind) {
(BoundKind::Root, BoundKind::Root) => self.0.dtype == other.0.dtype,
match (&self.0, &other.0) {
(
BoundExpression::Root { dtype: lhs_dtype },
BoundExpression::Root { dtype: rhs_dtype },
) => lhs_dtype == rhs_dtype,
(
BoundKind::Scalar {
BoundExpression::Scalar {
dtype: lhs_dtype,
scalar_fn: lhs_fn,
children: lhs_children,
},
BoundKind::Scalar {
BoundExpression::Scalar {
dtype: rhs_dtype,
scalar_fn: rhs_fn,
children: rhs_children,
},
) => {
lhs_fn == rhs_fn
&& Arc::ptr_eq(lhs_children, rhs_children)
&& self.0.dtype == other.0.dtype
&& lhs_dtype == rhs_dtype
}
_ => false,
// No catch-all: a new variant must state its own identity rather than silently
// comparing unequal, which would put `eq` out of step with `hash`.
(BoundExpression::Root { .. }, BoundExpression::Scalar { .. })
| (BoundExpression::Scalar { .. }, BoundExpression::Root { .. }) => false,
}
}
}
Expand All @@ -83,11 +88,12 @@ impl Hash for ExactBoundExpr {
fn hash<H: Hasher>(&self, state: &mut H) {
// DType differences are resolved by equality. Omitting the potentially lazy dtype keeps
// identity-keyed cache lookups from deserializing an entire schema just to compute a hash.
match &self.0.kind {
BoundKind::Root => state.write_u8(0),
BoundKind::Scalar {
match &self.0 {
BoundExpression::Root { .. } => state.write_u8(0),
BoundExpression::Scalar {
scalar_fn,
children,
..
} => {
state.write_u8(1);
scalar_fn.hash(state);
Expand All @@ -100,10 +106,7 @@ impl Hash for ExactBoundExpr {
impl BoundExpression {
/// Create a bound root expression with the given dtype.
pub fn new_root(dtype: DType) -> Self {
Self {
kind: BoundKind::Root,
dtype,
}
Self::Root { dtype }
}

/// Create a bound scalar node from a scalar function and already-bound children.
Expand All @@ -125,12 +128,10 @@ impl BoundExpression {
.collect_vec();
let dtype = scalar_fn.return_dtype(&arg_dtypes)?;

Ok(Self {
kind: BoundKind::Scalar {
scalar_fn,
children: children.into(),
},
Ok(Self::Scalar {
dtype,
scalar_fn,
children: children.into(),
})
}

Expand All @@ -140,7 +141,7 @@ impl BoundExpression {
children: impl IntoIterator<Item = BoundExpression>,
) -> VortexResult<Self> {
let children = Vec::from_iter(children);
let BoundKind::Scalar { scalar_fn, .. } = &self.kind else {
let BoundExpression::Scalar { scalar_fn, .. } = &self else {
vortex_ensure!(
children.is_empty(),
"Root expression cannot have {} children",
Expand All @@ -154,33 +155,30 @@ impl BoundExpression {

/// The dtype this expression evaluates to.
pub fn dtype(&self) -> &DType {
&self.dtype
}

/// The per-variant contents of this node.
pub fn kind(&self) -> &BoundKind {
&self.kind
match self {
Self::Scalar { dtype, .. } | Self::Root { dtype } => dtype,
}
}

/// The bound children of this node, in argument order. Empty for [`BoundKind::Root`].
/// The bound children of this node, in argument order. Empty for [`BoundExpression::Root`].
pub fn children(&self) -> &[BoundExpression] {
match &self.kind {
BoundKind::Scalar { children, .. } => children.as_slice(),
BoundKind::Root => &[],
match self {
Self::Scalar { children, .. } => children.as_slice(),
Self::Root { .. } => &[],
}
}

/// The scalar function for this node, or `None` if it is the scope root.
pub fn as_scalar(&self) -> Option<&ScalarFnRef> {
match &self.kind {
BoundKind::Scalar { scalar_fn, .. } => Some(scalar_fn),
BoundKind::Root => None,
match self {
Self::Scalar { scalar_fn, .. } => Some(scalar_fn),
Self::Root { .. } => None,
}
}

/// Whether this node is the scope root.
pub fn is_root(&self) -> bool {
matches!(self.kind, BoundKind::Root)
matches!(self, Self::Root { .. })
}

/// Display the bound expression as a formatted tree structure.
Expand All @@ -199,11 +197,12 @@ impl BoundExpression {
let mut expressions = Vec::new();

while let Some((node, visited)) = pending.pop() {
match node.kind() {
BoundKind::Root => expressions.push(crate::expr::root()),
BoundKind::Scalar {
match node {
BoundExpression::Root { .. } => expressions.push(crate::expr::root()),
BoundExpression::Scalar {
scalar_fn,
children,
..
} if visited => {
let child_start = expressions.len() - children.len();
let child_expressions = expressions.split_off(child_start);
Expand All @@ -212,7 +211,7 @@ impl BoundExpression {
.vortex_expect("a bound expression always has valid arity"),
);
}
BoundKind::Scalar { children, .. } => {
BoundExpression::Scalar { children, .. } => {
pending.push((node, true));
pending.extend(children.iter().rev().map(|child| (child, false)));
}
Expand All @@ -227,9 +226,9 @@ impl BoundExpression {

impl Display for BoundExpression {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
match self.kind() {
BoundKind::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f),
BoundKind::Root => f.write_str("$"),
match self {
Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f),
Self::Root { .. } => f.write_str("$"),
}
}
}
Expand Down Expand Up @@ -265,7 +264,7 @@ impl Expression {
/// Iterative drop to avoid stack overflows on deep trees.
impl Drop for BoundExpression {
fn drop(&mut self) {
let BoundKind::Scalar { children, .. } = &mut self.kind else {
let Self::Scalar { children, .. } = self else {
return;
};
let Some(children) = Arc::get_mut(children) else {
Expand All @@ -274,7 +273,7 @@ impl Drop for BoundExpression {

let mut to_drop = std::mem::take(children);
while let Some(mut child) = to_drop.pop() {
if let BoundKind::Scalar { children, .. } = &mut child.kind
if let BoundExpression::Scalar { children, .. } = &mut child
&& let Some(grandchildren) = Arc::get_mut(children)
{
to_drop.append(grandchildren);
Expand Down Expand Up @@ -355,8 +354,10 @@ mod tests {
let bound = eq(col("a"), lit(1_i32)).bind_scope(&scope())?;
let cloned = bound.clone();

let (BoundKind::Scalar { children: a, .. }, BoundKind::Scalar { children: b, .. }) =
(bound.kind(), cloned.kind())
let (
BoundExpression::Scalar { children: a, .. },
BoundExpression::Scalar { children: b, .. },
) = (&bound, &cloned)
else {
unreachable!("eq is a scalar node")
};
Expand Down
13 changes: 6 additions & 7 deletions vortex-array/src/expr/display.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@ use std::fmt::Display;
use std::fmt::Formatter;

use crate::expr::BoundExpression;
use crate::expr::BoundKind;
use crate::expr::Expression;
use crate::scalar_fn::ChildName;

Expand Down Expand Up @@ -84,16 +83,16 @@ impl DisplayTreeNode for BoundExpression {
}

fn tree_child_name(&self, index: usize) -> ChildName {
match self.kind() {
BoundKind::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index),
BoundKind::Root => unreachable!("the scope root has no children"),
match self {
BoundExpression::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index),
BoundExpression::Root { .. } => unreachable!("the scope root has no children"),
}
}

fn fmt_tree_node(&self, f: &mut Formatter<'_>) -> fmt::Result {
match self.kind() {
BoundKind::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f),
BoundKind::Root => write!(f, "{ROOT_DISPLAY}"),
match self {
BoundExpression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f),
BoundExpression::Root { .. } => write!(f, "{ROOT_DISPLAY}"),
}
}
}
Expand Down
17 changes: 8 additions & 9 deletions vortex-array/src/expr/traversal/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ use vortex_error::VortexResult;

use crate::expr::BoundExpression;
use crate::expr::Expression;
use crate::expr::bound_expression::BoundKind;
use crate::expr::traversal::fold::NodeFolderContextWrapper;

/// Signal to control a traversal's flow
Expand Down Expand Up @@ -534,7 +533,7 @@ impl Node for BoundExpression {
&'a self,
mut f: F,
) -> VortexResult<TraversalOrder> {
let BoundKind::Scalar { children, .. } = self.kind() else {
let BoundExpression::Scalar { children, .. } = self else {
return Ok(TraversalOrder::Continue);
};

Expand All @@ -552,7 +551,7 @@ impl Node for BoundExpression {
self,
mut f: F,
) -> VortexResult<Transformed<Self>> {
let BoundKind::Scalar { children, .. } = self.kind() else {
let BoundExpression::Scalar { children, .. } = &self else {
return Ok(Transformed::no(self));
};

Expand Down Expand Up @@ -583,16 +582,16 @@ impl Node for BoundExpression {
}

fn iter_children<T>(&self, f: impl FnOnce(&mut dyn Iterator<Item = &Self>) -> T) -> T {
match self.kind() {
BoundKind::Scalar { children, .. } => f(&mut children.iter()),
BoundKind::Root => f(&mut std::iter::empty()),
match self {
BoundExpression::Scalar { children, .. } => f(&mut children.iter()),
BoundExpression::Root { .. } => f(&mut std::iter::empty()),
}
}

fn children_count(&self) -> usize {
match self.kind() {
BoundKind::Scalar { children, .. } => children.len(),
BoundKind::Root => 0,
match self {
BoundExpression::Scalar { children, .. } => children.len(),
BoundExpression::Root { .. } => 0,
}
}
}
Expand Down
6 changes: 3 additions & 3 deletions vortex-array/src/expression.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,18 +10,18 @@ use crate::IntoArray;
use crate::arrays::ConstantArray;
use crate::arrays::ScalarFnArray;
use crate::expr::BoundExpression;
use crate::expr::BoundKind;
use crate::expr::Expression;
use crate::optimizer::ArrayOptimizer;
use crate::scalar_fn::fns::literal::Literal;

impl ArrayRef {
/// Apply a bound expression to this array, producing a new array in constant time.
pub fn apply_bound(self, expr: &BoundExpression) -> VortexResult<ArrayRef> {
let BoundKind::Scalar {
let BoundExpression::Scalar {
scalar_fn,
children,
} = expr.kind()
..
} = expr
else {
return Ok(self);
};
Expand Down
Loading