From 9296dac1a5eca62d1cbf188bcb16bc156459b838 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 6 Aug 2026 19:04:55 -0700 Subject: [PATCH] flatten BoundKind into BoundExpression Signed-off-by: Matt Katz --- vortex-array/src/expr/bound_expression.rs | 109 +++++++++++----------- vortex-array/src/expr/display.rs | 13 ++- vortex-array/src/expr/traversal/mod.rs | 17 ++-- vortex-array/src/expression.rs | 6 +- 4 files changed, 72 insertions(+), 73 deletions(-) diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index d93da7d2580..51aa5917dec 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -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. @@ -47,7 +41,10 @@ pub enum BoundKind { children: Arc>, }, /// 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. @@ -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, } } } @@ -83,11 +88,12 @@ impl Hash for ExactBoundExpr { fn hash(&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); @@ -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. @@ -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(), }) } @@ -140,7 +141,7 @@ impl BoundExpression { children: impl IntoIterator, ) -> VortexResult { 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", @@ -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. @@ -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); @@ -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))); } @@ -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("$"), } } } @@ -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 { @@ -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); @@ -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") }; diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index 202ec857bd0..a7642325b67 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -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; @@ -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}"), } } } diff --git a/vortex-array/src/expr/traversal/mod.rs b/vortex-array/src/expr/traversal/mod.rs index b97ae2536f7..952a73f2657 100644 --- a/vortex-array/src/expr/traversal/mod.rs +++ b/vortex-array/src/expr/traversal/mod.rs @@ -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 @@ -534,7 +533,7 @@ impl Node for BoundExpression { &'a self, mut f: F, ) -> VortexResult { - let BoundKind::Scalar { children, .. } = self.kind() else { + let BoundExpression::Scalar { children, .. } = self else { return Ok(TraversalOrder::Continue); }; @@ -552,7 +551,7 @@ impl Node for BoundExpression { self, mut f: F, ) -> VortexResult> { - let BoundKind::Scalar { children, .. } = self.kind() else { + let BoundExpression::Scalar { children, .. } = &self else { return Ok(Transformed::no(self)); }; @@ -583,16 +582,16 @@ impl Node for BoundExpression { } fn iter_children(&self, f: impl FnOnce(&mut dyn Iterator) -> 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, } } } diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index a1a46196766..d0590f2bf58 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -10,7 +10,6 @@ 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; @@ -18,10 +17,11 @@ 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 { - let BoundKind::Scalar { + let BoundExpression::Scalar { scalar_fn, children, - } = expr.kind() + .. + } = expr else { return Ok(self); };