diff --git a/vortex-array/src/expr/analysis/immediate_access.rs b/vortex-array/src/expr/analysis/immediate_access.rs index 6c2e4975a92..eacece3489e 100644 --- a/vortex-array/src/expr/analysis/immediate_access.rs +++ b/vortex-array/src/expr/analysis/immediate_access.rs @@ -40,24 +40,29 @@ use crate::scalar_fn::fns::select::Select; pub fn make_free_field_annotator( scope: &StructFields, ) -> impl AnnotationFn { - move |expr: &Expression| { - if let Some(selection) = expr.as_opt::() { + if children[0].is_root() { + return selection + .normalize_to_included_fields(scope.names()) + .vortex_expect("Select fields must be valid for scope") + .into_iter() + .collect(); + } + } else if let Some(field_name) = scalar_fn.as_opt::() + && children[0].is_root() + { return vec![field_name.clone()]; } - } else if expr.is_root() { - return scope.names().iter().cloned().collect(); - } - vec![] + vec![] + } } } @@ -65,25 +70,58 @@ pub fn make_free_field_annotator( pub fn make_bound_free_field_annotator( scope: &StructFields, ) -> impl AnnotationFn { - move |expr: &BoundExpression| { - let Some(scalar_fn) = expr.as_scalar() else { - return scope.names().iter().cloned().collect(); - }; - - if let Some(selection) = scalar_fn.as_opt::() { + if children[0].is_root() { + return selection + .normalize_to_included_fields(scope.names()) + .vortex_expect("Select fields must be valid for scope") + .into_iter() + .collect(); + } + } else if let Some(field_name) = scalar_fn.as_opt::() + && children[0].is_root() + { + return vec![field_name.clone()]; } - } else if let Some(field_name) = scalar_fn.as_opt::() - && expr.children()[0].is_root() - { - return vec![field_name.clone()]; + + vec![] } + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::*; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::expr::Scope; + use crate::expr::Variable; + use crate::expr::var; + + #[test] + fn variables_do_not_access_root_fields() -> VortexResult<()> { + let fields = StructFields::from_iter([("a", DType::Null)]); + let expression = var("value"); + + assert!(make_free_field_annotator(&fields)(&expression).is_empty()); - vec![] + let root_dtype = DType::Struct(fields.clone(), Nullability::NonNullable); + let variable_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let scope = + Scope::new(root_dtype).with_bindings([(Variable::new("value"), variable_dtype)])?; + let bound = expression.bind_scope(&scope)?; + assert!(make_bound_free_field_annotator(&fields)(&bound).is_empty()); + Ok(()) } } diff --git a/vortex-array/src/expr/analysis/infallible.rs b/vortex-array/src/expr/analysis/infallible.rs index fbf84e124f0..46b7a5af697 100644 --- a/vortex-array/src/expr/analysis/infallible.rs +++ b/vortex-array/src/expr/analysis/infallible.rs @@ -17,6 +17,8 @@ pub fn label_infallible(expr: &Expression) -> BooleanLabels<'_> { Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().is_infallible(), // The scope itself cannot fail. Expression::Root => true, + // References are vacuously infallible. + Expression::Variable(_) => true, }, |acc, &child| acc & child, ) diff --git a/vortex-array/src/expr/analysis/strict.rs b/vortex-array/src/expr/analysis/strict.rs index 01cc71d91c3..5e72e5541c8 100644 --- a/vortex-array/src/expr/analysis/strict.rs +++ b/vortex-array/src/expr/analysis/strict.rs @@ -15,7 +15,7 @@ pub fn label_strict(expr: &Expression) -> BooleanLabels<'_> { |expr| match expr { Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().is_strict(), // Vacuously strict. - Expression::Root => true, + Expression::Root | Expression::Variable(_) => true, }, |acc, &child| acc & child, ) diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 4bd276e191d..f8be1d2c9e7 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -11,6 +11,7 @@ use std::sync::Arc; use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_session::VortexSession; @@ -18,8 +19,10 @@ use crate::dtype::DType; use crate::expr::Expression; use crate::expr::display::DisplayTreeExpr; use crate::expr::scope::Scope; +use crate::expr::scope::VariableRef; use crate::expr::traversal::TraversalOrder; use crate::expr::traversal::pre_order_visit_down; +use crate::expr::variable::Variable; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::ScalarFnVTable; use crate::stats::rewrite::StatsRewriteCtx; @@ -50,6 +53,39 @@ pub enum BoundExpression { /// The dtype this node evaluates to. dtype: DType, }, + /// A variable resolved to a dtype and stable location in the bound scope. + Variable(BoundVariable), +} + +/// A variable resolved in a lexical scope. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct BoundVariable { + dtype: DType, + variable: Variable, + variable_ref: VariableRef, +} + +impl BoundVariable { + /// The dtype of the value bound to this variable. + pub fn dtype(&self) -> &DType { + &self.dtype + } + + /// The source-level variable name. + pub fn variable(&self) -> &Variable { + &self.variable + } + + /// The variable's stable location in the bound scope. + pub fn variable_ref(&self) -> VariableRef { + self.variable_ref + } +} + +impl Display for BoundVariable { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + Display::fmt(&self.variable, f) + } } /// A bound-expression wrapper that compares shared tree identity instead of structure. @@ -79,7 +115,10 @@ impl PartialEq for ExactBoundExpr { && Arc::ptr_eq(lhs_children, rhs_children) && lhs_dtype == rhs_dtype } - _ => false, + (BoundExpression::Variable(lhs), BoundExpression::Variable(rhs)) => lhs == rhs, + (BoundExpression::Root { .. }, _) + | (BoundExpression::Scalar { .. }, _) + | (BoundExpression::Variable(_), _) => false, } } } @@ -92,6 +131,11 @@ impl Hash for ExactBoundExpr { // identity-keyed cache lookups from deserializing an entire schema just to compute a hash. match &self.0 { BoundExpression::Root { .. } => state.write_u8(0), + BoundExpression::Variable(variable) => { + state.write_u8(2); + variable.variable().hash(state); + variable.variable_ref().hash(state); + } BoundExpression::Scalar { scalar_fn, children, @@ -146,30 +190,34 @@ impl BoundExpression { children: impl IntoIterator, ) -> VortexResult { let children = Vec::from_iter(children); - let BoundExpression::Scalar { scalar_fn, .. } = &self else { - vortex_ensure!( - children.is_empty(), - "Root expression cannot have {} children", - children.len() - ); - return Ok(self); - }; - - Self::try_new_vec(scalar_fn.clone(), children) + match &self { + BoundExpression::Scalar { scalar_fn, .. } => { + Self::try_new_vec(scalar_fn.clone(), children) + } + BoundExpression::Root { .. } | BoundExpression::Variable(_) => { + vortex_ensure!( + children.is_empty(), + "{self} cannot have {} children", + children.len() + ); + Ok(self) + } + } } /// The dtype this expression evaluates to. pub fn dtype(&self) -> &DType { match self { Self::Scalar { dtype, .. } | Self::Root { dtype } => dtype, + Self::Variable(variable) => variable.dtype(), } } - /// The bound children of this node, in argument order. Empty for [`BoundExpression::Root`]. + /// The bound children of this node, in argument order. Empty for leaf nodes. pub fn children(&self) -> &[BoundExpression] { match self { Self::Scalar { children, .. } => children.as_slice(), - Self::Root { .. } => &[], + Self::Root { .. } | Self::Variable(_) => &[], } } @@ -178,11 +226,19 @@ impl BoundExpression { &self.children()[index] } - /// The scalar function for this node, or `None` if it is the scope root. + /// The scalar function for this node, or `None` if it is a root or variable. pub fn as_scalar(&self) -> Option<&ScalarFnRef> { match self { Self::Scalar { scalar_fn, .. } => Some(scalar_fn), - Self::Root { .. } => None, + Self::Root { .. } | Self::Variable(_) => None, + } + } + + /// Return this node's bound variable, if it is a variable. + pub fn as_variable(&self) -> Option<&BoundVariable> { + match self { + Self::Variable(variable) => Some(variable), + Self::Scalar { .. } | Self::Root { .. } => None, } } @@ -213,7 +269,7 @@ impl BoundExpression { /// /// # Panics /// - /// Panics when this node is the scope root or uses a different scalar-function vtable. + /// Panics when this node is not a scalar or uses a different scalar-function vtable. pub fn as_(&self) -> &V::Options { self.as_opt::() .vortex_expect("Bound expression options type mismatch") @@ -261,6 +317,7 @@ impl Display for BoundExpression { match self { Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), Self::Root { .. } => f.write_str("$"), + Self::Variable(variable) => write!(f, "${variable}"), } } } @@ -277,19 +334,29 @@ impl Expression { /// Bind this expression against an explicit [`Scope`]. pub fn bind_scope(&self, scope: &Scope) -> VortexResult { - if self.is_root() { - return Ok(BoundExpression::new_root(scope.root().clone())); + match self { + Expression::Root => Ok(BoundExpression::new_root(scope.root().clone())), + Expression::Variable(variable) => { + let Some((dtype, variable_ref)) = scope.resolve(variable) else { + vortex_bail!("variable '{variable}' has no binder"); + }; + Ok(BoundExpression::Variable(BoundVariable { + dtype: dtype.clone(), + variable: variable.clone(), + variable_ref, + })) + } + Expression::Scalar { + scalar_fn, + children, + } => { + let children: Vec<_> = children + .iter() + .map(|child| child.bind_scope(scope)) + .try_collect()?; + BoundExpression::try_new(scalar_fn.clone(), children) + } } - - let children: Vec<_> = self - .children() - .iter() - .map(|child| child.bind_scope(scope)) - .try_collect()?; - let scalar_fn = self - .as_scalar() - .vortex_expect("root was handled above, so this is a scalar node"); - BoundExpression::try_new(scalar_fn.clone(), children) } } @@ -316,6 +383,7 @@ impl Drop for BoundExpression { #[cfg(test)] mod tests { + use vortex_error::VortexExpect; use vortex_error::VortexResult; use super::*; @@ -323,9 +391,12 @@ mod tests { use crate::dtype::PType; use crate::expr::col; use crate::expr::eq; + use crate::expr::is_not_null; use crate::expr::lit; use crate::expr::root; use crate::expr::test_harness::struct_dtype; + use crate::expr::var; + use crate::scalar_fn::fns::is_not_null::IsNotNull; use crate::scalar_fn::fns::literal::Literal; fn scope() -> Scope { @@ -404,6 +475,46 @@ mod tests { Ok(()) } + #[test] + fn variable_binds_to_its_scope() -> VortexResult<()> { + let value_dtype = DType::Primitive(PType::I64, Nullability::Nullable); + let scope = scope().with_bindings([(Variable::new("value"), value_dtype.clone())])?; + let expression = var("value"); + + assert!(expression.return_dtype(scope.root()).is_err()); + + let bound = expression.bind_scope(&scope)?; + let variable = bound + .as_variable() + .vortex_expect("variable must remain a variable after binding"); + assert_eq!(bound.dtype(), &value_dtype); + assert_eq!(variable.variable(), &Variable::new("value")); + assert_eq!(variable.variable_ref().frame(), 0); + assert_eq!(variable.variable_ref().slot(), 0); + assert_eq!(bound.to_string(), "$value"); + Ok(()) + } + + #[test] + fn unbound_variable_is_rejected() { + assert!(var("missing").bind_scope(&scope()).is_err()); + } + + #[test] + fn variable_validity_is_deferred_until_binding() -> VortexResult<()> { + let value_dtype = DType::Primitive(PType::I32, Nullability::Nullable); + let scope = scope().with_bindings([(Variable::new("value"), value_dtype)])?; + + let validity = var("value").validity()?; + assert_eq!(validity, is_not_null(var("value"))); + assert!(validity.contains::()?); + assert_eq!( + validity.bind_scope(&scope)?.dtype(), + &DType::Bool(Nullability::NonNullable) + ); + Ok(()) + } + #[test] fn clone_shares_children() -> VortexResult<()> { let bound = eq(col("a"), lit(1_i32)).bind_scope(&scope())?; diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index 250d7063834..1ab745c0169 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -69,6 +69,7 @@ impl DisplayTreeNode for Expression { match self { Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index), Expression::Root => unreachable!("the scope root has no children"), + Expression::Variable(_) => unreachable!("a variable has no children"), } } @@ -76,6 +77,7 @@ impl DisplayTreeNode for Expression { match self { Expression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), Expression::Root => write!(f, "{ROOT_DISPLAY}"), + Expression::Variable(var) => write!(f, "vortex.var({var})"), } } } @@ -89,6 +91,7 @@ impl DisplayTreeNode for BoundExpression { match self { BoundExpression::Scalar { scalar_fn, .. } => scalar_fn.signature().child_name(index), BoundExpression::Root { .. } => unreachable!("the scope root has no children"), + BoundExpression::Variable { .. } => unreachable!("a variable has no children"), } } @@ -96,6 +99,7 @@ impl DisplayTreeNode for BoundExpression { match self { BoundExpression::Scalar { scalar_fn, .. } => Display::fmt(scalar_fn, f), BoundExpression::Root { .. } => write!(f, "{ROOT_DISPLAY}"), + BoundExpression::Variable(var) => write!(f, "vortex.var({var})"), } } } @@ -152,6 +156,7 @@ mod tests { use crate::expr::root; use crate::expr::select; use crate::expr::select_exclude; + use crate::expr::var; use crate::scalar_fn::fns::between::BetweenOptions; use crate::scalar_fn::fns::between::StrictComparison; @@ -192,6 +197,12 @@ mod tests { assert_snapshot!(root_expr.display_tree().to_string(), @"vortex.root()"); } + #[test] + fn test_display_tree_variable() { + use insta::assert_snapshot; + assert_snapshot!(var("value").display_tree().to_string(), @"vortex.var(value)"); + } + #[test] fn test_display_tree_literal() { use insta::assert_snapshot; diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index 3ae9adb72d5..cf9ce929fb0 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -12,12 +12,15 @@ use std::sync::Arc; use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use crate::dtype::DType; use crate::expr::display::DisplayTreeExpr; +use crate::expr::is_not_null; use crate::expr::traversal::TraversalOrder; use crate::expr::traversal::pre_order_visit_down; +use crate::expr::variable::Variable; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::ScalarFnVTable; @@ -28,8 +31,8 @@ const NO_CHILDREN: &[Expression] = &[]; /// /// Most nodes are a scalar function applied to child expressions. [`Expression::Root`] is the scope /// itself: a language primitive rather than a registered function, because its dtype comes from the -/// scope rather than from children and it is not executable. A [`ScalarFnVTable`] can answer neither -/// of those, so `Root` is a variant instead. +/// scope rather than from children and it is not executable. [`Expression::Variable`] is likewise +/// resolved by the surrounding scope when the expression is bound. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum Expression { /// A scalar function applied to child expressions. @@ -41,6 +44,8 @@ pub enum Expression { }, /// The full scope of the expression evaluation. Root, + /// A named value resolved from the surrounding scope when the expression is bound. + Variable(Variable), } impl Expression { @@ -69,11 +74,19 @@ impl Expression { matches!(self, Self::Root) } + /// The variable referenced by this node, if this is a variable expression. + pub fn as_variable(&self) -> Option<&Variable> { + match self { + Self::Variable(variable) => Some(variable), + Self::Root | Self::Scalar { .. } => None, + } + } + /// Returns the scalar fn for this expression, or `None` if it is not a scalar node. pub fn as_scalar(&self) -> Option<&ScalarFnRef> { match self { Self::Scalar { scalar_fn, .. } => Some(scalar_fn), - Self::Root => None, + Self::Root | Self::Variable(_) => None, } } @@ -101,7 +114,7 @@ impl Expression { pub fn children(&self) -> &[Expression] { match self { Self::Scalar { children, .. } => children.as_slice(), - Self::Root => NO_CHILDREN, + Self::Root | Self::Variable(_) => NO_CHILDREN, } } @@ -117,13 +130,13 @@ impl Expression { ) -> VortexResult { let children = Vec::from_iter(children); match &self { - Self::Root => { + Self::Root | Self::Variable(_) => { vortex_ensure!( children.is_empty(), - "Expression arity mismatch: root expects 0 children but got {}", + "Expression arity mismatch: a leaf expects 0 children but got {}", children.len() ); - Ok(Self::Root) + Ok(self) } Self::Scalar { scalar_fn, .. } => { vortex_ensure!( @@ -140,17 +153,23 @@ impl Expression { } } - /// Computes the return dtype of this expression given the input dtype. - pub fn return_dtype(&self, scope: &DType) -> VortexResult { + /// Computes the return dtype of this expression given the root dtype. + /// + /// Returns an error for variables because a root dtype alone cannot resolve lexical bindings; + /// use [`Expression::bind_scope`] when the expression may contain variables. + pub fn return_dtype(&self, root_dtype: &DType) -> VortexResult { match self { - Self::Root => Ok(scope.clone()), + Self::Root => Ok(root_dtype.clone()), + Self::Variable(variable) => { + vortex_bail!("cannot determine dtype of unbound variable '{variable}'") + } Self::Scalar { scalar_fn, children, } => { let dtypes: Vec<_> = children .iter() - .map(|c| c.return_dtype(scope)) + .map(|c| c.return_dtype(root_dtype)) .try_collect()?; scalar_fn.return_dtype(&dtypes) } @@ -164,6 +183,8 @@ impl Expression { match self { // The scope is exactly as valid as itself. Self::Root => Ok(Self::Root), + // The binding supplies the variable's actual dtype later. + Self::Variable(_) => Ok(is_not_null(self.clone())), Self::Scalar { scalar_fn, .. } => scalar_fn.validity(self), } } @@ -175,6 +196,7 @@ impl Expression { pub fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { Self::Root => write!(f, "$"), + Self::Variable(variable) => write!(f, "${variable}"), Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), } } diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index fb8bfe227aa..1391ce2b427 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -17,6 +17,7 @@ use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::expr::BoundExpression; use crate::expr::Expression; +use crate::expr::Variable; use crate::scalar::Scalar; use crate::scalar::ScalarValue; use crate::scalar_fn::EmptyOptions; @@ -70,6 +71,11 @@ pub fn bound_root(dtype: DType) -> BoundExpression { BoundExpression::new_root(dtype) } +/// Creates an expression referencing the value bound to `name` in the surrounding scope. +pub fn var(name: impl AsRef) -> Expression { + Variable::new(name).into() +} + /// Return whether the expression is a root expression. pub fn is_root(expr: &Expression) -> bool { expr.is_root() diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index 59fd21c46ee..aff195d818c 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -68,6 +68,7 @@ pub mod scope; pub mod stats; pub mod transform; pub mod traversal; +pub mod variable; pub use analysis::*; pub use bound_expression::*; @@ -118,9 +119,11 @@ pub use exprs::root; pub use exprs::select; pub use exprs::select_exclude; pub use exprs::union_child_validities; +pub use exprs::var; pub use exprs::variant_get; pub use exprs::zip_expr; pub use scope::*; +pub use variable::*; pub trait VortexExprExt { /// Accumulate all field references from this expression and its children in a set @@ -163,6 +166,7 @@ impl PartialEq for ExactExpr { fn eq(&self, other: &Self) -> bool { match (&self.0, &other.0) { (Expression::Root, Expression::Root) => true, + (Expression::Variable(lhs), Expression::Variable(rhs)) => lhs == rhs, ( Expression::Scalar { scalar_fn: lhs_fn, @@ -173,7 +177,9 @@ impl PartialEq for ExactExpr { children: rhs_children, }, ) => lhs_fn == rhs_fn && Arc::ptr_eq(lhs_children, rhs_children), - _ => false, + (Expression::Root, _) + | (Expression::Scalar { .. }, _) + | (Expression::Variable(_), _) => false, } } } @@ -183,6 +189,10 @@ impl Hash for ExactExpr { fn hash(&self, state: &mut H) { match &self.0 { Expression::Root => state.write_u8(0), + Expression::Variable(variable) => { + state.write_u8(2); + variable.hash(state); + } Expression::Scalar { scalar_fn, children, diff --git a/vortex-array/src/expr/optimize.rs b/vortex-array/src/expr/optimize.rs index 3e625324a18..b47350007db 100644 --- a/vortex-array/src/expr/optimize.rs +++ b/vortex-array/src/expr/optimize.rs @@ -5,7 +5,7 @@ use std::cell::RefCell; use itertools::Itertools; use vortex_error::VortexResult; -use vortex_error::vortex_err; +use vortex_error::vortex_bail; use vortex_utils::aliases::hash_map::HashMap; use crate::dtype::DType; @@ -32,7 +32,7 @@ impl Expression { fn simplify_untyped_node(&self) -> VortexResult> { match self { Expression::Scalar { scalar_fn, .. } => scalar_fn.simplify_untyped(self), - Expression::Root => Ok(None), + Expression::Root | Expression::Variable(_) => Ok(None), } } @@ -40,7 +40,7 @@ impl Expression { fn simplify_node(&self, ctx: &dyn SimplifyCtx) -> VortexResult> { match self { Expression::Scalar { scalar_fn, .. } => scalar_fn.simplify(self, ctx), - Expression::Root => Ok(None), + Expression::Root | Expression::Variable(_) => Ok(None), } } @@ -51,7 +51,7 @@ impl Expression { ) -> VortexResult>> { match self { Expression::Scalar { scalar_fn, .. } => scalar_fn.reduce_expression(node), - Expression::Root => Ok(None), + Expression::Root | Expression::Variable(_) => Ok(None), } } @@ -186,16 +186,22 @@ impl SimplifyCtx for SimplifyCache<'_> { return Ok(dtype.clone()); } - // Otherwise, compute dtype from children - let input_dtypes: Vec<_> = expr - .children() - .iter() - .map(|c| self.return_dtype(c)) - .try_collect()?; - let dtype = expr - .as_scalar() - .ok_or_else(|| vortex_err!("cannot type a non-scalar expression: {expr}"))? - .return_dtype(&input_dtypes)?; + let dtype = match expr { + Expression::Variable(variable) => { + vortex_bail!("cannot determine dtype of unbound variable '{variable}'") + } + Expression::Scalar { + scalar_fn, + children, + } => { + let input_dtypes: Vec<_> = children + .iter() + .map(|child| self.return_dtype(child)) + .try_collect()?; + scalar_fn.return_dtype(&input_dtypes)? + } + Expression::Root => unreachable!("handled above"), + }; self.dtype_cache .borrow_mut() .insert(expr.clone(), dtype.clone()); diff --git a/vortex-array/src/expr/proto.rs b/vortex-array/src/expr/proto.rs index 4f544fecafb..b6ec9bc771b 100644 --- a/vortex-array/src/expr/proto.rs +++ b/vortex-array/src/expr/proto.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools; +use prost::Message; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -9,8 +10,10 @@ use vortex_proto::expr as pb; use vortex_session::VortexSession; use crate::expr::Expression; +use crate::expr::Variable; use crate::scalar_fn::ForeignScalarFnVTable; use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::session::ScalarFnSessionExt; pub trait ExprSerializeProtoExt { @@ -22,35 +25,74 @@ pub trait ExprSerializeProtoExt { /// already-serialized expressions keep round-tripping. pub(crate) const ROOT_ID: &str = "vortex.root"; +/// The wire id for [`Expression::Variable`]. +pub(crate) const VARIABLE_ID: &str = "vortex.var"; + +impl Variable { + /// Serialize this variable to its protobuf representation. + fn serialize_proto(&self) -> pb::Expr { + pb::Expr { + id: VARIABLE_ID.to_string(), + children: vec![], + metadata: Some( + pb::VariableOpts { + name: self.name().to_string(), + } + .encode_to_vec(), + ), + } + } + + /// Deserialize a variable expression whose id is [`VARIABLE_ID`]. + fn from_proto(expr: &pb::Expr) -> VortexResult { + vortex_ensure!( + expr.children.is_empty(), + "a variable must have no children, got {}", + expr.children.len() + ); + let options = pb::VariableOpts::decode(expr.metadata())?; + Ok(Self::new(options.name)) + } +} + +fn serialize_scalar( + expression: &Expression, + scalar_fn: &ScalarFnRef, + children: &[Expression], +) -> VortexResult { + let children = children + .iter() + .map(|child| child.serialize_proto()) + .try_collect()?; + + let metadata = scalar_fn.options().serialize()?.ok_or_else(|| { + vortex_err!( + "Expression '{}' is not serializable: {expression}", + scalar_fn.id() + ) + })?; + + Ok(pb::Expr { + id: scalar_fn.id().to_string(), + children, + metadata: Some(metadata), + }) +} + impl ExprSerializeProtoExt for Expression { fn serialize_proto(&self) -> VortexResult { - let Some(scalar_fn) = self.as_scalar() else { - return Ok(pb::Expr { + match self { + Expression::Root => Ok(pb::Expr { id: ROOT_ID.to_string(), children: vec![], metadata: Some(vec![]), - }); - }; - - let children = self - .children() - .iter() - .map(|child| child.serialize_proto()) - .try_collect()?; - - let metadata = scalar_fn.options().serialize()?.ok_or_else(|| { - vortex_err!( - "Expression '{}' is not serializable: {}", - scalar_fn.id(), - self - ) - })?; - - Ok(pb::Expr { - id: scalar_fn.id().to_string(), - children, - metadata: Some(metadata), - }) + }), + Expression::Variable(variable) => Ok(variable.serialize_proto()), + Expression::Scalar { + scalar_fn, + children, + } => serialize_scalar(self, scalar_fn, children), + } } } @@ -66,6 +108,10 @@ impl Expression { return Ok(Expression::Root); } + if expr.id == VARIABLE_ID { + return Ok(Variable::from_proto(expr)?.into()); + } + #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] let expr_id = ScalarFnId::new(expr.id.as_str()); let children = expr @@ -98,6 +144,7 @@ pub fn deserialize_expr_proto( #[cfg(test)] mod tests { use prost::Message; + use vortex_error::VortexResult; use vortex_proto::expr as pb; use vortex_session::VortexSession; @@ -111,6 +158,7 @@ mod tests { use crate::expr::lit; use crate::expr::or; use crate::expr::root; + use crate::expr::var; use crate::scalar_fn::fns::between::BetweenOptions; use crate::scalar_fn::fns::between::StrictComparison; use crate::scalar_fn::session::ScalarFnSession; @@ -141,6 +189,28 @@ mod tests { assert_eq!(&deser_expr, &expr); } + #[test] + fn variable_serde() -> VortexResult<()> { + let expression = var("value"); + let encoded = expression.serialize_proto()?.encode_to_vec(); + let proto = pb::Expr::decode(encoded.as_slice())?; + + assert_eq!( + Expression::from_proto(&proto, &array_session())?, + expression + ); + Ok(()) + } + + #[test] + fn variable_rejects_children() -> VortexResult<()> { + let mut proto = var("value").serialize_proto()?; + proto.children.push(root().serialize_proto()?); + + assert!(Expression::from_proto(&proto, &array_session()).is_err()); + Ok(()) + } + #[test] fn unknown_expression_id_allow_unknown() { let session = VortexSession::empty().with::(); diff --git a/vortex-array/src/expr/scope.rs b/vortex-array/src/expr/scope.rs index 04de3202cf5..3c71401d69a 100644 --- a/vortex-array/src/expr/scope.rs +++ b/vortex-array/src/expr/scope.rs @@ -1,28 +1,133 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_utils::aliases::hash_set::HashSet; + use crate::dtype::DType; +use crate::expr::Variable; + +/// A stable location for a variable resolved in a lexical [`Scope`]. +/// +/// Frames are indexed from the outermost binding and slots are indexed in declaration order. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct VariableRef { + frame: usize, + slot: usize, +} + +impl VariableRef { + fn new(frame: usize, slot: usize) -> Self { + Self { frame, slot } + } + + /// The lexical frame containing this binding, counted from the outermost frame. + pub fn frame(&self) -> usize { + self.frame + } + + /// The binding's position in its frame, in declaration order. + pub fn slot(&self) -> usize { + self.slot + } +} + +/// A set of named bindings introduced by a single binder. +/// +/// Names within a frame must be unique. A name in an inner frame shadows the same name in an +/// outer frame. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct Frame(Arc>); + +impl Frame { + /// Create a frame from its bindings, in declaration order. + pub fn try_new(bindings: impl IntoIterator) -> VortexResult { + let mut seen = HashSet::new(); + let mut frame = Vec::new(); + for (variable, dtype) in bindings { + if !seen.insert(variable.clone()) { + vortex_bail!("duplicate binding '{variable}' in a single frame"); + } + frame.push((variable, dtype)); + } + Ok(Self(Arc::new(frame))) + } + + /// The bindings in this frame, in declaration order. + pub fn bindings(&self) -> &[(Variable, DType)] { + &self.0 + } + + fn get(&self, name: &Variable) -> Option<(&DType, usize)> { + self.0 + .iter() + .enumerate() + .find_map(|(slot, (bound, dtype))| (bound == name).then_some((dtype, slot))) + } +} /// The context an [`Expression`](crate::expr::Expression) is bound against. /// -/// Today a scope is just the dtype that [`root`](crate::expr::root) resolves to. It is an opaque -/// struct rather than a bare [`DType`] so that lexical bindings can be added later without changing -/// [`Expression::bind_scope`](crate::expr::Expression::bind_scope)'s signature. +/// A scope is the dtype that [`root`](crate::expr::root) resolves to, plus a stack of [`Frame`]s +/// holding named bindings. Names resolve from the innermost frame outward. #[derive(Clone, Debug, PartialEq, Eq)] pub struct Scope { root: DType, + frames: Vec, } impl Scope { /// Create a scope in which `root` resolves to the given dtype. pub fn new(root: DType) -> Self { - Self { root } + Self { + root, + frames: Vec::new(), + } } /// The dtype that `root` resolves to. pub fn root(&self) -> &DType { &self.root } + + /// The number of lexical frames in this scope. + pub fn depth(&self) -> usize { + self.frames.len() + } + + /// Return this scope extended with `frame` as its innermost frame. + pub fn push_frame(&self, frame: Frame) -> Self { + let mut frames = self.frames.clone(); + frames.push(frame); + Self { + root: self.root.clone(), + frames, + } + } + + /// Return this scope extended with one frame containing `bindings`. + pub fn with_bindings( + self, + bindings: impl IntoIterator, + ) -> VortexResult { + Ok(self.push_frame(Frame::try_new(bindings)?)) + } + + /// Resolve `name`, searching innermost-first so inner bindings shadow outer ones. + pub fn resolve(&self, name: &Variable) -> Option<(&DType, VariableRef)> { + self.frames + .iter() + .enumerate() + .rev() + .find_map(|(frame, bindings)| { + bindings + .get(name) + .map(|(dtype, slot)| (dtype, VariableRef::new(frame, slot))) + }) + } } impl From for Scope { @@ -33,8 +138,19 @@ impl From for Scope { #[cfg(test)] mod tests { + use vortex_error::VortexResult; + use super::*; use crate::dtype::Nullability; + use crate::dtype::PType; + + fn i32_() -> DType { + DType::Primitive(PType::I32, Nullability::NonNullable) + } + + fn utf8() -> DType { + DType::Utf8(Nullability::NonNullable) + } #[test] fn root_round_trips() { @@ -42,4 +158,63 @@ mod tests { assert_eq!(Scope::new(dtype.clone()).root(), &dtype); assert_eq!(Scope::from(dtype.clone()).root(), &dtype); } + + #[test] + fn resolve_tracks_frames_slots_and_shadowing() -> VortexResult<()> { + let scope = Scope::new(i32_()) + .push_frame(Frame::try_new([(Variable::new("x"), i32_())])?) + .push_frame(Frame::try_new([ + (Variable::new("x"), utf8()), + (Variable::new("y"), i32_()), + ])?); + + assert_eq!(scope.depth(), 2); + assert_eq!( + scope.resolve(&Variable::new("x")), + Some((&utf8(), VariableRef::new(1, 0))) + ); + assert_eq!( + scope.resolve(&Variable::new("y")), + Some((&i32_(), VariableRef::new(1, 1))) + ); + assert!(scope.resolve(&Variable::new("missing")).is_none()); + Ok(()) + } + + #[test] + fn duplicate_names_in_one_frame_are_rejected() { + assert!( + Frame::try_new([(Variable::new("x"), i32_()), (Variable::new("x"), utf8())]).is_err() + ); + } + + #[test] + fn pushing_a_frame_does_not_mutate_the_original() -> VortexResult<()> { + let outer = Scope::new(i32_()); + let inner = outer.push_frame(Frame::try_new([(Variable::new("x"), i32_())])?); + + assert!(outer.resolve(&Variable::new("x")).is_none()); + assert!(inner.resolve(&Variable::new("x")).is_some()); + Ok(()) + } + + #[test] + fn pushing_an_inner_frame_preserves_outer_variable_refs() -> VortexResult<()> { + let captured = Variable::new("captured"); + let outer = Scope::new(i32_()).push_frame(Frame::try_new([(captured.clone(), utf8())])?); + let outer_ref = outer + .resolve(&captured) + .map(|(_, variable_ref)| variable_ref); + + let inner = outer.push_frame(Frame::try_new([(Variable::new("parameter"), i32_())])?); + + assert_eq!(outer_ref, Some(VariableRef::new(0, 0))); + assert_eq!( + inner + .resolve(&captured) + .map(|(_, variable_ref)| variable_ref), + outer_ref + ); + Ok(()) + } } diff --git a/vortex-array/src/expr/traversal/mod.rs b/vortex-array/src/expr/traversal/mod.rs index 0f55600afe5..987d420c9a0 100644 --- a/vortex-array/src/expr/traversal/mod.rs +++ b/vortex-array/src/expr/traversal/mod.rs @@ -533,8 +533,11 @@ impl Node for BoundExpression { &'a self, mut f: F, ) -> VortexResult { - let BoundExpression::Scalar { children, .. } = self else { - return Ok(TraversalOrder::Continue); + let children = match self { + BoundExpression::Scalar { children, .. } => children, + BoundExpression::Root { .. } | BoundExpression::Variable(_) => { + return Ok(TraversalOrder::Continue); + } }; for child in children.iter() { @@ -551,8 +554,11 @@ impl Node for BoundExpression { self, mut f: F, ) -> VortexResult> { - let BoundExpression::Scalar { children, .. } = &self else { - return Ok(Transformed::no(self)); + let children = match &self { + BoundExpression::Scalar { children, .. } => children, + BoundExpression::Root { .. } | BoundExpression::Variable(_) => { + return Ok(Transformed::no(self)); + } }; let mut order = TraversalOrder::Continue; @@ -593,14 +599,16 @@ impl Node for BoundExpression { fn iter_children(&self, f: impl FnOnce(&mut dyn Iterator) -> T) -> T { match self { BoundExpression::Scalar { children, .. } => f(&mut children.iter()), - BoundExpression::Root { .. } => f(&mut std::iter::empty()), + BoundExpression::Root { .. } | BoundExpression::Variable(_) => { + f(&mut std::iter::empty()) + } } } fn children_count(&self) -> usize { match self { BoundExpression::Scalar { children, .. } => children.len(), - BoundExpression::Root { .. } => 0, + BoundExpression::Root { .. } | BoundExpression::Variable(_) => 0, } } } diff --git a/vortex-array/src/expr/variable.rs b/vortex-array/src/expr/variable.rs new file mode 100644 index 00000000000..96a751a30d5 --- /dev/null +++ b/vortex-array/src/expr/variable.rs @@ -0,0 +1,101 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; +use std::sync::Arc; + +use crate::dtype::FieldName; +use crate::expr::Expression; + +/// The name of a value bound in a [`Scope`](crate::expr::Scope). +#[derive(Clone, Debug)] +pub struct Variable(Arc); + +impl Variable { + /// Create a variable with the given name. + pub fn new(name: impl AsRef) -> Self { + Self(Arc::from(name.as_ref())) + } + + /// The variable's name. + pub fn name(&self) -> &str { + &self.0 + } +} + +/// Compares by name, with a pointer-equality fast path for cloned variables. +impl PartialEq for Variable { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) || self.0 == other.0 + } +} + +impl Eq for Variable {} + +impl Hash for Variable { + fn hash(&self, state: &mut H) { + self.0.hash(state); + } +} + +impl Display for Variable { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + Display::fmt(&self.0, f) + } +} + +impl From for Expression { + fn from(variable: Variable) -> Self { + Expression::Variable(variable) + } +} + +impl From for Variable { + fn from(field_name: FieldName) -> Self { + Self(Arc::clone(field_name.inner())) + } +} + +impl From<&str> for Variable { + fn from(value: &str) -> Self { + Self(value.into()) + } +} + +impl AsRef for Variable { + fn as_ref(&self) -> &str { + &self.0 + } +} + +#[cfg(test)] +mod tests { + use std::collections::hash_map::RandomState; + use std::hash::BuildHasher; + + use super::*; + + #[test] + fn equality_and_hashing_are_by_name() { + let state = RandomState::new(); + let x1 = Variable::new("x"); + let x2 = Variable::new("x"); + let y = Variable::new("y"); + + assert_eq!(x1, x2); + assert_ne!(x1, y); + assert_eq!(state.hash_one(x1), state.hash_one(x2)); + } + + #[test] + fn cloning_shares_the_name() { + let original = Variable::new("x"); + let cloned = original.clone(); + + assert!(Arc::ptr_eq(&original.0, &cloned.0)); + } +} diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index d0590f2bf58..83c2e26f804 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -2,8 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools; -use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use crate::ArrayRef; use crate::IntoArray; @@ -12,62 +12,96 @@ use crate::arrays::ScalarFnArray; use crate::expr::BoundExpression; use crate::expr::Expression; use crate::optimizer::ArrayOptimizer; +use crate::scalar_fn::ScalarFnRef; 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 BoundExpression::Scalar { - scalar_fn, - children, - .. - } = expr - else { - return Ok(self); - }; + match expr { + BoundExpression::Root { .. } => Ok(self), + BoundExpression::Variable(variable) => { + vortex_bail!("cannot apply variable '{variable}' without a provided value") + } + BoundExpression::Scalar { + scalar_fn, + children, + .. + } => apply_bound_scalar_fn(self, scalar_fn, children), + } + } - if let Some(scalar) = scalar_fn.as_opt::() { - return Ok(ConstantArray::new(scalar.clone(), self.len()).into_array()); + /// Apply the expression to this array, producing a new array in constant time. + pub fn apply(self, expr: &Expression) -> VortexResult { + match expr { + Expression::Root => Ok(self), + Expression::Variable(variable) => { + vortex_bail!("cannot apply unbound variable '{variable}'") + } + Expression::Scalar { + scalar_fn, + children, + } => apply_scalar_fn(self, scalar_fn, children), } + } +} - let children: Vec<_> = children - .iter() - .map(|child| self.clone().apply_bound(child)) - .try_collect()?; +fn apply_bound_scalar_fn( + root: ArrayRef, + scalar_fn: &ScalarFnRef, + children: &[BoundExpression], +) -> VortexResult { + if let Some(scalar) = scalar_fn.as_opt::() { + return Ok(ConstantArray::new(scalar.clone(), root.len()).into_array()); + } - let array = - ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, self.len())?.into_array(); + let children: Vec<_> = children + .iter() + .map(|child| root.clone().apply_bound(child)) + .try_collect()?; + let array = + ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, root.len())?.into_array(); + array.optimize() +} - array.optimize() +fn apply_scalar_fn( + root: ArrayRef, + scalar_fn: &ScalarFnRef, + children: &[Expression], +) -> VortexResult { + if let Some(scalar) = scalar_fn.as_opt::() { + return Ok(ConstantArray::new(scalar.clone(), root.len()).into_array()); } - /// Apply the expression to this array, producing a new array in constant time. - pub fn apply(self, expr: &Expression) -> VortexResult { - // If the expression is a root, return self. - if expr.is_root() { - return Ok(self); - } + let children: Vec<_> = children + .iter() + .map(|child| root.clone().apply(child)) + .try_collect()?; + let array = + ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, root.len())?.into_array(); + array.optimize() +} - // Manually convert literals to ConstantArray. - if let Some(scalar) = expr.as_opt::() { - return Ok(ConstantArray::new(scalar.clone(), self.len()).into_array()); - } +#[cfg(test)] +mod tests { + use vortex_buffer::buffer; + use vortex_error::VortexResult; - // Otherwise, collect the child arrays. - let children: Vec<_> = expr - .children() - .iter() - .map(|e| self.clone().apply(e)) - .try_collect()?; + use crate::IntoArray; + use crate::expr::Scope; + use crate::expr::Variable; + use crate::expr::var; - // And wrap the scalar function up in an array. - let scalar_fn = expr - .as_scalar() - .vortex_expect("root and literal were handled above, so this is a scalar node"); - let array = - ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, self.len())?.into_array(); + #[test] + fn variable_application_requires_a_runtime_binding() -> VortexResult<()> { + let root = buffer![1_i32, 2, 3].into_array(); + let expression = var("value"); + assert!(root.clone().apply(&expression).is_err()); - // Optimize the resulting array's root. - array.optimize() + let scope = Scope::new(root.dtype().clone()) + .with_bindings([(Variable::new("value"), root.dtype().clone())])?; + let bound = expression.bind_scope(&scope)?; + assert!(root.apply_bound(&bound).is_err()); + Ok(()) } } diff --git a/vortex-proto/proto/expr.proto b/vortex-proto/proto/expr.proto index 00ffaac433c..d506214c4a2 100644 --- a/vortex-proto/proto/expr.proto +++ b/vortex-proto/proto/expr.proto @@ -125,3 +125,8 @@ message SelectOpts { message CaseWhenOpts { uint32 num_children = 1; } + +// Options for `vortex.var`, a reference to a name bound in an enclosing scope. +message VariableOpts { + string name = 1; +} diff --git a/vortex-proto/src/generated/vortex.expr.rs b/vortex-proto/src/generated/vortex.expr.rs index a44328623e3..2660ce852b7 100644 --- a/vortex-proto/src/generated/vortex.expr.rs +++ b/vortex-proto/src/generated/vortex.expr.rs @@ -207,3 +207,9 @@ pub struct CaseWhenOpts { #[prost(uint32, tag = "1")] pub num_children: u32, } +/// Options for `vortex.var`, a reference to a name bound in an enclosing scope. +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct VariableOpts { + #[prost(string, tag = "1")] + pub name: ::prost::alloc::string::String, +}