From 315d461298418071852718cd588b1b363d1ce934 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 19 Aug 2026 14:05:28 +0000 Subject: [PATCH] refactor(expr): share bound expressions by Arc handle `BoundExpression` was passed around by value, with its children behind a single `Arc>` so that clones stayed cheap. That made node identity a property of the child vector rather than of the node, and it meant every node clone still copied a `DType` and a `ScalarFnRef`. Introduce `BoundExpressionRef = Arc` as the currency of the bound tree, and hold children as `Box<[BoundExpressionRef]>`. Cloning a subtree is now a refcount bump on the node itself, and rebuilding a node keeps the untouched children in place. Threaded through every crate that builds, walks, or caches bound expressions: constructors, `bound::*` builders, traversal, analysis, partitioning, stats rewrites, layouts, scans, pruning, and CUDA. Notable details: - `Node` is implemented for `BoundExpressionRef` rather than for the enum, so traversals hand out shared handles and identity survives a walk. - `with_children`, `falsify`, and `satisfy` take `self: BoundExpressionRef`. - `ExactBoundExpr` gains an `Arc::ptr_eq` fast path and compares child handles elementwise, preserving its previous identity semantics. - The iterative `Drop` walks shared handles, descending only into subtrees it is the last owner of. Covered by a new deep-tree regression test that overflows the stack without it. Signed-off-by: Joe Isaacs --- vortex-array/src/expr/analysis/annotation.rs | 27 +-- .../src/expr/analysis/immediate_access.rs | 6 +- vortex-array/src/expr/analysis/labeling.rs | 17 +- .../expr/analysis/referenced_field_paths.rs | 15 +- vortex-array/src/expr/bound_expression.rs | 186 +++++++++++------- vortex-array/src/expr/display.rs | 32 ++- vortex-array/src/expr/exprs.rs | 163 ++++++++------- vortex-array/src/expr/mod.rs | 22 ++- .../src/expr/transform/bound_partition.rs | 44 +++-- vortex-array/src/expr/traversal/mod.rs | 57 +++--- vortex-array/src/expression.rs | 5 +- vortex-array/src/scalar_fn/fns/dynamic.rs | 6 +- vortex-array/src/scalar_fn/vtable.rs | 5 +- vortex-array/src/stats/bind.rs | 22 +-- vortex-array/src/stats/expr.rs | 43 ++-- vortex-array/src/stats/rewrite.rs | 46 ++--- vortex-array/src/stats/rewrite/builtins.rs | 155 ++++++++------- vortex-cuda/src/layout.rs | 12 +- vortex-file/src/pruning.rs | 16 +- vortex-file/src/tests.rs | 4 +- vortex-file/src/v2/file_stats_reader.rs | 12 +- vortex-layout/benches/zone_map_prune.rs | 16 +- vortex-layout/src/layouts/chunked/reader.rs | 8 +- vortex-layout/src/layouts/dict/reader.rs | 39 ++-- vortex-layout/src/layouts/flat/reader.rs | 12 +- vortex-layout/src/layouts/list/expr.rs | 23 ++- vortex-layout/src/layouts/list/reader.rs | 26 +-- vortex-layout/src/layouts/partitioned.rs | 14 +- vortex-layout/src/layouts/row_idx/mod.rs | 39 ++-- vortex-layout/src/layouts/struct_/reader.rs | 60 +++--- vortex-layout/src/layouts/zoned/pruning.rs | 21 +- vortex-layout/src/layouts/zoned/reader.rs | 12 +- vortex-layout/src/layouts/zoned/zone_map.rs | 18 +- vortex-layout/src/plan/plans/concat.rs | 2 +- vortex-layout/src/plan/plans/eval.rs | 12 +- vortex-layout/src/plan/plans/pack.rs | 58 +++--- vortex-layout/src/plan/plans/row_idx.rs | 16 +- vortex-layout/src/plan/plans/take.rs | 5 +- vortex-layout/src/reader.rs | 8 +- vortex-layout/src/scan/filter.rs | 13 +- vortex-layout/src/scan/layout.rs | 12 +- vortex-layout/src/scan/multi.rs | 6 +- vortex-layout/src/scan/repeated_scan.rs | 12 +- vortex-layout/src/scan/scan_builder.rs | 39 ++-- vortex-layout/src/scan/split_by.rs | 8 +- vortex-layout/src/scan/tasks.rs | 4 +- vortex-spatial/src/prune/distance.rs | 14 +- vortex-spatial/src/prune/intersects.rs | 10 +- vortex-spatial/src/prune/mod.rs | 36 ++-- 49 files changed, 779 insertions(+), 659 deletions(-) diff --git a/vortex-array/src/expr/analysis/annotation.rs b/vortex-array/src/expr/analysis/annotation.rs index f2c63d7c4d7..cd888700056 100644 --- a/vortex-array/src/expr/analysis/annotation.rs +++ b/vortex-array/src/expr/analysis/annotation.rs @@ -2,13 +2,14 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::hash::Hash; +use std::sync::Arc; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::aliases::hash_set::HashSet; -use crate::expr::BoundExpression; +use crate::expr::BoundExpressionRef; use crate::expr::ExactBoundExpr; use crate::expr::Expression; use crate::expr::traversal::Node; @@ -85,11 +86,11 @@ where /// Unlike [`descendent_annotations`], this uses [`ExactBoundExpr`] keys to preserve the cheap /// identity semantics of an already-bound tree. pub fn descendent_bound_annotations( - expr: &BoundExpression, + expr: &BoundExpressionRef, annotate: A, ) -> BoundAnnotations where - A: AnnotationFn, + A: AnnotationFn, { bound_annotations(expr, annotate, true) } @@ -98,22 +99,22 @@ where /// /// The returned map uses [`ExactBoundExpr`] keys so lookups do not structurally hash node dtypes. pub fn direct_bound_annotations( - expr: &BoundExpression, + expr: &BoundExpressionRef, annotate: A, ) -> BoundAnnotations where - A: AnnotationFn, + A: AnnotationFn, { bound_annotations(expr, annotate, false) } fn bound_annotations( - expr: &BoundExpression, + expr: &BoundExpressionRef, annotate: A, propagate_up: bool, ) -> BoundAnnotations where - A: AnnotationFn, + A: AnnotationFn, { let mut visitor = BoundAnnotationVisitor { annotations: Default::default(), @@ -176,7 +177,7 @@ where struct BoundAnnotationVisitor where - A: AnnotationFn, + A: AnnotationFn, { annotations: BoundAnnotations, annotate: A, @@ -185,9 +186,9 @@ where impl<'a, A> NodeVisitor<'a> for BoundAnnotationVisitor where - A: AnnotationFn, + A: AnnotationFn, { - type NodeTy = BoundExpression; + type NodeTy = BoundExpressionRef; fn visit_down(&mut self, node: &'a Self::NodeTy) -> VortexResult { let annotations = (self.annotate)(node); @@ -196,7 +197,7 @@ where } self.annotations - .entry(ExactBoundExpr(node.clone())) + .entry(ExactBoundExpr(Arc::clone(node))) .or_default() .extend(annotations); Ok(TraversalOrder::Skip) @@ -212,13 +213,13 @@ where .iter() .filter_map(|child| { self.annotations - .get(&ExactBoundExpr(child.clone())) + .get(&ExactBoundExpr(Arc::clone(child))) .cloned() }) .collect::>(); let annotations = self .annotations - .entry(ExactBoundExpr(node.clone())) + .entry(ExactBoundExpr(Arc::clone(node))) .or_default(); child_annotations .into_iter() diff --git a/vortex-array/src/expr/analysis/immediate_access.rs b/vortex-array/src/expr/analysis/immediate_access.rs index 6c2e4975a92..ec992174d72 100644 --- a/vortex-array/src/expr/analysis/immediate_access.rs +++ b/vortex-array/src/expr/analysis/immediate_access.rs @@ -5,7 +5,7 @@ use vortex_error::VortexExpect; use crate::dtype::FieldName; use crate::dtype::StructFields; -use crate::expr::BoundExpression; +use crate::expr::BoundExpressionRef; use crate::expr::Expression; use crate::expr::analysis::AnnotationFn; use crate::scalar_fn::fns::get_item::GetItem; @@ -64,8 +64,8 @@ pub fn make_free_field_annotator( /// Returns the free top-level fields for bound expression nodes. pub fn make_bound_free_field_annotator( scope: &StructFields, -) -> impl AnnotationFn { - move |expr: &BoundExpression| { +) -> impl AnnotationFn { + move |expr: &BoundExpressionRef| { let Some(scalar_fn) = expr.as_scalar() else { return scope.names().iter().cloned().collect(); }; diff --git a/vortex-array/src/expr/analysis/labeling.rs b/vortex-array/src/expr/analysis/labeling.rs index 51957f26190..5f247a90b9f 100644 --- a/vortex-array/src/expr/analysis/labeling.rs +++ b/vortex-array/src/expr/analysis/labeling.rs @@ -2,12 +2,13 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::hash::Hash; +use std::sync::Arc; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_utils::aliases::hash_map::HashMap; -use crate::expr::BoundExpression; +use crate::expr::BoundExpressionRef; use crate::expr::ExactBoundExpr; use crate::expr::Expression; use crate::expr::traversal::Node; @@ -62,8 +63,8 @@ where /// /// This avoids structurally hashing bound dtypes, which may deserialize a lazy schema. pub fn label_bound_tree( - expr: &BoundExpression, - self_label: impl Fn(&BoundExpression) -> L, + expr: &BoundExpressionRef, + self_label: impl Fn(&BoundExpressionRef) -> L, mut merge_child: impl FnMut(L, &L) -> L, ) -> BoundLabels { let mut visitor = BoundLabelingVisitor { @@ -120,7 +121,7 @@ where struct BoundLabelingVisitor<'a, L, F, G> where - F: Fn(&BoundExpression) -> L, + F: Fn(&BoundExpressionRef) -> L, G: FnMut(L, &L) -> L, { labels: BoundLabels, @@ -130,10 +131,10 @@ where impl<'node, 'visitor, L: Clone, F, G> NodeVisitor<'node> for BoundLabelingVisitor<'visitor, L, F, G> where - F: Fn(&BoundExpression) -> L, + F: Fn(&BoundExpressionRef) -> L, G: FnMut(L, &L) -> L, { - type NodeTy = BoundExpression; + type NodeTy = BoundExpressionRef; fn visit_down(&mut self, _node: &'node Self::NodeTy) -> VortexResult { Ok(TraversalOrder::Continue) @@ -144,12 +145,12 @@ where let final_label = node.children().iter().fold(self_label, |acc, child| { let child_label = self .labels - .get(&ExactBoundExpr(child.clone())) + .get(&ExactBoundExpr(Arc::clone(child))) .vortex_expect("child must have label"); (self.merge_child)(acc, child_label) }); self.labels - .insert(ExactBoundExpr(node.clone()), final_label); + .insert(ExactBoundExpr(Arc::clone(node)), final_label); Ok(TraversalOrder::Continue) } } diff --git a/vortex-array/src/expr/analysis/referenced_field_paths.rs b/vortex-array/src/expr/analysis/referenced_field_paths.rs index 48ef285c20a..6f2b467a096 100644 --- a/vortex-array/src/expr/analysis/referenced_field_paths.rs +++ b/vortex-array/src/expr/analysis/referenced_field_paths.rs @@ -1,13 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + use vortex_error::VortexResult; use vortex_error::vortex_err; use crate::dtype::Field; use crate::dtype::FieldPath; use crate::dtype::FieldPathSet; -use crate::expr::BoundExpression; +use crate::expr::BoundExpressionRef; use crate::expr::traversal::FoldDownContext; use crate::expr::traversal::FoldUp; use crate::expr::traversal::NodeExt; @@ -22,12 +24,11 @@ use crate::scalar_fn::fns::select::Select; /// expression is represented by [`FieldPath::root`], which conservatively selects all fields. /// Scalar functions other than `GetItem` and `Select` conservatively reference each complete child /// output. -pub fn referenced_field_paths(expr: &BoundExpression) -> VortexResult { +pub fn referenced_field_paths(expr: &BoundExpressionRef) -> VortexResult { let mut collector = ReferencedFieldPaths { field_paths: FieldPathSet::default(), }; - expr.clone() - .fold_context(&vec![FieldPath::root()], &mut collector)?; + Arc::clone(expr).fold_context(&vec![FieldPath::root()], &mut collector)?; Ok(collector.field_paths) } @@ -47,14 +48,14 @@ struct ReferencedFieldPaths { } impl NodeFolderContext for ReferencedFieldPaths { - type NodeTy = BoundExpression; + type NodeTy = BoundExpressionRef; type Result = (); type Context = Vec; fn visit_down( &mut self, requested: &Self::Context, - node: &BoundExpression, + node: &BoundExpressionRef, ) -> VortexResult> { if node.is_root() { self.field_paths.extend( @@ -115,7 +116,7 @@ impl NodeFolderContext for ReferencedFieldPaths { fn visit_up( &mut self, - _node: BoundExpression, + _node: BoundExpressionRef, _requested: &Self::Context, _children: Vec<()>, ) -> VortexResult> { diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 627774f442f..b3a5599a534 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -18,17 +18,24 @@ use crate::dtype::DType; use crate::expr::Expression; use crate::expr::display::DisplayTreeExpr; use crate::expr::scope::Scope; -use crate::expr::traversal::TraversalOrder; -use crate::expr::traversal::pre_order_visit_down; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::ScalarFnVTable; use crate::stats::rewrite::StatsRewriteCtx; +/// A shared handle to a [`BoundExpression`]. +/// +/// Bound trees are immutable and shared by handle: every node is reference counted, so cloning a +/// subtree, storing it in a cache, or handing it to another thread is a refcount bump rather than +/// a copy of the tree. +pub type BoundExpressionRef = Arc; + /// An [`Expression`] that has been type-checked against a [`Scope`]. /// /// Every node carries its own dtype, so reading one is a field access rather than a walk of the /// subtree. Holding a `BoundExpression` is proof that the whole tree type-checked. /// +/// Nodes are handed around as [`BoundExpressionRef`] rather than by value. +/// /// 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)] @@ -41,9 +48,8 @@ pub enum BoundExpression { scalar_fn: ScalarFnRef, /// The bound children, in argument order. /// - /// Sharing keeps clones cheap even though the iterative [`Drop`] implementation prevents - /// consumers from destructuring a `BoundExpression` by value. - children: Arc>, + /// Each child is shared, so rebuilding a node keeps the untouched subtrees in place. + children: Box<[BoundExpressionRef]>, }, /// The scope itself. Its dtype is the scope's root dtype. Root { @@ -53,12 +59,20 @@ pub enum BoundExpression { } /// A bound-expression wrapper that compares shared tree identity instead of structure. +/// +/// Two wrappers are equal when they hold the same node, or when they hold nodes built from the +/// same scalar function over the very same child handles. Structurally equal trees built +/// independently are not equal, which is what keeps identity-keyed caches from walking a tree on +/// every lookup. #[derive(Clone, Debug)] -pub struct ExactBoundExpr(pub BoundExpression); +pub struct ExactBoundExpr(pub BoundExpressionRef); impl PartialEq for ExactBoundExpr { fn eq(&self, other: &Self) -> bool { - match (&self.0, &other.0) { + if Arc::ptr_eq(&self.0, &other.0) { + return true; + } + match (&*self.0, &*other.0) { ( BoundExpression::Root { dtype: lhs_dtype }, BoundExpression::Root { dtype: rhs_dtype }, @@ -76,7 +90,11 @@ impl PartialEq for ExactBoundExpr { }, ) => { lhs_fn == rhs_fn - && Arc::ptr_eq(lhs_children, rhs_children) + && lhs_children.len() == rhs_children.len() + && lhs_children + .iter() + .zip(rhs_children.iter()) + .all(|(lhs, rhs)| Arc::ptr_eq(lhs, rhs)) && lhs_dtype == rhs_dtype } _ => false, @@ -90,7 +108,7 @@ 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 { + match &*self.0 { BoundExpression::Root { .. } => state.write_u8(0), BoundExpression::Scalar { scalar_fn, @@ -99,7 +117,9 @@ impl Hash for ExactBoundExpr { } => { state.write_u8(1); scalar_fn.hash(state); - Arc::as_ptr(children).hash(state); + for child in children.iter() { + Arc::as_ptr(child).hash(state); + } } } } @@ -107,16 +127,16 @@ impl Hash for ExactBoundExpr { impl BoundExpression { /// Create a bound root expression with the given dtype. - pub fn new_root(dtype: DType) -> Self { - Self::Root { dtype } + pub fn new_root(dtype: DType) -> BoundExpressionRef { + Arc::new(Self::Root { dtype }) } /// Create a bound scalar node from a scalar function and already-bound children. pub fn try_new( scalar_fn: ScalarFnRef, - children: impl IntoIterator, - ) -> VortexResult { - let children = Vec::from_iter(children); + children: impl IntoIterator, + ) -> VortexResult { + let children: Box<[_]> = children.into_iter().collect(); vortex_ensure!( scalar_fn.signature().arity().matches(children.len()), "Expression arity mismatch: expected {} children but got {}", @@ -130,24 +150,24 @@ impl BoundExpression { .collect_vec(); let dtype = scalar_fn.return_dtype(&arg_dtypes)?; - Ok(Self::Scalar { + Ok(Arc::new(Self::Scalar { dtype, scalar_fn, - children: children.into(), - }) + children, + })) } /// Rebuild this node with new bound children, recomputing its dtype. pub fn with_children( - self, - children: impl IntoIterator, - ) -> VortexResult { - let children = Vec::from_iter(children); + self: BoundExpressionRef, + children: impl IntoIterator, + ) -> VortexResult { + let children: Box<[_]> = children.into_iter().collect(); let BoundExpression::Scalar { dtype, scalar_fn, children: old_children, - } = &self + } = self.as_ref() else { vortex_ensure!( children.is_empty(), @@ -164,11 +184,11 @@ impl BoundExpression { .zip(old_children.iter()) .all(|(new, old)| new.dtype() == old.dtype()) { - return Ok(Self::Scalar { + return Ok(Arc::new(Self::Scalar { dtype: dtype.clone(), scalar_fn: scalar_fn.clone(), - children: children.into(), - }); + children, + })); } Self::try_new(scalar_fn.clone(), children) @@ -182,15 +202,15 @@ impl BoundExpression { } /// The bound children of this node, in argument order. Empty for [`BoundExpression::Root`]. - pub fn children(&self) -> &[BoundExpression] { + pub fn children(&self) -> &[BoundExpressionRef] { match self { - Self::Scalar { children, .. } => children.as_slice(), + Self::Scalar { children, .. } => children, Self::Root { .. } => &[], } } /// Return the child at `index`. - pub fn child(&self, index: usize) -> &BoundExpression { + pub fn child(&self, index: usize) -> &BoundExpressionRef { &self.children()[index] } @@ -209,15 +229,7 @@ impl BoundExpression { /// Return whether this expression tree contains a node using the given scalar-function vtable. pub fn contains(&self) -> VortexResult { - let mut contains = false; - pre_order_visit_down(self, |node| { - if node.is::() { - contains = true; - return Ok(TraversalOrder::Stop); - } - Ok(TraversalOrder::Continue) - })?; - Ok(contains) + Ok(self.any_node(|node| node.is::())) } /// Return the typed scalar-function options when this node uses the given vtable. @@ -244,32 +256,41 @@ impl BoundExpression { /// /// Expressions without a scope root, such as literals, match every dtype. pub fn is_root_bound_to(&self, dtype: &DType) -> bool { - let mut is_bound_to = true; - pre_order_visit_down(self, |node| { - if node.is_root() && node.dtype() != dtype { - is_bound_to = false; - return Ok(TraversalOrder::Stop); - } - Ok(TraversalOrder::Continue) - }) - .vortex_expect("bound expression traversal cannot not fail"); - is_bound_to + !self.any_node(|node| node.is_root() && node.dtype() != dtype) } /// Return an expression that proves this predicate is definitely false from statistics. - pub fn falsify(&self, session: &VortexSession) -> VortexResult> { - StatsRewriteCtx::new(session).falsify(self) + pub fn falsify( + self: BoundExpressionRef, + session: &VortexSession, + ) -> VortexResult> { + StatsRewriteCtx::new(session).falsify(&self) } /// Return an expression that proves this predicate is definitely true from statistics. - pub fn satisfy(&self, session: &VortexSession) -> VortexResult> { - StatsRewriteCtx::new(session).satisfy(self) + pub fn satisfy( + self: BoundExpressionRef, + session: &VortexSession, + ) -> VortexResult> { + StatsRewriteCtx::new(session).satisfy(&self) } /// Display the bound expression as a formatted tree structure. pub fn display_tree(&self) -> impl Display { DisplayTreeExpr(self) } + + /// Return whether any node of this tree satisfies `predicate`, walking iteratively. + fn any_node(&self, mut predicate: impl FnMut(&BoundExpression) -> bool) -> bool { + let mut stack = vec![self]; + while let Some(node) = stack.pop() { + if predicate(node) { + return true; + } + stack.extend(node.children().iter().map(Arc::as_ref)); + } + false + } } impl Display for BoundExpression { @@ -287,12 +308,12 @@ impl Expression { /// The returned tree carries a dtype on each node, so callers needing types at more than one /// node should bind once and read fields rather than calling /// [`return_dtype`](Expression::return_dtype) repeatedly. - pub fn bind(&self, dtype: &DType) -> VortexResult { + pub fn bind(&self, dtype: &DType) -> VortexResult { self.bind_scope(&Scope::new(dtype.clone())) } /// Bind this expression against an explicit [`Scope`]. - pub fn bind_scope(&self, scope: &Scope) -> VortexResult { + pub fn bind_scope(&self, scope: &Scope) -> VortexResult { if self.is_root() { return Ok(BoundExpression::new_root(scope.root().clone())); } @@ -315,16 +336,15 @@ impl Drop for BoundExpression { let Self::Scalar { children, .. } = self else { return; }; - let Some(children) = Arc::get_mut(children) else { - return; - }; - let mut to_drop = std::mem::take(children); - while let Some(mut child) = to_drop.pop() { - if let BoundExpression::Scalar { children, .. } = &mut child - && let Some(grandchildren) = Arc::get_mut(children) + let mut to_drop = std::mem::take(children).into_vec(); + while let Some(child) = to_drop.pop() { + // Descending is only useful for the last owner of a subtree; releasing a shared + // handle is O(1) and leaves the children alone. + if let Some(mut node) = Arc::into_inner(child) + && let Self::Scalar { children, .. } = &mut node { - to_drop.append(grandchildren); + to_drop.append(&mut std::mem::take(children).into_vec()); } } } @@ -337,6 +357,7 @@ mod tests { use super::*; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::expr::bound; use crate::expr::col; use crate::expr::eq; use crate::expr::lit; @@ -421,18 +442,24 @@ mod tests { } #[test] - fn clone_shares_children() -> VortexResult<()> { + fn clone_shares_the_tree() -> VortexResult<()> { let bound = eq(col("a"), lit(1_i32)).bind_scope(&scope())?; - let cloned = bound.clone(); + let cloned = Arc::clone(&bound); - let ( - BoundExpression::Scalar { children: a, .. }, - BoundExpression::Scalar { children: b, .. }, - ) = (&bound, &cloned) - else { - unreachable!("eq is a scalar node") - }; - assert!(Arc::ptr_eq(a, b)); + assert!(Arc::ptr_eq(&bound, &cloned)); + Ok(()) + } + + #[test] + fn rebuilding_a_node_shares_untouched_children() -> VortexResult<()> { + let bound = eq(col("a"), lit(1_i32)).bind_scope(&scope())?; + let children = bound.children().to_vec(); + let rebuilt = Arc::clone(&bound).with_children(children)?; + + assert!(!Arc::ptr_eq(&bound, &rebuilt)); + for (old, new) in bound.children().iter().zip(rebuilt.children()) { + assert!(Arc::ptr_eq(old, new)); + } Ok(()) } @@ -452,11 +479,24 @@ mod tests { let independently_bound = expr.bind_scope(&scope())?; assert_eq!(bound, independently_bound); - assert_eq!(ExactBoundExpr(bound.clone()), ExactBoundExpr(bound.clone())); + assert_eq!( + ExactBoundExpr(Arc::clone(&bound)), + ExactBoundExpr(Arc::clone(&bound)) + ); assert_ne!(ExactBoundExpr(bound), ExactBoundExpr(independently_bound)); Ok(()) } + #[test] + fn deep_trees_drop_without_overflowing_the_stack() -> VortexResult<()> { + let mut expr = lit(true).bind(&struct_dtype())?; + for _ in 0..100_000 { + expr = bound::not(expr); + } + drop(expr); + Ok(()) + } + #[test] fn binding_reports_a_type_error() { let expr = eq(col("a"), lit("nope")); diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index 250d7063834..c08ddaac977 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -41,7 +41,7 @@ impl ExprDisplay for Expression { impl ExprDisplay for BoundExpression { fn display_child(&self, index: usize) -> &dyn ExprDisplay { - &self.children()[index] + self.children()[index].as_ref() } fn display_children_count(&self) -> usize { @@ -50,7 +50,9 @@ impl ExprDisplay for BoundExpression { } trait DisplayTreeNode: Sized { - fn tree_children(&self) -> &[Self]; + fn tree_child(&self, index: usize) -> &Self; + + fn tree_children_count(&self) -> usize; fn tree_child_name(&self, index: usize) -> ChildName; @@ -61,8 +63,12 @@ trait DisplayTreeNode: Sized { const ROOT_DISPLAY: &str = "vortex.root()"; impl DisplayTreeNode for Expression { - fn tree_children(&self) -> &[Self] { - Expression::children(self) + fn tree_child(&self, index: usize) -> &Self { + Expression::child(self, index) + } + + fn tree_children_count(&self) -> usize { + Expression::children(self).len() } fn tree_child_name(&self, index: usize) -> ChildName { @@ -81,8 +87,12 @@ impl DisplayTreeNode for Expression { } impl DisplayTreeNode for BoundExpression { - fn tree_children(&self) -> &[Self] { - BoundExpression::children(self) + fn tree_child(&self, index: usize) -> &Self { + BoundExpression::child(self, index) + } + + fn tree_children_count(&self) -> usize { + BoundExpression::children(self).len() } fn tree_child_name(&self, index: usize) -> ChildName { @@ -120,10 +130,14 @@ impl TreeDisplayAdapter for DisplayTreeExpr<'_, T> { node: &Self::Node, visit: &mut dyn FnMut(&str, &Self::Node, bool) -> fmt::Result, ) -> fmt::Result { - let children = node.tree_children(); - for (index, child) in children.iter().enumerate() { + let children_count = node.tree_children_count(); + for index in 0..children_count { let child_name = node.tree_child_name(index); - visit(child_name.as_ref(), child, index + 1 == children.len())?; + visit( + child_name.as_ref(), + node.tree_child(index), + index + 1 == children_count, + )?; } Ok(()) } diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index fb8bfe227aa..7a28780252c 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -16,6 +16,7 @@ use crate::dtype::FieldName; use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::expr::BoundExpression; +use crate::expr::BoundExpressionRef; use crate::expr::Expression; use crate::scalar::Scalar; use crate::scalar::ScalarValue; @@ -66,7 +67,7 @@ pub fn root() -> Expression { } /// Creates a bound expression that references a root scope with the given dtype. -pub fn bound_root(dtype: DType) -> BoundExpression { +pub fn bound_root(dtype: DType) -> BoundExpressionRef { BoundExpression::new_root(dtype) } @@ -99,7 +100,7 @@ pub fn lit(value: impl Into) -> Expression { } /// Creates a bound literal expression. -pub fn bound_lit(value: impl Into) -> BoundExpression { +pub fn bound_lit(value: impl Into) -> BoundExpressionRef { Literal .try_new_bound_expr(value.into(), []) .vortex_expect("literal expressions are always well-typed") @@ -120,7 +121,7 @@ pub fn col(field: impl Into) -> Expression { } /// Creates a bound expression that accesses a field from a root scope with the given dtype. -pub fn bound_col(field: impl Into, scope: DType) -> BoundExpression { +pub fn bound_col(field: impl Into, scope: DType) -> BoundExpressionRef { bound_get_item(field, bound_root(scope)) } @@ -137,7 +138,10 @@ pub fn get_item(field: impl Into, child: Expression) -> Expression { } /// Creates a bound expression that extracts a named field from a struct expression. -pub fn bound_get_item(field: impl Into, child: BoundExpression) -> BoundExpression { +pub fn bound_get_item( + field: impl Into, + child: BoundExpressionRef, +) -> BoundExpressionRef { GetItem .try_new_bound_expr(field.into(), [child]) .vortex_expect("get-item expressions must reference a field in the child dtype") @@ -159,10 +163,10 @@ pub fn variant_get( /// Creates a bound expression that extracts a path from a Variant expression. pub fn bound_variant_get( - child: BoundExpression, + child: BoundExpressionRef, path: impl Into, dtype: Option, -) -> BoundExpression { +) -> BoundExpressionRef { VariantGet .try_new_bound_expr(VariantGetOptions::new(path.into(), dtype), [child]) .vortex_expect("variant-get expressions require a Variant child") @@ -185,10 +189,10 @@ pub fn case_when( /// Creates a bound CASE WHEN expression with one WHEN/THEN pair and an ELSE value. pub fn bound_case_when( - condition: BoundExpression, - then_value: BoundExpression, - else_value: BoundExpression, -) -> BoundExpression { + condition: BoundExpressionRef, + then_value: BoundExpressionRef, + else_value: BoundExpressionRef, +) -> BoundExpressionRef { let options = CaseWhenOptions { num_when_then_pairs: 1, has_else: true, @@ -209,9 +213,9 @@ pub fn case_when_no_else(condition: Expression, then_value: Expression) -> Expre /// Creates a bound CASE WHEN expression with one WHEN/THEN pair and no ELSE value. pub fn bound_case_when_no_else( - condition: BoundExpression, - then_value: BoundExpression, -) -> BoundExpression { + condition: BoundExpressionRef, + then_value: BoundExpressionRef, +) -> BoundExpressionRef { let options = CaseWhenOptions { num_when_then_pairs: 1, has_else: false, @@ -253,9 +257,9 @@ pub fn nested_case_when( /// Creates a bound n-ary CASE WHEN expression from WHEN/THEN pairs and an optional ELSE value. pub fn bound_nested_case_when( - when_then_pairs: Vec<(BoundExpression, BoundExpression)>, - else_value: Option, -) -> BoundExpression { + when_then_pairs: Vec<(BoundExpressionRef, BoundExpressionRef)>, + else_value: Option, +) -> BoundExpressionRef { assert!( !when_then_pairs.is_empty(), "nested_case_when requires at least one when/then pair" @@ -295,9 +299,9 @@ pub fn binary(operator: Operator, lhs: Expression, rhs: Expression) -> Expressio /// Creates a bound binary expression with the given operator. pub fn bound_binary( operator: Operator, - lhs: BoundExpression, - rhs: BoundExpression, -) -> BoundExpression { + lhs: BoundExpressionRef, + rhs: BoundExpressionRef, +) -> BoundExpressionRef { Binary .try_new_bound_expr(operator, [lhs, rhs]) .vortex_expect("binary expressions must have compatible operand dtypes") @@ -331,7 +335,7 @@ pub fn eq(lhs: Expression, rhs: Expression) -> Expression { } /// Creates a bound equality expression. -pub fn bound_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { +pub fn bound_eq(lhs: BoundExpressionRef, rhs: BoundExpressionRef) -> BoundExpressionRef { bound_binary(Operator::Eq, lhs, rhs) } @@ -363,7 +367,7 @@ pub fn not_eq(lhs: Expression, rhs: Expression) -> Expression { } /// Creates a bound inequality expression. -pub fn bound_not_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { +pub fn bound_not_eq(lhs: BoundExpressionRef, rhs: BoundExpressionRef) -> BoundExpressionRef { bound_binary(Operator::NotEq, lhs, rhs) } @@ -395,7 +399,7 @@ pub fn gt_eq(lhs: Expression, rhs: Expression) -> Expression { } /// Creates a bound greater-than-or-equal expression. -pub fn bound_gt_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { +pub fn bound_gt_eq(lhs: BoundExpressionRef, rhs: BoundExpressionRef) -> BoundExpressionRef { bound_binary(Operator::Gte, lhs, rhs) } @@ -427,7 +431,7 @@ pub fn gt(lhs: Expression, rhs: Expression) -> Expression { } /// Creates a bound greater-than expression. -pub fn bound_gt(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { +pub fn bound_gt(lhs: BoundExpressionRef, rhs: BoundExpressionRef) -> BoundExpressionRef { bound_binary(Operator::Gt, lhs, rhs) } @@ -459,7 +463,7 @@ pub fn lt_eq(lhs: Expression, rhs: Expression) -> Expression { } /// Creates a bound less-than-or-equal expression. -pub fn bound_lt_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { +pub fn bound_lt_eq(lhs: BoundExpressionRef, rhs: BoundExpressionRef) -> BoundExpressionRef { bound_binary(Operator::Lte, lhs, rhs) } @@ -491,7 +495,7 @@ pub fn lt(lhs: Expression, rhs: Expression) -> Expression { } /// Creates a bound less-than expression. -pub fn bound_lt(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { +pub fn bound_lt(lhs: BoundExpressionRef, rhs: BoundExpressionRef) -> BoundExpressionRef { bound_binary(Operator::Lt, lhs, rhs) } @@ -521,7 +525,7 @@ pub fn or(lhs: Expression, rhs: Expression) -> Expression { } /// Creates a bound boolean OR expression. -pub fn bound_or(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { +pub fn bound_or(lhs: BoundExpressionRef, rhs: BoundExpressionRef) -> BoundExpressionRef { bound_binary(Operator::Or, lhs, rhs) } @@ -539,9 +543,9 @@ where } /// Collects bound expressions into a balanced tree of boolean OR expressions. -pub fn bound_or_collect(iter: I) -> Option +pub fn bound_or_collect(iter: I) -> Option where - I: IntoIterator, + I: IntoIterator, { iter.into_iter().reduce_balanced(bound_or) } @@ -572,7 +576,7 @@ pub fn and(lhs: Expression, rhs: Expression) -> Expression { } /// Creates a bound boolean AND expression. -pub fn bound_and(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { +pub fn bound_and(lhs: BoundExpressionRef, rhs: BoundExpressionRef) -> BoundExpressionRef { bound_binary(Operator::And, lhs, rhs) } @@ -590,9 +594,9 @@ where } /// Collects bound expressions into a balanced tree of boolean AND expressions. -pub fn bound_and_collect(iter: I) -> Option +pub fn bound_and_collect(iter: I) -> Option where - I: IntoIterator, + I: IntoIterator, { iter.into_iter().reduce_balanced(bound_and) } @@ -638,7 +642,7 @@ pub fn checked_add(lhs: Expression, rhs: Expression) -> Expression { } /// Creates a bound checked-add expression. -pub fn bound_checked_add(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { +pub fn bound_checked_add(lhs: BoundExpressionRef, rhs: BoundExpressionRef) -> BoundExpressionRef { bound_binary(Operator::Add, lhs, rhs) } @@ -657,7 +661,7 @@ pub fn not(operand: Expression) -> Expression { } /// Creates a bound expression that logically inverts boolean values. -pub fn bound_not(operand: BoundExpression) -> BoundExpression { +pub fn bound_not(operand: BoundExpressionRef) -> BoundExpressionRef { Not.try_new_bound_expr(EmptyOptions, [operand]) .vortex_expect("not expressions require a boolean operand") } @@ -692,11 +696,11 @@ pub fn between( /// Creates a bound expression that checks if values are between two bounds. pub fn bound_between( - arr: BoundExpression, - lower: BoundExpression, - upper: BoundExpression, + arr: BoundExpressionRef, + lower: BoundExpressionRef, + upper: BoundExpressionRef, options: BetweenOptions, -) -> BoundExpression { +) -> BoundExpressionRef { Between .try_new_bound_expr(options, [arr, lower, upper]) .vortex_expect("between expressions require compatible operand dtypes") @@ -718,7 +722,10 @@ pub fn select(field_names: impl Into, child: Expression) -> Expressi } /// Creates a bound expression that selects specific fields from a struct expression. -pub fn bound_select(field_names: impl Into, child: BoundExpression) -> BoundExpression { +pub fn bound_select( + field_names: impl Into, + child: BoundExpressionRef, +) -> BoundExpressionRef { Select .try_new_bound_expr(FieldSelection::Include(field_names.into()), [child]) .vortex_expect("select expressions require fields from a struct child") @@ -741,8 +748,8 @@ pub fn select_exclude(fields: impl Into, child: Expression) -> Expre /// Creates a bound expression that excludes specific fields from a struct expression. pub fn bound_select_exclude( fields: impl Into, - child: BoundExpression, -) -> BoundExpression { + child: BoundExpressionRef, +) -> BoundExpressionRef { Select .try_new_bound_expr(FieldSelection::Exclude(fields.into()), [child]) .vortex_expect("select expressions require fields from a struct child") @@ -776,9 +783,9 @@ pub fn pack( /// Creates a bound expression that packs values into a struct with named fields. pub fn bound_pack( - elements: impl IntoIterator, BoundExpression)>, + elements: impl IntoIterator, BoundExpressionRef)>, nullability: Nullability, -) -> BoundExpression { +) -> BoundExpressionRef { let (names, values): (Vec<_>, Vec<_>) = elements .into_iter() .map(|(name, value)| (name.into(), value)) @@ -810,7 +817,7 @@ pub fn cast(child: Expression, target: DType) -> Expression { } /// Creates a bound expression that casts values to a target dtype. -pub fn bound_cast(child: BoundExpression, target: DType) -> BoundExpression { +pub fn bound_cast(child: BoundExpressionRef, target: DType) -> BoundExpressionRef { Cast.try_new_bound_expr(target, [child]) .vortex_expect("cast expressions require a supported source and target dtype") } @@ -828,7 +835,10 @@ pub fn fill_null(child: Expression, fill_value: Expression) -> Expression { } /// Creates a bound expression that replaces null values with a fill value. -pub fn bound_fill_null(child: BoundExpression, fill_value: BoundExpression) -> BoundExpression { +pub fn bound_fill_null( + child: BoundExpressionRef, + fill_value: BoundExpressionRef, +) -> BoundExpressionRef { FillNull .try_new_bound_expr(EmptyOptions, [child, fill_value]) .vortex_expect("fill-null expressions require compatible child and fill dtypes") @@ -849,7 +859,7 @@ pub fn is_null(child: Expression) -> Expression { } /// Creates a bound expression that checks for null values. -pub fn bound_is_null(child: BoundExpression) -> BoundExpression { +pub fn bound_is_null(child: BoundExpressionRef) -> BoundExpressionRef { IsNull .try_new_bound_expr(EmptyOptions, [child]) .vortex_expect("is-null expressions are always well-typed") @@ -870,7 +880,7 @@ pub fn is_not_null(child: Expression) -> Expression { } /// Creates a bound expression that checks for non-null values. -pub fn bound_is_not_null(child: BoundExpression) -> BoundExpression { +pub fn bound_is_not_null(child: BoundExpressionRef) -> BoundExpressionRef { IsNotNull .try_new_bound_expr(EmptyOptions, [child]) .vortex_expect("is-not-null expressions are always well-typed") @@ -890,7 +900,7 @@ pub fn like(child: Expression, pattern: Expression) -> Expression { } /// Creates a bound SQL LIKE expression. -pub fn bound_like(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { +pub fn bound_like(child: BoundExpressionRef, pattern: BoundExpressionRef) -> BoundExpressionRef { bound_like_with_options(child, pattern, false, false) } @@ -906,7 +916,7 @@ pub fn ilike(child: Expression, pattern: Expression) -> Expression { } /// Creates a bound case-insensitive SQL ILIKE expression. -pub fn bound_ilike(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { +pub fn bound_ilike(child: BoundExpressionRef, pattern: BoundExpressionRef) -> BoundExpressionRef { bound_like_with_options(child, pattern, false, true) } @@ -922,7 +932,10 @@ pub fn not_like(child: Expression, pattern: Expression) -> Expression { } /// Creates a bound negated SQL NOT LIKE expression. -pub fn bound_not_like(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { +pub fn bound_not_like( + child: BoundExpressionRef, + pattern: BoundExpressionRef, +) -> BoundExpressionRef { bound_like_with_options(child, pattern, true, false) } @@ -938,16 +951,19 @@ pub fn not_ilike(child: Expression, pattern: Expression) -> Expression { } /// Creates a bound negated case-insensitive SQL NOT ILIKE expression. -pub fn bound_not_ilike(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { +pub fn bound_not_ilike( + child: BoundExpressionRef, + pattern: BoundExpressionRef, +) -> BoundExpressionRef { bound_like_with_options(child, pattern, true, true) } fn bound_like_with_options( - child: BoundExpression, - pattern: BoundExpression, + child: BoundExpressionRef, + pattern: BoundExpressionRef, negated: bool, case_insensitive: bool, -) -> BoundExpression { +) -> BoundExpressionRef { Like.try_new_bound_expr( LikeOptions { negated, @@ -966,7 +982,7 @@ pub fn mask(array: Expression, mask: Expression) -> Expression { } /// Creates a bound mask expression. -pub fn bound_mask(array: BoundExpression, mask: BoundExpression) -> BoundExpression { +pub fn bound_mask(array: BoundExpressionRef, mask: BoundExpressionRef) -> BoundExpressionRef { Mask.try_new_bound_expr(EmptyOptions, [array, mask]) .vortex_expect("mask expressions require a boolean mask") } @@ -990,7 +1006,7 @@ pub fn merge(elements: impl IntoIterator>) -> Expre } /// Creates a bound expression that merges struct expressions. -pub fn bound_merge(elements: impl IntoIterator) -> BoundExpression { +pub fn bound_merge(elements: impl IntoIterator) -> BoundExpressionRef { bound_merge_opts(elements, DuplicateHandling::default()) } @@ -1006,9 +1022,9 @@ pub fn merge_opts( /// Creates a bound merge expression with explicit duplicate handling. pub fn bound_merge_opts( - elements: impl IntoIterator, + elements: impl IntoIterator, duplicate_handling: DuplicateHandling, -) -> BoundExpression { +) -> BoundExpressionRef { Merge .try_new_bound_expr(duplicate_handling, elements) .vortex_expect("merge expressions require non-nullable struct children") @@ -1028,10 +1044,10 @@ pub fn zip_expr(mask: Expression, if_true: Expression, if_false: Expression) -> /// Creates a bound zip expression that conditionally selects between two arrays. pub fn bound_zip_expr( - mask: BoundExpression, - if_true: BoundExpression, - if_false: BoundExpression, -) -> BoundExpression { + mask: BoundExpressionRef, + if_true: BoundExpressionRef, + if_false: BoundExpressionRef, +) -> BoundExpressionRef { Zip.try_new_bound_expr(EmptyOptions, [if_true, if_false, mask]) .vortex_expect("zip expressions require a boolean mask and compatible value dtypes") } @@ -1046,8 +1062,8 @@ pub fn dynamic_with_options(options: DynamicComparisonExpr, lhs: Expression) -> /// Creates a bound dynamic comparison expression from its complete options. pub fn bound_dynamic_with_options( options: DynamicComparisonExpr, - lhs: BoundExpression, -) -> BoundExpression { + lhs: BoundExpressionRef, +) -> BoundExpressionRef { DynamicComparison .try_new_bound_expr(options, [lhs]) .vortex_expect("dynamic comparisons require a compatible left-hand dtype") @@ -1080,8 +1096,8 @@ pub fn bound_dynamic( rhs_value: impl Fn() -> Option + Send + Sync + 'static, rhs_dtype: DType, default: bool, - lhs: BoundExpression, -) -> BoundExpression { + lhs: BoundExpressionRef, +) -> BoundExpressionRef { bound_dynamic_with_options( DynamicComparisonExpr { operator, @@ -1110,7 +1126,10 @@ pub fn list_contains(list: Expression, value: Expression) -> Expression { } /// Creates a bound expression that checks if a value is contained in a list. -pub fn bound_list_contains(list: BoundExpression, value: BoundExpression) -> BoundExpression { +pub fn bound_list_contains( + list: BoundExpressionRef, + value: BoundExpressionRef, +) -> BoundExpressionRef { ListContains .try_new_bound_expr(EmptyOptions, [list, value]) .vortex_expect("list-contains expressions require a compatible list and value dtype") @@ -1130,7 +1149,7 @@ pub fn byte_length(input: Expression) -> Expression { } /// Creates a bound expression that computes each element's byte length. -pub fn bound_byte_length(input: BoundExpression) -> BoundExpression { +pub fn bound_byte_length(input: BoundExpressionRef) -> BoundExpressionRef { ByteLength .try_new_bound_expr(EmptyOptions, [input]) .vortex_expect("byte-length expressions require a variable-length binary child") @@ -1149,7 +1168,7 @@ pub fn ext_storage(input: Expression) -> Expression { } /// Creates a bound expression that extracts an extension array's storage values. -pub fn bound_ext_storage(input: BoundExpression) -> BoundExpression { +pub fn bound_ext_storage(input: BoundExpressionRef) -> BoundExpressionRef { ExtStorage .try_new_bound_expr(EmptyOptions, [input]) .vortex_expect("extension-storage expressions require an extension child") @@ -1170,7 +1189,7 @@ pub fn list_length(input: Expression) -> Expression { } /// Creates a bound expression that computes the number of elements in each list. -pub fn bound_list_length(input: BoundExpression) -> BoundExpression { +pub fn bound_list_length(input: BoundExpressionRef) -> BoundExpressionRef { ListLength .try_new_bound_expr(EmptyOptions, [input]) .vortex_expect("list-length expressions require a list child") @@ -1195,7 +1214,7 @@ pub fn list_sum(input: Expression) -> Expression { } /// Creates a bound expression that sums the elements of each list. -pub fn bound_list_sum(input: BoundExpression) -> BoundExpression { +pub fn bound_list_sum(input: BoundExpressionRef) -> BoundExpressionRef { ListSum .try_new_bound_expr(NumericalAggregateOpts::default(), [input]) .vortex_expect("list-sum expressions require a numeric list child") @@ -1209,9 +1228,9 @@ pub fn list_sum_opts(input: Expression, options: NumericalAggregateOpts) -> Expr /// Creates a bound list-sum expression with explicit aggregate options. pub fn bound_list_sum_opts( - input: BoundExpression, + input: BoundExpressionRef, options: NumericalAggregateOpts, -) -> BoundExpression { +) -> BoundExpressionRef { ListSum .try_new_bound_expr(options, [input]) .vortex_expect("list-sum expressions require a numeric list child") diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index 59fd21c46ee..06bfa410342 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -296,19 +296,29 @@ mod tests { let root = bound::root(scope.clone()); let value = bound::get_item("value", root); let literal = bound::lit(5i32); - let condition = bound::gt(value.clone(), literal.clone()); + let condition = bound::gt(Arc::clone(&value), Arc::clone(&literal)); assert_eq!(condition.dtype(), &DType::Bool(Nullability::NonNullable)); - assert_eq!(condition.children(), &[value.clone(), literal.clone()]); + assert_eq!( + condition.children(), + &[Arc::clone(&value), Arc::clone(&literal)] + ); - let case = bound::case_when(condition.clone(), value.clone(), literal.clone()); + let case = bound::case_when( + Arc::clone(&condition), + Arc::clone(&value), + Arc::clone(&literal), + ); assert_eq!(case.dtype(), &value_dtype); - assert_eq!(case.children(), &[condition.clone(), value, literal]); + assert_eq!(case.children(), &[Arc::clone(&condition), value, literal]); let packed = bound::pack( - [("condition", condition.clone()), ("value", case.clone())], + [ + ("condition", Arc::clone(&condition)), + ("value", Arc::clone(&case)), + ], Nullability::NonNullable, ); - assert_eq!(packed.children(), &[condition, case.clone()]); + assert_eq!(packed.children(), &[condition, Arc::clone(&case)]); assert_eq!( packed.dtype(), &DType::Struct( diff --git a/vortex-array/src/expr/transform/bound_partition.rs b/vortex-array/src/expr/transform/bound_partition.rs index 3272be862bc..d66898a6000 100644 --- a/vortex-array/src/expr/transform/bound_partition.rs +++ b/vortex-array/src/expr/transform/bound_partition.rs @@ -4,6 +4,7 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hash; +use std::sync::Arc; use itertools::Itertools; use vortex_error::VortexExpect; @@ -17,6 +18,7 @@ use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::StructFields; use crate::expr::BoundExpression; +use crate::expr::BoundExpressionRef; use crate::expr::ExactBoundExpr; use crate::expr::analysis::Annotation; use crate::expr::analysis::AnnotationFn; @@ -41,8 +43,8 @@ use crate::expr::traversal::TraversalOrder; /// root. /// /// See . -pub fn partition_bound>( - expr: BoundExpression, +pub fn partition_bound>( + expr: BoundExpressionRef, annotate_fn: A, ) -> VortexResult> where @@ -58,7 +60,7 @@ where /// /// Prefer [`partition_bound`] when annotations can be derived by an [`AnnotationFn`]. pub fn partition_bound_annotations( - expr: BoundExpression, + expr: BoundExpressionRef, annotations: BoundAnnotations, ) -> VortexResult> where @@ -66,7 +68,7 @@ where FieldName: From, { let mut collector = PartitionCollector::::new(&annotations); - expr.clone().rewrite(&mut collector)?; + Arc::clone(&expr).rewrite(&mut collector)?; let mut partitions = Vec::with_capacity(collector.sub_expressions.len()); let mut partition_annotations = Vec::with_capacity(collector.sub_expressions.len()); @@ -107,9 +109,9 @@ where #[derive(Debug)] pub struct BoundPartitionedExpr { /// The root expression used to re-assemble the results. - pub root: BoundExpression, + pub root: BoundExpressionRef, /// The partition expressions themselves. - pub partitions: Box<[BoundExpression]>, + pub partitions: Box<[BoundExpressionRef]>, /// The field name of each partition as referenced in the root expression. pub partition_names: FieldNames, /// The annotation associated with each partition. @@ -137,7 +139,7 @@ where { /// Return the partition for a given field, if it exists. // FIXME(ngates): this should return an iterator since an annotation may have multiple partitions. - pub fn find_partition(&self, id: &A) -> Option<&BoundExpression> { + pub fn find_partition(&self, id: &A) -> Option<&BoundExpressionRef> { let id = FieldName::from(id.clone()); self.partition_names .iter() @@ -146,7 +148,10 @@ where } /// Replace the partition expressions and update every root dtype in the recombination tree. - pub fn replace_partitions(&mut self, partitions: Box<[BoundExpression]>) -> VortexResult<()> { + pub fn replace_partitions( + &mut self, + partitions: Box<[BoundExpressionRef]>, + ) -> VortexResult<()> { vortex_ensure!( partitions.len() == self.partition_names.len(), "Expected {} partitions, got {}", @@ -155,7 +160,7 @@ where ); let root_dtype = partition_root_dtype(&self.partition_names, &partitions); - let root = replace_root_dtype(self.root.clone(), root_dtype)?; + let root = replace_root_dtype(Arc::clone(&self.root), root_dtype)?; self.partitions = partitions; self.root = root; Ok(()) @@ -165,7 +170,7 @@ where #[derive(Debug)] struct PartitionCollector<'a, A: Annotation> { annotations: &'a BoundAnnotations, - sub_expressions: HashMap>, + sub_expressions: HashMap>, } impl<'a, A: Annotation + Display> PartitionCollector<'a, A> { @@ -187,10 +192,10 @@ impl NodeRewriter for PartitionCollector<'_, A> where FieldName: From, { - type NodeTy = BoundExpression; + type NodeTy = BoundExpressionRef; fn visit_down(&mut self, node: Self::NodeTy) -> VortexResult> { - match self.annotations.get(&ExactBoundExpr(node.clone())) { + match self.annotations.get(&ExactBoundExpr(Arc::clone(&node))) { // If this expression only accesses a single field, then we can skip the children Some(annotations) if annotations.len() == 1 => { let annotation = annotations @@ -198,7 +203,7 @@ where .next() .vortex_expect("expected one field"); let sub_exprs = self.sub_expressions.entry(annotation.clone()).or_default(); - sub_exprs.push(node.clone()); + sub_exprs.push(Arc::clone(&node)); Ok(Transformed { value: node, changed: false, @@ -236,10 +241,10 @@ impl NodeRewriter for PartitionRootRewriter<'_, A> where FieldName: From, { - type NodeTy = BoundExpression; + type NodeTy = BoundExpressionRef; fn visit_down(&mut self, node: Self::NodeTy) -> VortexResult> { - let Some(annotations) = self.annotations.get(&ExactBoundExpr(node.clone())) else { + let Some(annotations) = self.annotations.get(&ExactBoundExpr(Arc::clone(&node))) else { return Ok(Transformed::no(node)); }; if annotations.len() != 1 { @@ -271,7 +276,7 @@ where } } -fn partition_root_dtype(names: &FieldNames, partitions: &[BoundExpression]) -> DType { +fn partition_root_dtype(names: &FieldNames, partitions: &[BoundExpressionRef]) -> DType { DType::Struct( StructFields::new( names.clone(), @@ -284,7 +289,10 @@ fn partition_root_dtype(names: &FieldNames, partitions: &[BoundExpression]) -> D ) } -fn replace_root_dtype(expr: BoundExpression, root_dtype: DType) -> VortexResult { +fn replace_root_dtype( + expr: BoundExpressionRef, + root_dtype: DType, +) -> VortexResult { Ok(expr .transform_down(|node| { if node.is_root() { @@ -340,7 +348,7 @@ mod tests { } fn partition_by_field( - expr: BoundExpression, + expr: BoundExpressionRef, dtype: &DType, ) -> VortexResult> { let fields = dtype.as_struct_fields_opt().unwrap(); diff --git a/vortex-array/src/expr/traversal/mod.rs b/vortex-array/src/expr/traversal/mod.rs index 0f55600afe5..616e280d912 100644 --- a/vortex-array/src/expr/traversal/mod.rs +++ b/vortex-array/src/expr/traversal/mod.rs @@ -22,7 +22,7 @@ pub use visitor::pre_order_visit_down; pub use visitor::pre_order_visit_up; use vortex_error::VortexResult; -use crate::expr::BoundExpression; +use crate::expr::BoundExpressionRef; use crate::expr::Expression; use crate::expr::traversal::fold::NodeFolderContextWrapper; @@ -528,16 +528,12 @@ impl Node for Expression { } } -impl Node for BoundExpression { +impl Node for BoundExpressionRef { fn apply_children<'a, F: FnMut(&'a Self) -> VortexResult>( &'a self, mut f: F, ) -> VortexResult { - let BoundExpression::Scalar { children, .. } = self else { - return Ok(TraversalOrder::Continue); - }; - - for child in children.iter() { + for child in self.children() { match f(child)? { TraversalOrder::Continue | TraversalOrder::Skip => {} TraversalOrder::Stop => return Ok(TraversalOrder::Stop), @@ -551,32 +547,35 @@ impl Node for BoundExpression { self, mut f: F, ) -> VortexResult> { - let BoundExpression::Scalar { children, .. } = &self else { + if self.children().is_empty() { return Ok(Transformed::no(self)); - }; + } let mut order = TraversalOrder::Continue; // Stays `None` until a child actually changes. Most nodes of a rewritten tree are // untouched. let mut rewritten: Option> = None; - for (index, child) in children.iter().enumerate() { - let value = match order { - TraversalOrder::Continue | TraversalOrder::Skip => { - let result = f(child.clone())?; - order = result.order; - if result.changed && rewritten.is_none() { - let mut prefix = Vec::with_capacity(children.len()); - prefix.extend_from_slice(&children[..index]); - rewritten = Some(prefix); + { + let children = self.children(); + for (index, child) in children.iter().enumerate() { + let value = match order { + TraversalOrder::Continue | TraversalOrder::Skip => { + let result = f(Arc::clone(child))?; + order = result.order; + if result.changed && rewritten.is_none() { + let mut prefix = Vec::with_capacity(children.len()); + prefix.extend_from_slice(&children[..index]); + rewritten = Some(prefix); + } + result.value } - result.value - } - TraversalOrder::Stop => child.clone(), - }; + TraversalOrder::Stop => Arc::clone(child), + }; - if let Some(rewritten) = &mut rewritten { - rewritten.push(value); + if let Some(rewritten) = &mut rewritten { + rewritten.push(value); + } } } @@ -591,17 +590,11 @@ 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()), - } + f(&mut self.children().iter()) } fn children_count(&self) -> usize { - match self { - BoundExpression::Scalar { children, .. } => children.len(), - BoundExpression::Root { .. } => 0, - } + self.children().len() } } diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index d0590f2bf58..d49e12fb988 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -10,18 +10,19 @@ use crate::IntoArray; use crate::arrays::ConstantArray; use crate::arrays::ScalarFnArray; use crate::expr::BoundExpression; +use crate::expr::BoundExpressionRef; 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 { + pub fn apply_bound(self, expr: &BoundExpressionRef) -> VortexResult { let BoundExpression::Scalar { scalar_fn, children, .. - } = expr + } = expr.as_ref() else { return Ok(self); }; diff --git a/vortex-array/src/scalar_fn/fns/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index 87eb843bc02..9c6c48c07f0 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/vortex-array/src/scalar_fn/fns/dynamic.rs @@ -19,7 +19,7 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ConstantArray; use crate::dtype::DType; -use crate::expr::BoundExpression; +use crate::expr::BoundExpressionRef; use crate::expr::display::ExprDisplay; use crate::expr::traversal::NodeExt; use crate::expr::traversal::NodeVisitor; @@ -206,12 +206,12 @@ pub struct DynamicExprUpdates { impl DynamicExprUpdates { /// Track dynamic scalar functions contained in a bound expression tree. - pub fn new(expr: &BoundExpression) -> Option { + pub fn new(expr: &BoundExpressionRef) -> Option { #[derive(Default)] struct Visitor(Vec); impl NodeVisitor<'_> for Visitor { - type NodeTy = BoundExpression; + type NodeTy = BoundExpressionRef; fn visit_down(&mut self, node: &'_ Self::NodeTy) -> VortexResult { if let Some(dynamic) = node diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index 88b01e72e57..3c1c89fd140 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -22,6 +22,7 @@ use crate::arrays::ScalarFn; use crate::arrays::ScalarFnArray; use crate::dtype::DType; use crate::expr::BoundExpression; +use crate::expr::BoundExpressionRef; use crate::expr::Expression; use crate::expr::display::ExprDisplay; use crate::scalar_fn::ScalarFnId; @@ -512,8 +513,8 @@ pub trait ScalarFnVTableExt: ScalarFnVTable { fn try_new_bound_expr( &self, options: Self::Options, - children: impl IntoIterator, - ) -> VortexResult { + children: impl IntoIterator, + ) -> VortexResult { BoundExpression::try_new(self.bind(options), children) } } diff --git a/vortex-array/src/stats/bind.rs b/vortex-array/src/stats/bind.rs index e07a588de94..7f1a2dd9710 100644 --- a/vortex-array/src/stats/bind.rs +++ b/vortex-array/src/stats/bind.rs @@ -17,7 +17,7 @@ use vortex_error::VortexResult; use crate::aggregate_fn::AggregateFnRef; use crate::dtype::DType; -use crate::expr::BoundExpression; +use crate::expr::BoundExpressionRef; use crate::expr::bound::lit; use crate::expr::traversal::NodeExt; use crate::expr::traversal::Transformed; @@ -37,16 +37,16 @@ pub trait StatBinder { /// statistic is unavailable in their backing representation. fn bind_aggregate( &self, - input: &BoundExpression, + input: &BoundExpressionRef, aggregate_fn: &AggregateFnRef, stat_dtype: &DType, - ) -> VortexResult>; + ) -> VortexResult>; /// Expression to use when a stat is unavailable. /// /// The default is a nullable null literal, which preserves three-valued /// pruning semantics for stats-table execution. - fn missing_stat(&self, dtype: DType) -> VortexResult { + fn missing_stat(&self, dtype: DType) -> VortexResult { null_expr(dtype) } } @@ -57,9 +57,9 @@ pub trait StatBinder { /// are responsible for expressing stat semantics; binding maps aggregate-backed /// stat requests to the concrete stats representation supported by the binder. pub fn bind_stats( - predicate: BoundExpression, + predicate: BoundExpressionRef, binder: &B, -) -> VortexResult { +) -> VortexResult { Ok(predicate .transform_down(|expr| { if !expr.is::() { @@ -75,9 +75,9 @@ pub fn bind_stats( } fn bind_stat_fn( - expr: &BoundExpression, + expr: &BoundExpressionRef, binder: &(impl StatBinder + ?Sized), -) -> VortexResult> { +) -> VortexResult> { let options = expr.as_::(); let aggregate_fn = options.aggregate_fn(); // `StatFn` has exactly one child: the expression the aggregate statistic is computed over. @@ -86,7 +86,7 @@ fn bind_stat_fn( binder.bind_aggregate(input, aggregate_fn, expr.dtype()) } -fn null_expr(dtype: DType) -> VortexResult { +fn null_expr(dtype: DType) -> VortexResult { Ok(lit(Scalar::null(dtype.as_nullable()))) } @@ -140,10 +140,10 @@ mod tests { impl StatBinder for TestBinder { fn bind_aggregate( &self, - _input: &BoundExpression, + _input: &BoundExpressionRef, aggregate_fn: &AggregateFnRef, _stat_dtype: &DType, - ) -> VortexResult> { + ) -> VortexResult> { let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) else { return Ok(None); }; diff --git a/vortex-array/src/stats/expr.rs b/vortex-array/src/stats/expr.rs index 1e0ceef02d3..4b8ac9a26dd 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -17,7 +17,7 @@ use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::nan_count::NanCount; use crate::aggregate_fn::fns::null_count::NullCount; use crate::aggregate_fn::fns::sum::Sum; -use crate::expr::BoundExpression; +use crate::expr::BoundExpressionRef; use crate::expr::Expression; use crate::scalar_fn::ScalarFnVTableExt; pub use crate::scalar_fn::fns::stat::StatFn; @@ -31,7 +31,7 @@ pub fn stat(expr: Expression, aggregate_fn: AggregateFnRef) -> Expression { StatFn.new_expr(StatOptions::new(aggregate_fn), [expr]) } -fn bound_stat(expr: BoundExpression, aggregate_fn: AggregateFnRef) -> BoundExpression { +fn bound_stat(expr: BoundExpressionRef, aggregate_fn: AggregateFnRef) -> BoundExpressionRef { StatFn .try_new_bound_expr(StatOptions::new(aggregate_fn), [expr]) .vortex_expect("stat expressions must use an aggregate supported by the child dtype") @@ -43,7 +43,7 @@ pub fn min_max(expr: Expression) -> Expression { stat(expr, MinMax.bind(NumericalAggregateOpts::skip_nans())) } -fn bound_min_max(expr: BoundExpression) -> BoundExpression { +fn bound_min_max(expr: BoundExpressionRef) -> BoundExpressionRef { bound_stat(expr, MinMax.bind(NumericalAggregateOpts::skip_nans())) } @@ -53,7 +53,7 @@ pub fn sum(expr: Expression) -> Expression { stat(expr, Sum.bind(NumericalAggregateOpts::skip_nans())) } -fn bound_sum(expr: BoundExpression) -> BoundExpression { +fn bound_sum(expr: BoundExpressionRef) -> BoundExpressionRef { bound_stat(expr, Sum.bind(NumericalAggregateOpts::skip_nans())) } @@ -62,7 +62,7 @@ pub fn null_count(expr: Expression) -> Expression { stat(expr, NullCount.bind(EmptyOptions)) } -fn bound_null_count(expr: BoundExpression) -> BoundExpression { +fn bound_null_count(expr: BoundExpressionRef) -> BoundExpressionRef { bound_stat(expr, NullCount.bind(EmptyOptions)) } @@ -71,7 +71,7 @@ pub fn all_null(expr: Expression) -> Expression { stat(expr, AllNull.bind(EmptyOptions)) } -fn bound_all_null(expr: BoundExpression) -> BoundExpression { +fn bound_all_null(expr: BoundExpressionRef) -> BoundExpressionRef { bound_stat(expr, AllNull.bind(EmptyOptions)) } @@ -80,7 +80,7 @@ pub fn all_nan(expr: Expression) -> Expression { stat(expr, AllNan.bind(EmptyOptions)) } -fn bound_all_nan(expr: BoundExpression) -> BoundExpression { +fn bound_all_nan(expr: BoundExpressionRef) -> BoundExpressionRef { bound_stat(expr, AllNan.bind(EmptyOptions)) } @@ -89,7 +89,7 @@ pub fn all_non_null(expr: Expression) -> Expression { stat(expr, AllNonNull.bind(EmptyOptions)) } -fn bound_all_non_null(expr: BoundExpression) -> BoundExpression { +fn bound_all_non_null(expr: BoundExpressionRef) -> BoundExpressionRef { bound_stat(expr, AllNonNull.bind(EmptyOptions)) } @@ -98,7 +98,7 @@ pub fn all_non_nan(expr: Expression) -> Expression { stat(expr, AllNonNan.bind(EmptyOptions)) } -fn bound_all_non_nan(expr: BoundExpression) -> BoundExpression { +fn bound_all_non_nan(expr: BoundExpressionRef) -> BoundExpressionRef { bound_stat(expr, AllNonNan.bind(EmptyOptions)) } @@ -107,7 +107,7 @@ pub fn nan_count(expr: Expression) -> Expression { stat(expr, NanCount.bind(EmptyOptions)) } -fn bound_nan_count(expr: BoundExpression) -> BoundExpression { +fn bound_nan_count(expr: BoundExpressionRef) -> BoundExpressionRef { bound_stat(expr, NanCount.bind(EmptyOptions)) } @@ -117,56 +117,57 @@ fn bound_nan_count(expr: BoundExpression) -> BoundExpression { /// the input dtype. pub mod bound { use crate::aggregate_fn::AggregateFnRef; - use crate::expr::BoundExpression; + use crate::expr::BoundExpressionRef; /// Creates a bound expression that reads a stored aggregate statistic. - pub fn stat(expr: BoundExpression, aggregate_fn: AggregateFnRef) -> BoundExpression { + pub fn stat(expr: BoundExpressionRef, aggregate_fn: AggregateFnRef) -> BoundExpressionRef { super::bound_stat(expr, aggregate_fn) } /// Creates a bound nullable `{ min, max }` statistic expression. - pub fn min_max(expr: BoundExpression) -> BoundExpression { + pub fn min_max(expr: BoundExpressionRef) -> BoundExpressionRef { super::bound_min_max(expr) } /// Creates a bound nullable sum statistic expression. - pub fn sum(expr: BoundExpression) -> BoundExpression { + pub fn sum(expr: BoundExpressionRef) -> BoundExpressionRef { super::bound_sum(expr) } /// Creates a bound nullable null-count statistic expression. - pub fn null_count(expr: BoundExpression) -> BoundExpression { + pub fn null_count(expr: BoundExpressionRef) -> BoundExpressionRef { super::bound_null_count(expr) } /// Creates a bound nullable all-null statistic expression. - pub fn all_null(expr: BoundExpression) -> BoundExpression { + pub fn all_null(expr: BoundExpressionRef) -> BoundExpressionRef { super::bound_all_null(expr) } /// Creates a bound nullable all-NaN statistic expression. - pub fn all_nan(expr: BoundExpression) -> BoundExpression { + pub fn all_nan(expr: BoundExpressionRef) -> BoundExpressionRef { super::bound_all_nan(expr) } /// Creates a bound nullable all-non-null statistic expression. - pub fn all_non_null(expr: BoundExpression) -> BoundExpression { + pub fn all_non_null(expr: BoundExpressionRef) -> BoundExpressionRef { super::bound_all_non_null(expr) } /// Creates a bound nullable all-non-NaN statistic expression. - pub fn all_non_nan(expr: BoundExpression) -> BoundExpression { + pub fn all_non_nan(expr: BoundExpressionRef) -> BoundExpressionRef { super::bound_all_non_nan(expr) } /// Creates a bound nullable NaN-count statistic expression. - pub fn nan_count(expr: BoundExpression) -> BoundExpression { + pub fn nan_count(expr: BoundExpressionRef) -> BoundExpressionRef { super::bound_nan_count(expr) } } #[cfg(test)] mod tests { + use std::sync::Arc; use std::sync::LazyLock; use vortex_buffer::buffer; @@ -209,7 +210,7 @@ mod tests { fn bound_stats_constructor_preserves_child_and_dtype() -> VortexResult<()> { let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let root = bound_expr::root(input_dtype.clone()); - let bound = bound_stats::sum(root.clone()); + let bound = bound_stats::sum(Arc::clone(&root)); assert_eq!(bound.children(), &[root]); assert_eq!( diff --git a/vortex-array/src/stats/rewrite.rs b/vortex-array/src/stats/rewrite.rs index ddf74ee5dab..7c3d2b849fa 100644 --- a/vortex-array/src/stats/rewrite.rs +++ b/vortex-array/src/stats/rewrite.rs @@ -12,7 +12,7 @@ use vortex_session::VortexSession; use vortex_utils::iter::ReduceBalancedIterExt; use crate::dtype::DType; -use crate::expr::BoundExpression; +use crate::expr::BoundExpressionRef; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::fns::binary::Binary; @@ -56,9 +56,9 @@ pub trait StatsRewriteRule: Debug + Send + Sync + 'static { /// Returns `Ok(None)` when this rule cannot construct a sound falsity proof for `expr`. fn falsify( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { _ = expr; _ = ctx; Ok(None) @@ -76,9 +76,9 @@ pub trait StatsRewriteRule: Debug + Send + Sync + 'static { /// Returns `Ok(None)` when this rule cannot construct a sound truth proof for `expr`. fn satisfy( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { _ = expr; _ = ctx; Ok(None) @@ -102,23 +102,23 @@ impl<'a> StatsRewriteCtx<'a> { } /// Return the dtype of `expr` within this rewrite scope. - pub fn return_dtype(&self, expr: &BoundExpression) -> VortexResult { + pub fn return_dtype(&self, expr: &BoundExpressionRef) -> VortexResult { Ok(expr.dtype().clone()) } /// Rewrite `expr` into a stats-backed falsifier. - pub fn falsify(&self, expr: &BoundExpression) -> VortexResult> { + pub fn falsify(&self, expr: &BoundExpressionRef) -> VortexResult> { self.ensure_predicate(expr)?; rewrite(expr, self, StatsRewriteRule::falsify) } /// Rewrite `expr` into a stats-backed satisfier. - pub fn satisfy(&self, expr: &BoundExpression) -> VortexResult> { + pub fn satisfy(&self, expr: &BoundExpressionRef) -> VortexResult> { self.ensure_predicate(expr)?; rewrite(expr, self, StatsRewriteRule::satisfy) } - fn ensure_predicate(&self, expr: &BoundExpression) -> VortexResult<()> { + fn ensure_predicate(&self, expr: &BoundExpressionRef) -> VortexResult<()> { let dtype = self.return_dtype(expr)?; vortex_ensure!( matches!(dtype, DType::Bool(_)), @@ -129,14 +129,14 @@ impl<'a> StatsRewriteCtx<'a> { } fn rewrite( - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, apply: fn( &dyn StatsRewriteRule, - &BoundExpression, + &BoundExpressionRef, &StatsRewriteCtx<'_>, - ) -> VortexResult>, -) -> VortexResult> { + ) -> VortexResult>, +) -> VortexResult> { // The scope alone proves nothing about the rows it contains. let Some(scalar_fn) = expr.as_scalar() else { return Ok(None); @@ -160,6 +160,8 @@ fn rewrite( #[cfg(test)] mod tests { + use std::sync::Arc; + use vortex_error::VortexResult; use super::StatsRewriteCtx; @@ -167,7 +169,7 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; - use crate::expr::BoundExpression; + use crate::expr::BoundExpressionRef; use crate::expr::lit; use crate::expr::or; use crate::scalar_fn::ScalarFnId; @@ -177,8 +179,8 @@ mod tests { #[derive(Debug)] struct StaticLiteralRule { - falsifier: Option, - satisfier: Option, + falsifier: Option, + satisfier: Option, } impl StatsRewriteRule for StaticLiteralRule { @@ -188,17 +190,17 @@ mod tests { fn falsify( &self, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(self.falsifier.clone()) } fn satisfy( &self, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(self.satisfier.clone()) } } @@ -249,7 +251,7 @@ mod tests { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let expr = lit(true).bind(&dtype)?; - assert_eq!(expr.falsify(&session)?, None); + assert_eq!(Arc::clone(&expr).falsify(&session)?, None); assert_eq!(expr.satisfy(&session)?, None); Ok(()) } @@ -260,7 +262,7 @@ mod tests { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let expr = lit(7).bind(&dtype)?; - assert!(expr.falsify(&session).is_err()); + assert!(Arc::clone(&expr).falsify(&session).is_err()); assert!(expr.satisfy(&session).is_err()); Ok(()) } diff --git a/vortex-array/src/stats/rewrite/builtins.rs b/vortex-array/src/stats/rewrite/builtins.rs index 3cb5fdb06df..ae9302c0f01 100644 --- a/vortex-array/src/stats/rewrite/builtins.rs +++ b/vortex-array/src/stats/rewrite/builtins.rs @@ -13,7 +13,7 @@ use crate::aggregate_fn::fns::all_non_nan::AllNonNan; use crate::aggregate_fn::fns::all_non_null::AllNonNull; use crate::aggregate_fn::fns::all_null::AllNull; use crate::dtype::DType; -use crate::expr::BoundExpression; +use crate::expr::BoundExpressionRef; use crate::expr::bound::and; use crate::expr::bound::and_collect; use crate::expr::bound::binary; @@ -70,7 +70,7 @@ pub(crate) fn register_builtins(session: &StatsSession) { session.register_rewrite(DynamicComparisonAllNonNanStatsRewrite); } -fn row_count() -> BoundExpression { +fn row_count() -> BoundExpressionRef { RowCount .try_new_bound_expr(EmptyOptions, []) .vortex_expect("row-count expressions are always well-typed") @@ -86,9 +86,9 @@ impl StatsRewriteRule for BinaryNanCountStatsRewrite { fn falsify( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { binary_falsify::(expr, ctx) } } @@ -103,17 +103,17 @@ impl StatsRewriteRule for BinaryAllNonNanStatsRewrite { fn falsify( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { binary_falsify::(expr, ctx) } } fn binary_falsify( - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, -) -> VortexResult> { +) -> VortexResult> { let operator = expr.as_::(); let lhs = expr.child(0); let rhs = expr.child(1); @@ -195,15 +195,15 @@ impl StatsRewriteRule for BetweenStatsRewrite { fn falsify( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { let options = expr.as_::(); - let arr = expr.child(0).clone(); - let lower = expr.child(1).clone(); - let upper = expr.child(2).clone(); + let arr = Arc::clone(expr.child(0)); + let lower = Arc::clone(expr.child(1)); + let upper = Arc::clone(expr.child(2)); - let lhs = binary(options.lower_strict.to_operator(), lower, arr.clone()); + let lhs = binary(options.lower_strict.to_operator(), lower, Arc::clone(&arr)); let rhs = binary(options.upper_strict.to_operator(), arr, upper); ctx.falsify(&and(lhs, rhs)) } @@ -219,17 +219,17 @@ impl StatsRewriteRule for IsNullNullCountStatsRewrite { fn falsify( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, lit(0u64)))) } fn satisfy( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, row_count()))) } } @@ -244,9 +244,9 @@ impl StatsRewriteRule for IsNullAllNonNullStatsRewrite { fn falsify( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(Some(all_non_null(expr.child(0)))) } } @@ -261,9 +261,9 @@ impl StatsRewriteRule for IsNullAllNullStatsRewrite { fn satisfy( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(Some(all_null(expr.child(0)))) } } @@ -278,17 +278,17 @@ impl StatsRewriteRule for IsNotNullNullCountStatsRewrite { fn falsify( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, row_count()))) } fn satisfy( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, lit(0u64)))) } } @@ -303,9 +303,9 @@ impl StatsRewriteRule for IsNotNullAllNullStatsRewrite { fn falsify( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(Some(all_null(expr.child(0)))) } } @@ -320,9 +320,9 @@ impl StatsRewriteRule for IsNotNullAllNonNullStatsRewrite { fn satisfy( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(Some(all_non_null(expr.child(0)))) } } @@ -337,9 +337,9 @@ impl StatsRewriteRule for LikeStatsRewrite { fn falsify( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { let like_options = expr.as_::(); if like_options.negated || like_options.case_insensitive { return Ok(None); @@ -392,9 +392,9 @@ impl StatsRewriteRule for ListContainsNanCountStatsRewrite { fn falsify( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { list_contains_falsify::(expr, ctx) } } @@ -409,17 +409,17 @@ impl StatsRewriteRule for ListContainsAllNonNanStatsRewrite { fn falsify( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { list_contains_falsify::(expr, ctx) } } fn list_contains_falsify( - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, -) -> VortexResult> { +) -> VortexResult> { let list = expr.child(0); let needle = expr.child(1); @@ -446,8 +446,8 @@ fn list_contains_falsify( let value_predicate = and_collect(elements.iter().map(|value| { or( - lt(value_max.clone(), lit(value.clone())), - gt(value_min.clone(), lit(value.clone())), + lt(Arc::clone(&value_max), lit(value.clone())), + gt(Arc::clone(&value_min), lit(value.clone())), ) })); value_predicate @@ -466,9 +466,9 @@ impl StatsRewriteRule for DynamicComparisonNanCountStatsRewrite { fn falsify( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { dynamic_comparison_falsify::(expr, ctx) } } @@ -483,17 +483,17 @@ impl StatsRewriteRule for DynamicComparisonAllNonNanStatsRewrite { fn falsify( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { dynamic_comparison_falsify::(expr, ctx) } } fn dynamic_comparison_falsify( - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, -) -> VortexResult> { +) -> VortexResult> { let dynamic = expr.as_::(); let lhs = expr.child(0); @@ -518,36 +518,36 @@ fn dynamic_comparison_falsify( with_non_nan_guards::

(ctx, [lhs], value_predicate) } -fn min(expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>) -> Option { +fn min(expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>) -> Option { stat_expr(expr, Stat::Min, ctx) } -fn max(expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>) -> Option { +fn max(expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>) -> Option { stat_expr(expr, Stat::Max, ctx) } -fn null_count(expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>) -> Option { +fn null_count(expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>) -> Option { stat_expr(expr, Stat::NullCount, ctx) } -fn all_null(expr: &BoundExpression) -> BoundExpression { - stat_fn(expr.clone(), AllNull.bind(AggregateEmptyOptions)) +fn all_null(expr: &BoundExpressionRef) -> BoundExpressionRef { + stat_fn(Arc::clone(expr), AllNull.bind(AggregateEmptyOptions)) } -fn all_non_null(expr: &BoundExpression) -> BoundExpression { - stat_fn(expr.clone(), AllNonNull.bind(AggregateEmptyOptions)) +fn all_non_null(expr: &BoundExpressionRef) -> BoundExpressionRef { + stat_fn(Arc::clone(expr), AllNonNull.bind(AggregateEmptyOptions)) } enum NanCheck { NotNeeded, - Check(BoundExpression), + Check(BoundExpressionRef), Unavailable, } trait NonNanProof { const EMIT_UNGUARDED_REWRITES: bool; - fn check(ctx: &StatsRewriteCtx<'_>, expr: &BoundExpression) -> VortexResult; + fn check(ctx: &StatsRewriteCtx<'_>, expr: &BoundExpressionRef) -> VortexResult; } struct NanCountProof; @@ -555,7 +555,7 @@ struct NanCountProof; impl NonNanProof for NanCountProof { const EMIT_UNGUARDED_REWRITES: bool = true; - fn check(ctx: &StatsRewriteCtx<'_>, expr: &BoundExpression) -> VortexResult { + fn check(ctx: &StatsRewriteCtx<'_>, expr: &BoundExpressionRef) -> VortexResult { non_nan_check(ctx, expr, |expr| { match stat_expr(expr, Stat::NaNCount, ctx) { Some(nan_count) => NanCheck::Check(eq(nan_count, lit(0u64))), @@ -570,9 +570,12 @@ struct AllNonNanProof; impl NonNanProof for AllNonNanProof { const EMIT_UNGUARDED_REWRITES: bool = false; - fn check(ctx: &StatsRewriteCtx<'_>, expr: &BoundExpression) -> VortexResult { + fn check(ctx: &StatsRewriteCtx<'_>, expr: &BoundExpressionRef) -> VortexResult { non_nan_check(ctx, expr, |expr| { - NanCheck::Check(stat_fn(expr.clone(), AllNonNan.bind(AggregateEmptyOptions))) + NanCheck::Check(stat_fn( + Arc::clone(expr), + AllNonNan.bind(AggregateEmptyOptions), + )) }) } } @@ -582,8 +585,8 @@ impl NonNanProof for AllNonNanProof { // from float to non-float still needs a proof about the float source values. fn non_nan_check( ctx: &StatsRewriteCtx<'_>, - expr: &BoundExpression, - proof: impl FnOnce(&BoundExpression) -> NanCheck, + expr: &BoundExpressionRef, + proof: impl FnOnce(&BoundExpressionRef) -> NanCheck, ) -> VortexResult { if let Some(scalar) = expr.as_opt::() { if !scalar.dtype().is_float() { @@ -616,10 +619,10 @@ fn has_nans(dtype: &DType) -> bool { } fn stat_expr( - expr: &BoundExpression, + expr: &BoundExpressionRef, stat: Stat, ctx: &StatsRewriteCtx<'_>, -) -> Option { +) -> Option { if let Some(literal) = literal_stat(expr, stat) { return Some(literal); } @@ -643,14 +646,14 @@ fn stat_expr( aggregate_fn .return_dtype(&input_dtype) .is_some() - .then(|| stat_fn(expr.clone(), aggregate_fn)) + .then(|| stat_fn(Arc::clone(expr), aggregate_fn)) } fn with_non_nan_guards<'a, P: NonNanProof>( ctx: &StatsRewriteCtx<'_>, - exprs: impl IntoIterator, - value_predicate: BoundExpression, -) -> VortexResult> { + exprs: impl IntoIterator, + value_predicate: BoundExpressionRef, +) -> VortexResult> { let mut nan_checks = Vec::new(); for expr in exprs { match P::check(ctx, expr)? { @@ -671,7 +674,7 @@ fn with_non_nan_guards<'a, P: NonNanProof>( }) } -fn literal_stat(expr: &BoundExpression, stat: Stat) -> Option { +fn literal_stat(expr: &BoundExpressionRef, stat: Stat) -> Option { let scalar = expr.as_opt::()?; match stat { Stat::Min | Stat::Max => Some(lit(scalar.clone())), @@ -693,11 +696,11 @@ fn literal_stat(expr: &BoundExpression, stat: Stat) -> Option { } fn cast_stat( - expr: &BoundExpression, + expr: &BoundExpressionRef, dtype: &DType, stat: Stat, ctx: &StatsRewriteCtx<'_>, -) -> Option { +) -> Option { match stat { Stat::Min | Stat::Max => stat_expr(expr, stat, ctx).map(|stat| cast(stat, dtype.clone())), Stat::NaNCount | Stat::Sum | Stat::UncompressedSizeInBytes => stat_expr(expr, stat, ctx), @@ -705,7 +708,7 @@ fn cast_stat( } } -fn stat_fn(expr: BoundExpression, aggregate_fn: AggregateFnRef) -> BoundExpression { +fn stat_fn(expr: BoundExpressionRef, aggregate_fn: AggregateFnRef) -> BoundExpressionRef { stat(expr, aggregate_fn) } @@ -728,7 +731,7 @@ mod tests { use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; - use crate::expr::BoundExpression; + use crate::expr::BoundExpressionRef; use crate::expr::Expression; use crate::expr::and; use crate::expr::between; @@ -797,15 +800,15 @@ mod tests { ) } - fn falsify(expr: &Expression) -> VortexResult> { + fn falsify(expr: &Expression) -> VortexResult> { expr.bind(&test_scope())?.falsify(&SESSION) } - fn satisfy(expr: &Expression) -> VortexResult> { + fn satisfy(expr: &Expression) -> VortexResult> { expr.bind(&test_scope())?.satisfy(&SESSION) } - fn bind_expected(expr: Option) -> VortexResult> { + fn bind_expected(expr: Option) -> VortexResult> { expr.map(|expr| expr.bind(&test_scope())).transpose() } @@ -897,9 +900,9 @@ mod tests { fn falsify( &self, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { self.0.fetch_add(1, Ordering::Relaxed); Ok(None) } diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 216b3369bd1..979fe063828 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -19,7 +19,7 @@ use vortex::array::MaskFuture; use vortex::array::ProstMetadata; use vortex::array::VortexSessionExecute; use vortex::array::arrays::Constant; -use vortex::array::expr::BoundExpression; +use vortex::array::expr::BoundExpressionRef; use vortex::array::expr::stats::Precision; use vortex::array::expr::stats::Stat; use vortex::array::expr::stats::StatsProvider; @@ -268,7 +268,7 @@ impl LayoutReader for CudaFlatReader { fn pruning_evaluation( &self, _row_range: &Range, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, mask: Mask, ) -> VortexResult { Ok(MaskFuture::ready(mask)) @@ -277,7 +277,7 @@ impl LayoutReader for CudaFlatReader { fn filter_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult { let row_range = usize::try_from(row_range.start) @@ -286,7 +286,7 @@ impl LayoutReader for CudaFlatReader { .vortex_expect("Row range end must fit within CudaFlatLayout size"); let name = Arc::clone(&self.name); let array = self.array_future(); - let expr = expr.clone(); + let expr = Arc::clone(expr); let session = self.session.clone(); Ok(MaskFuture::new(mask.len(), async move { @@ -326,7 +326,7 @@ impl LayoutReader for CudaFlatReader { fn projection_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult>> { let row_range = usize::try_from(row_range.start) @@ -335,7 +335,7 @@ impl LayoutReader for CudaFlatReader { .vortex_expect("Row range end must fit within CudaFlatLayout size"); let name = Arc::clone(&self.name); let array = self.array_future(); - let expr = expr.clone(); + let expr = Arc::clone(expr); Ok(async move { tracing::debug!("CudaFlat array evaluation {} - {}", name, expr); diff --git a/vortex-file/src/pruning.rs b/vortex-file/src/pruning.rs index df327638d00..a7e01b36203 100644 --- a/vortex-file/src/pruning.rs +++ b/vortex-file/src/pruning.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -10,7 +12,7 @@ use vortex_array::arrays::NullArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; use vortex_array::dtype::StructFields; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::bound::lit; use vortex_array::expr::stats::Stat; use vortex_array::scalar::Scalar; @@ -26,13 +28,13 @@ use vortex_session::VortexSession; use crate::FileStatistics; pub(crate) fn can_prune_file_stats( - expr: &BoundExpression, + expr: &BoundExpressionRef, row_count: u64, file_stats: &FileStatistics, struct_fields: &StructFields, session: &VortexSession, ) -> VortexResult { - let Some(pruning_expr) = expr.falsify(session)? else { + let Some(pruning_expr) = Arc::clone(expr).falsify(session)? else { return Ok(false); }; @@ -68,10 +70,10 @@ struct FileStatsBinder<'a> { impl StatBinder for FileStatsBinder<'_> { fn bind_aggregate( &self, - input: &BoundExpression, + input: &BoundExpressionRef, aggregate_fn: &AggregateFnRef, _stat_dtype: &DType, - ) -> VortexResult> { + ) -> VortexResult> { let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) else { return Ok(None); }; @@ -83,7 +85,7 @@ impl StatBinder for FileStatsBinder<'_> { } impl FileStatsBinder<'_> { - fn stat_ref(&self, field_path: &FieldPath, stat: Stat) -> Option { + fn stat_ref(&self, field_path: &FieldPath, stat: Stat) -> Option { // FileStats currently only holds top-level field statistics. if field_path.parts().len() != 1 { return None; @@ -102,7 +104,7 @@ impl FileStatsBinder<'_> { } } -fn direct_field_path(expr: &BoundExpression) -> Option { +fn direct_field_path(expr: &BoundExpressionRef) -> Option { if expr.is_root() { return Some(FieldPath::root()); } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index f5c177c9cdf..3b77018a035 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -38,7 +38,7 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::PType::I32; use vortex_array::dtype::StructFields; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::Expression; use vortex_array::expr::and; use vortex_array::expr::cast; @@ -115,7 +115,7 @@ fn strict_sorted(indices: Buffer) -> StrictSortedBuffer { StrictSortedBuffer::try_new(indices).expect("test indices should be strictly increasing") } -fn bind_scan_expr(file: &VortexFile, expr: Expression) -> BoundExpression { +fn bind_scan_expr(file: &VortexFile, expr: Expression) -> BoundExpressionRef { expr.optimize_recursive(file.dtype()) .and_then(|expr| expr.bind(file.dtype())) .vortex_expect("scan expression should bind") diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index 03ad2ab88d4..dc57fc3d0e2 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -14,7 +14,7 @@ use vortex_array::MaskFuture; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::StructFields; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::ExactBoundExpr; use vortex_error::VortexResult; use vortex_layout::ArrayFuture; @@ -72,7 +72,7 @@ impl FileStatsLayoutReader { /// /// Row-count placeholders are resolved against the full file row count, /// independent of the requested row range. - fn evaluate_file_stats(&self, expr: &BoundExpression) -> VortexResult { + fn evaluate_file_stats(&self, expr: &BoundExpressionRef) -> VortexResult { can_prune_file_stats( expr, self.child.row_count(), @@ -113,10 +113,10 @@ impl LayoutReader for FileStatsLayoutReader { fn pruning_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: Mask, ) -> VortexResult { - let key = ExactBoundExpr(expr.clone()); + let key = ExactBoundExpr(Arc::clone(expr)); // Check cache first with read-only lock. if let Some(pruned) = self.prune_cache.get(&key) { @@ -140,7 +140,7 @@ impl LayoutReader for FileStatsLayoutReader { fn filter_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult { self.child.filter_evaluation(row_range, expr, mask) @@ -149,7 +149,7 @@ impl LayoutReader for FileStatsLayoutReader { fn projection_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult { self.child.projection_evaluation(row_range, expr, mask) diff --git a/vortex-layout/benches/zone_map_prune.rs b/vortex-layout/benches/zone_map_prune.rs index 87eb21fa9e4..29dc34c368f 100644 --- a/vortex-layout/benches/zone_map_prune.rs +++ b/vortex-layout/benches/zone_map_prune.rs @@ -28,7 +28,7 @@ use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::Expression; use vortex_array::expr::eq; use vortex_array::expr::gt; @@ -248,7 +248,7 @@ fn f64_dtype() -> DType { DType::Primitive(PType::F64, Nullability::Nullable) } -fn falsify(expr: Expression, column_dtype: &DType) -> BoundExpression { +fn falsify(expr: Expression, column_dtype: &DType) -> BoundExpressionRef { expr.bind(column_dtype) .unwrap() .falsify(&SESSION) @@ -256,7 +256,7 @@ fn falsify(expr: Expression, column_dtype: &DType) -> BoundExpression { .unwrap() } -fn run(bencher: Bencher, zone_map: ZoneMap, predicate: BoundExpression) { +fn run(bencher: Bencher, zone_map: ZoneMap, predicate: BoundExpressionRef) { bencher.bench(|| { divan::black_box( zone_map @@ -269,7 +269,7 @@ fn run(bencher: Bencher, zone_map: ZoneMap, predicate: BoundExpression) { /// Integer range predicate: binds to `max` only, no row count, no NaN guard. #[divan::bench(args = ZONE_COUNTS)] fn int_gt(bencher: Bencher, num_zones: usize) { - static PREDICATE: LazyLock = + static PREDICATE: LazyLock = LazyLock::new(|| falsify(gt(root(), lit(5_000i32)), &i32_dtype())); run( bencher, @@ -282,7 +282,7 @@ fn int_gt(bencher: Bencher, num_zones: usize) { /// stores `nan_count` they lower to the same expression. #[divan::bench(args = ZONE_COUNTS)] fn float_gt(bencher: Bencher, num_zones: usize) { - static PREDICATE: LazyLock = + static PREDICATE: LazyLock = LazyLock::new(|| falsify(gt(root(), lit(5_000f64)), &f64_dtype())); run( bencher, @@ -294,7 +294,7 @@ fn float_gt(bencher: Bencher, num_zones: usize) { /// Null predicate: lowers to `null_count == row_count`, exercising the row-count path. #[divan::bench(args = ZONE_COUNTS)] fn is_not_null_pred(bencher: Bencher, num_zones: usize) { - static PREDICATE: LazyLock = + static PREDICATE: LazyLock = LazyLock::new(|| falsify(is_not_null(root()), &i32_dtype())); run( bencher, @@ -306,7 +306,7 @@ fn is_not_null_pred(bencher: Bencher, num_zones: usize) { /// A 16-term `OR` chain, which is where lowering cost grows relative to evaluation cost. #[divan::bench(args = ZONE_COUNTS)] fn or_chain(bencher: Bencher, num_zones: usize) { - static PREDICATE: LazyLock = LazyLock::new(|| { + static PREDICATE: LazyLock = LazyLock::new(|| { let expr = (0..16i32) .map(|i| eq(root(), lit(i * 500))) .reduce(or) @@ -324,7 +324,7 @@ fn or_chain(bencher: Bencher, num_zones: usize) { /// constant. #[divan::bench(args = ZONE_COUNTS)] fn missing_stats(bencher: Bencher, num_zones: usize) { - static PREDICATE: LazyLock = + static PREDICATE: LazyLock = LazyLock::new(|| falsify(gt(root(), lit(5_000i32)), &i32_dtype())); run( bencher, diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index b5f4f5b438d..f72e1bfe9b5 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -20,7 +20,7 @@ use vortex_array::MaskFuture; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -295,7 +295,7 @@ impl LayoutReader for ChunkedReader { fn pruning_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: Mask, ) -> VortexResult { if row_range.is_empty() { @@ -341,7 +341,7 @@ impl LayoutReader for ChunkedReader { fn filter_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult { if row_range.is_empty() { @@ -380,7 +380,7 @@ impl LayoutReader for ChunkedReader { fn projection_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult>> { if row_range.is_empty() { diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index c3033063d9c..125ca0392e2 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -19,7 +19,7 @@ use vortex_array::arrays::SharedArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::bound::pack as bound_pack; use vortex_array::expr::direct_bound_annotations; @@ -145,13 +145,13 @@ impl DictReader { }) } - fn values_eval(&self, expr: BoundExpression) -> SharedArrayFuture { + fn values_eval(&self, expr: BoundExpressionRef) -> SharedArrayFuture { // This is unsound since we cannot be sure that all the values are referenced in the query // after applying the filter, so if the expression is fallible this might fail when it // shouldn't. // TODO(joe): fixme - let key = ExactBoundExpr(expr.clone()); + let key = ExactBoundExpr(Arc::clone(&expr)); // Check cache first with read-only lock if let Some(fut) = self.values_evals.get(&key) { @@ -186,11 +186,10 @@ const PUSHDOWN_ANNOTATION: &str = ""; /// We want to push to the array only if the expression has a negative cost, is infallible, and is /// strict. Strictness ensures dictionary null codes still force a null result after pushdown. fn split_expression_for_pushdown( - expr: &BoundExpression, -) -> VortexResult<(BoundExpression, Option)> { - let references_root = - label_bound_tree(expr, BoundExpression::is_root, |acc, &child| acc | child); - let annotations = direct_bound_annotations(expr, |expr: &BoundExpression| { + expr: &BoundExpressionRef, +) -> VortexResult<(BoundExpressionRef, Option)> { + let references_root = label_bound_tree(expr, |node| node.is_root(), |acc, &child| acc | child); + let annotations = direct_bound_annotations(expr, |expr: &BoundExpressionRef| { let Some(scalar_fn) = expr.as_scalar() else { return vec![]; }; @@ -199,7 +198,7 @@ fn split_expression_for_pushdown( && signature.is_strict() && is_negative_cost(scalar_fn.id()) && references_root - .get(&ExactBoundExpr(expr.clone())) + .get(&ExactBoundExpr(Arc::clone(expr))) .copied() .unwrap_or(true) { @@ -208,12 +207,12 @@ fn split_expression_for_pushdown( vec![] } }); - let partition = partition_bound_annotations(expr.clone(), annotations)?; + let partition = partition_bound_annotations(Arc::clone(expr), annotations)?; if partition.partitions.is_empty() { Ok((partition.root, None)) } else { debug_assert_eq!(1, partition.partitions.len()); - Ok((partition.root, Some(partition.partitions[0].clone()))) + Ok((partition.root, Some(Arc::clone(&partition.partitions[0])))) } } @@ -242,7 +241,7 @@ impl LayoutReader for DictReader { fn pruning_evaluation( &self, _row_range: &Range, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, mask: Mask, ) -> VortexResult { // NOTE: we can get the values here, convert expression to the codes domain, and push down @@ -255,11 +254,11 @@ impl LayoutReader for DictReader { fn filter_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult { // TODO(joe): fix up expr partitioning with fallibility and strictness annotations - let values_eval = self.values_eval(expr.clone()); + let values_eval = self.values_eval(Arc::clone(expr)); // We register interest on the entire codes row_range for now, there // is no straightforward shift into the codes domain we can do to the expression @@ -286,7 +285,7 @@ impl LayoutReader for DictReader { fn projection_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult>> { // TODO: fix up expr partitioning with fallibility and strictness annotations @@ -371,7 +370,7 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; - use vortex_array::expr::BoundExpression; + use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::Expression; use vortex_array::expr::bound::pack as bound_pack; use vortex_array::expr::byte_length; @@ -745,8 +744,8 @@ mod tests { fn test_apply( original: Expression, - outer: BoundExpression, - inner: BoundExpression, + outer: BoundExpressionRef, + inner: BoundExpressionRef, ) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = VarBinArray::from_iter( @@ -769,12 +768,12 @@ mod tests { fn split_bound( expr: Expression, dtype: &DType, - ) -> VortexResult<(BoundExpression, Option)> { + ) -> VortexResult<(BoundExpressionRef, Option)> { let bound = expr.bind(dtype)?; split_expression_for_pushdown(&bound) } - fn pushed_scope(inner: &BoundExpression) -> DType { + fn pushed_scope(inner: &BoundExpressionRef) -> DType { DType::Struct( StructFields::from_iter([(PUSHDOWN_ANNOTATION, inner.dtype().clone())]), Nullability::NonNullable, diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index aa7609f1659..6e3c9cda6f9 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -13,7 +13,7 @@ use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::serde::SerializedArray; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -115,7 +115,7 @@ impl LayoutReader for FlatReader { fn pruning_evaluation( &self, _row_range: &Range, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, mask: Mask, ) -> VortexResult { Ok(MaskFuture::ready(mask)) @@ -124,7 +124,7 @@ impl LayoutReader for FlatReader { fn filter_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult { let row_range = usize::try_from(row_range.start) @@ -133,7 +133,7 @@ impl LayoutReader for FlatReader { .vortex_expect("Row range end must fit within FlatLayout size"); let name = Arc::clone(&self.name); let array = self.array_future(); - let expr = expr.clone(); + let expr = Arc::clone(expr); let session = self.session.clone(); Ok(MaskFuture::new(mask.len(), async move { @@ -183,7 +183,7 @@ impl LayoutReader for FlatReader { fn projection_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult>> { let row_range = usize::try_from(row_range.start) @@ -192,7 +192,7 @@ impl LayoutReader for FlatReader { .vortex_expect("Row range end must fit within FlatLayout size"); let name = Arc::clone(&self.name); let array = self.array_future(); - let expr = expr.clone(); + let expr = Arc::clone(expr); Ok(async move { trace!("Flat array evaluation {} - {}", name, expr); diff --git a/vortex-layout/src/layouts/list/expr.rs b/vortex-layout/src/layouts/list/expr.rs index 8c10e400ad8..848259c69ff 100644 --- a/vortex-layout/src/layouts/list/expr.rs +++ b/vortex-layout/src/layouts/list/expr.rs @@ -1,9 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::bound::not; use vortex_array::scalar_fn::fns::is_not_null::IsNotNull; use vortex_array::scalar_fn::fns::is_null::IsNull; @@ -27,7 +30,7 @@ pub(super) enum ListChildrenNeeded { } /// The minimal set of list children needed to evaluate a bound expression. -pub(super) fn get_necessary_bound_list_children(expr: &BoundExpression) -> ListChildrenNeeded { +pub(super) fn get_necessary_bound_list_children(expr: &BoundExpressionRef) -> ListChildrenNeeded { if is_bound_null_root(expr) { return ListChildrenNeeded::Validity; } @@ -47,14 +50,14 @@ pub(super) fn get_necessary_bound_list_children(expr: &BoundExpression) -> ListC .unwrap_or(ListChildrenNeeded::Validity) } -fn is_bound_null_root(expr: &BoundExpression) -> bool { +fn is_bound_null_root(expr: &BoundExpressionRef) -> bool { (expr.as_scalar().is_some_and(|f| f.is::()) || expr.as_scalar().is_some_and(|f| f.is::())) && expr.children().len() == 1 && expr.children()[0].is_root() } -fn is_bound_list_length_root(expr: &BoundExpression) -> bool { +fn is_bound_list_length_root(expr: &BoundExpressionRef) -> bool { expr.as_scalar().is_some_and(|f| f.is::()) && expr.children().len() == 1 && expr.children()[0].is_root() @@ -63,15 +66,15 @@ fn is_bound_list_length_root(expr: &BoundExpression) -> bool { /// Rewrite a validity-class expression so it can be evaluated against the list's validity bool /// array (`true` == valid row): `is_not_null(root())` becomes `root()` and `is_null(root())` /// becomes `not(root())`. All other nodes are rebuilt with rewritten children. -pub(super) fn rewrite_validity_expr(expr: &BoundExpression) -> VortexResult { +pub(super) fn rewrite_validity_expr(expr: &BoundExpressionRef) -> VortexResult { let validity_dtype = DType::Bool(Nullability::NonNullable); rewrite_validity_expr_with_root(expr, &validity_dtype) } fn rewrite_validity_expr_with_root( - expr: &BoundExpression, + expr: &BoundExpressionRef, root_dtype: &DType, -) -> VortexResult { +) -> VortexResult { if expr.as_scalar().is_some_and(|f| f.is::()) && expr.children().len() == 1 && expr.children()[0].is_root() @@ -93,7 +96,7 @@ fn rewrite_validity_expr_with_root( .iter() .map(|child| rewrite_validity_expr_with_root(child, root_dtype)) .collect::>>()?; - expr.clone().with_children(children) + Arc::clone(expr).with_children(children) } /// Rewrite an offsets-class expression so it can be evaluated against an array of list lengths. @@ -101,9 +104,9 @@ fn rewrite_validity_expr_with_root( /// offsets-class expressions they can only be validity checks, and the lengths array carries the /// same validity as the original list. pub(super) fn rewrite_offsets_expr( - expr: &BoundExpression, + expr: &BoundExpressionRef, lengths_dtype: &DType, -) -> VortexResult { +) -> VortexResult { if is_bound_list_length_root(expr) || expr.is_root() { return Ok(BoundExpression::new_root(lengths_dtype.clone())); } @@ -113,5 +116,5 @@ pub(super) fn rewrite_offsets_expr( .iter() .map(|child| rewrite_offsets_expr(child, lengths_dtype)) .collect::>>()?; - expr.clone().with_children(children) + Arc::clone(expr).with_children(children) } diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 53227635de6..6b42b6b5b18 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -19,7 +19,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::validity::Validity; @@ -108,7 +108,7 @@ impl ListReader { fn project_validity( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult { let validity_reader = self.validity.clone(); @@ -151,13 +151,13 @@ impl ListReader { fn project_all( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult { let is_full_range = row_range.start == 0 && row_range.end == self.layout.row_count(); let reader = self.clone(); let row_range = row_range.clone(); - let expr = expr.clone(); + let expr = Arc::clone(expr); Ok(async move { let mask = mask.await?; if is_full_range && mask.all_true() { @@ -170,11 +170,11 @@ impl ListReader { } /// Fetch the complete `elements`, `offsets`, and `validity` children concurrently. - fn project_all_full(&self, expr: &BoundExpression) -> VortexResult { + fn project_all_full(&self, expr: &BoundExpressionRef) -> VortexResult { let row_count = self.layout.row_count(); let elements_row_count = self.elements.row_count(); let nullability = self.layout.dtype().nullability(); - let expr = expr.clone(); + let expr = Arc::clone(expr); let offsets_fut = self.fetch_raw_offsets(&(0..row_count))?; let elements_fut = self.fetch_raw_elements(&(0..elements_row_count))?; @@ -205,7 +205,7 @@ impl ListReader { fn project_all_bounded( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: Mask, ) -> VortexResult { // Crop to the smallest contiguous row range containing every selected list. @@ -219,7 +219,7 @@ impl ListReader { ..(row_range.start + u64::try_from(selected_rows.end)?); let nullability = self.layout.dtype().nullability(); - let expr = expr.clone(); + let expr = Arc::clone(expr); let reader = self.clone(); let offsets_fut = self.fetch_raw_offsets(&selected_row_range)?; @@ -257,7 +257,7 @@ impl ListReader { fn project_offsets_validity( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult { let offsets = self.fetch_raw_offsets(row_range)?; @@ -412,7 +412,7 @@ impl LayoutReader for ListReader { fn pruning_evaluation( &self, _row_range: &Range, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, mask: Mask, ) -> VortexResult { Ok(MaskFuture::ready(mask)) @@ -421,13 +421,13 @@ impl LayoutReader for ListReader { fn filter_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult { let len = mask.len(); let reader = self.clone(); let row_range = row_range.clone(); - let expr = expr.clone(); + let expr = Arc::clone(expr); let session = self.session.clone(); Ok(MaskFuture::new(len, async move { @@ -460,7 +460,7 @@ impl LayoutReader for ListReader { fn projection_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult { // Read as little as possible based on which list children the expression needs. diff --git a/vortex-layout/src/layouts/partitioned.rs b/vortex-layout/src/layouts/partitioned.rs index 94d1c3f782d..fd0b2b37b1f 100644 --- a/vortex-layout/src/layouts/partitioned.rs +++ b/vortex-layout/src/layouts/partitioned.rs @@ -13,7 +13,7 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::transform::BoundPartitionedExpr; use vortex_array::validity::Validity; use vortex_error::VortexError; @@ -27,15 +27,15 @@ pub(crate) trait BoundPartitionedExprEval

{ fn into_mask_future( self: Arc, mask: MaskFuture, - mask_fn: impl Fn(&P, &BoundExpression, MaskFuture) -> VortexResult, - array_fn: impl Fn(&P, &BoundExpression, MaskFuture) -> VortexResult, + mask_fn: impl Fn(&P, &BoundExpressionRef, MaskFuture) -> VortexResult, + array_fn: impl Fn(&P, &BoundExpressionRef, MaskFuture) -> VortexResult, session: VortexSession, ) -> VortexResult; fn into_array_future( self: Arc, mask: MaskFuture, - array_fn: impl Fn(&P, &BoundExpression, MaskFuture) -> VortexResult, + array_fn: impl Fn(&P, &BoundExpressionRef, MaskFuture) -> VortexResult, ) -> VortexResult; } @@ -43,8 +43,8 @@ impl BoundPartitionedExprEval

for BoundPartitionedE fn into_mask_future( self: Arc, mask: MaskFuture, - mask_fn: impl Fn(&P, &BoundExpression, MaskFuture) -> VortexResult, - array_fn: impl Fn(&P, &BoundExpression, MaskFuture) -> VortexResult, + mask_fn: impl Fn(&P, &BoundExpressionRef, MaskFuture) -> VortexResult, + array_fn: impl Fn(&P, &BoundExpressionRef, MaskFuture) -> VortexResult, session: VortexSession, ) -> VortexResult { // Construct evaluations for each child. @@ -105,7 +105,7 @@ impl BoundPartitionedExprEval

for BoundPartitionedE fn into_array_future( self: Arc, mask: MaskFuture, - array_fn: impl Fn(&P, &BoundExpression, MaskFuture) -> VortexResult, + array_fn: impl Fn(&P, &BoundExpressionRef, MaskFuture) -> VortexResult, ) -> VortexResult { // Construct evaluations for each child. let field_evals: Vec<_> = self diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index e7c83ec2950..394985524e7 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -25,6 +25,7 @@ use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::transform::BoundPartitionedExpr; use vortex_array::expr::transform::partition_bound; @@ -65,8 +66,8 @@ impl RowIdxLayoutReader { } } - fn partition_expr(&self, expr: &BoundExpression) -> VortexResult { - let key = ExactBoundExpr(expr.clone()); + fn partition_expr(&self, expr: &BoundExpressionRef) -> VortexResult { + let key = ExactBoundExpr(Arc::clone(expr)); // Check cache first with read-only lock. if let Some(entry) = self.partition_cache.get(&key) @@ -85,9 +86,9 @@ impl RowIdxLayoutReader { Ok(result) } - fn compute_partitioning(&self, expr: &BoundExpression) -> VortexResult { + fn compute_partitioning(&self, expr: &BoundExpressionRef) -> VortexResult { // Partition the expression into row idx and child expressions. - let mut partitioned = partition_bound(expr.clone(), |expr: &BoundExpression| { + let mut partitioned = partition_bound(Arc::clone(expr), |expr: &BoundExpressionRef| { if expr .as_scalar() .is_some_and(|scalar_fn| scalar_fn.is::()) @@ -103,8 +104,8 @@ impl RowIdxLayoutReader { // If there's only a single partition, we can directly return the expression. if partitioned.partitions.len() == 1 { return Ok(match &partitioned.partition_annotations[0] { - Partition::RowIdx => Partitioning::RowIdx(replace_row_idx(expr.clone())?), - Partition::Child => Partitioning::Child(expr.clone()), + Partition::RowIdx => Partitioning::RowIdx(replace_row_idx(Arc::clone(expr))?), + Partition::Child => Partitioning::Child(Arc::clone(expr)), }); } @@ -125,9 +126,9 @@ impl RowIdxLayoutReader { #[derive(Clone)] enum Partitioning { // An expression that only references the row index (e.g., `row_idx == 5`). - RowIdx(BoundExpression), + RowIdx(BoundExpressionRef), // An expression that does not reference the row index. - Child(BoundExpression), + Child(BoundExpressionRef), // Contains both the RowIdx and Child expressions, (e.g., `row_idx < child.some_field`). Partitioned(Arc>), } @@ -184,14 +185,14 @@ impl LayoutReader for RowIdxLayoutReader { fn pruning_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: Mask, ) -> VortexResult { Ok(match &self.partition_expr(expr)? { Partitioning::RowIdx(expr) => row_idx_mask_future( self.row_offset, row_range, - expr.clone(), + Arc::clone(expr), MaskFuture::ready(mask), self.session.clone(), ), @@ -203,7 +204,7 @@ impl LayoutReader for RowIdxLayoutReader { fn filter_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult { match &self.partition_expr(expr)? { @@ -217,7 +218,7 @@ impl LayoutReader for RowIdxLayoutReader { Partition::RowIdx => Ok(row_idx_mask_future( self.row_offset, row_range, - expr.clone(), + Arc::clone(expr), mask, self.session.clone(), )), @@ -227,7 +228,7 @@ impl LayoutReader for RowIdxLayoutReader { Partition::RowIdx => Ok(row_idx_array_future( self.row_offset, row_range, - expr.clone(), + Arc::clone(expr), mask, self.session.clone(), )), @@ -241,14 +242,14 @@ impl LayoutReader for RowIdxLayoutReader { fn projection_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult>> { match &self.partition_expr(expr)? { Partitioning::RowIdx(expr) => Ok(row_idx_array_future( self.row_offset, row_range, - expr.clone(), + Arc::clone(expr), mask, self.session.clone(), )), @@ -258,7 +259,7 @@ impl LayoutReader for RowIdxLayoutReader { Partition::RowIdx => Ok(row_idx_array_future( self.row_offset, row_range, - expr.clone(), + Arc::clone(expr), mask, self.session.clone(), )), @@ -273,7 +274,7 @@ impl LayoutReader for RowIdxLayoutReader { } } -fn replace_row_idx(expr: BoundExpression) -> VortexResult { +fn replace_row_idx(expr: BoundExpressionRef) -> VortexResult { Ok(expr .transform_down(|node| { if node @@ -312,7 +313,7 @@ fn idx_array(row_offset: u64, row_range: &Range) -> SequenceArray { fn row_idx_mask_future( row_offset: u64, row_range: &Range, - expr: BoundExpression, + expr: BoundExpressionRef, mask: MaskFuture, session: VortexSession, ) -> MaskFuture { @@ -333,7 +334,7 @@ fn row_idx_mask_future( fn row_idx_array_future( row_offset: u64, row_range: &Range, - expr: BoundExpression, + expr: BoundExpressionRef, mask: MaskFuture, session: VortexSession, ) -> ArrayFuture { diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 005a4e47fed..346584f5925 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -20,6 +20,7 @@ use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::bound::get_item; use vortex_array::expr::bound::pack; @@ -60,7 +61,7 @@ pub struct StructReader { /// A `pack` expression that holds each individual field of the root DType. This expansion /// ensures we can correctly partition expressions over the fields of the struct. - expanded_root_expr: BoundExpression, + expanded_root_expr: BoundExpressionRef, field_lookup: Option>, partitioned_expr_cache: DashMap>>, @@ -159,8 +160,8 @@ impl StructReader { } /// Utility for partitioning an expression over the fields of a struct. - fn partition_expr(&self, expr: &BoundExpression) -> VortexResult { - let key = ExactBoundExpr(expr.clone()); + fn partition_expr(&self, expr: &BoundExpressionRef) -> VortexResult { + let key = ExactBoundExpr(Arc::clone(expr)); // Look up the cell under a shared shard lock; only a miss takes the write lock, and // only for as long as it takes to insert an empty cell. @@ -185,15 +186,18 @@ impl StructReader { Ok(cell.get_or_init(|| result).clone()) } - fn compute_partitioned_expr(&self, expr: &BoundExpression) -> VortexResult { + fn compute_partitioned_expr(&self, expr: &BoundExpressionRef) -> VortexResult { // First, we expand the root scope into the fields of the struct to ensure // that partitioning works correctly. - let expr = - expand_struct_root(expr.clone(), &self.expanded_root_expr, self.struct_fields())?; + let expr = expand_struct_root( + Arc::clone(expr), + &self.expanded_root_expr, + self.struct_fields(), + )?; // Partition the expression into expressions that can be evaluated over individual fields let mut partitioned = partition_bound( - expr.clone(), + Arc::clone(&expr), make_bound_free_field_annotator( self.dtype() .as_struct_fields_opt() @@ -224,7 +228,11 @@ impl StructReader { .iter() .zip_eq(partitioned.partition_names.iter()) .map(|(expr, name)| { - step_into_struct_field(expr.clone(), name, self.field_reader(name)?.dtype().clone()) + step_into_struct_field( + Arc::clone(expr), + name, + self.field_reader(name)?.dtype().clone(), + ) }) .try_collect::<_, Vec<_>, _>()? .into_boxed_slice(); @@ -237,12 +245,12 @@ impl StructReader { fn expanded_struct_root( root_dtype: &DType, fields: &StructFields, -) -> VortexResult { +) -> VortexResult { let root = BoundExpression::new_root(root_dtype.clone()); let children = fields .names() .iter() - .map(|name| get_item(name.clone(), root.clone())) + .map(|name| get_item(name.clone(), Arc::clone(&root))) .collect::>(); Ok(pack( fields.names().iter().cloned().zip(children), @@ -251,15 +259,15 @@ fn expanded_struct_root( } fn expand_struct_root( - expr: BoundExpression, - expanded_root: &BoundExpression, + expr: BoundExpressionRef, + expanded_root: &BoundExpressionRef, fields: &StructFields, -) -> VortexResult { +) -> VortexResult { Ok(expr .transform_down(|node| { if node.is_root() { return Ok(Transformed { - value: expanded_root.clone(), + value: Arc::clone(expanded_root), changed: true, order: TraversalOrder::Skip, }); @@ -268,11 +276,7 @@ fn expand_struct_root( let Some(scalar_fn) = node.as_scalar() else { return Ok(Transformed::no(node)); }; - if !node - .children() - .first() - .is_some_and(BoundExpression::is_root) - { + if !node.children().first().is_some_and(|child| child.is_root()) { return Ok(Transformed::no(node)); } @@ -281,7 +285,7 @@ fn expand_struct_root( vortex_err!("Field {field_name} not found while expanding struct root") })?; return Ok(Transformed { - value: expanded_root.children()[idx].clone(), + value: Arc::clone(&expanded_root.children()[idx]), changed: true, order: TraversalOrder::Skip, }); @@ -295,7 +299,7 @@ fn expand_struct_root( let idx = fields.find(name).vortex_expect( "normalized selection fields must exist in the struct root", ); - expanded_root.children()[idx].clone() + Arc::clone(&expanded_root.children()[idx]) }) .collect(); return Ok(Transformed { @@ -311,10 +315,10 @@ fn expand_struct_root( } fn step_into_struct_field( - expr: BoundExpression, + expr: BoundExpressionRef, field_name: &FieldName, field_dtype: DType, -) -> VortexResult { +) -> VortexResult { Ok(expr .transform_down(|node| { let is_field_access = node @@ -336,7 +340,7 @@ fn step_into_struct_field( .into_inner()) } -fn is_pack_or_merge(expr: &BoundExpression) -> bool { +fn is_pack_or_merge(expr: &BoundExpressionRef) -> bool { expr.as_scalar() .is_some_and(|scalar_fn| scalar_fn.is::() || scalar_fn.is::()) } @@ -347,7 +351,7 @@ fn is_pack_or_merge(expr: &BoundExpression) -> bool { #[derive(Clone)] enum Partitioned { /// An expression which only operates over a single field - Single(FieldName, BoundExpression), + Single(FieldName, BoundExpressionRef), /// An expression which operates over multiple fields Multi(Arc>), } @@ -391,7 +395,7 @@ impl LayoutReader for StructReader { fn pruning_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: Mask, ) -> VortexResult { // Partition the expression into expressions that can be evaluated over individual fields @@ -417,7 +421,7 @@ impl LayoutReader for StructReader { fn filter_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult { // Partition the expression into expressions that can be evaluated over individual fields @@ -458,7 +462,7 @@ impl LayoutReader for StructReader { fn projection_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask_fut: MaskFuture, ) -> VortexResult { let validity_fut = self diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index 700aab86ce4..264fb8928f0 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -18,7 +18,7 @@ use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::dynamic::DynamicExprUpdates; @@ -38,7 +38,7 @@ use crate::layouts::zoned::zone_map::ZoneMap; type SharedZoneMap = Shared>>; pub(super) type SharedPruningResult = Shared>>>; -type PredicateCache = Arc>>; +type PredicateCache = Arc>>; pub(super) struct PruningState { zone_count: usize, @@ -78,8 +78,11 @@ impl PruningState { } } - pub(super) fn pruning_mask_future(&self, expr: BoundExpression) -> Option { - let key = ExactBoundExpr(expr.clone()); + pub(super) fn pruning_mask_future( + &self, + expr: BoundExpressionRef, + ) -> Option { + let key = ExactBoundExpr(Arc::clone(&expr)); if let Some(result) = self.pruning_result.get(&key) { return result.value().clone(); @@ -89,7 +92,7 @@ impl PruningState { .entry(key) .or_insert_with(|| { let dynamic_updates = DynamicExprUpdates::new(&expr); - match self.pruning_predicate(expr.clone()) { + match self.pruning_predicate(Arc::clone(&expr)) { None => { trace!(%expr, "no pruning predicate"); None @@ -126,13 +129,13 @@ impl PruningState { .clone() } - fn pruning_predicate(&self, expr: BoundExpression) -> Option { - let key = ExactBoundExpr(expr.clone()); + fn pruning_predicate(&self, expr: BoundExpressionRef) -> Option { + let key = ExactBoundExpr(Arc::clone(&expr)); self.pruning_predicates .entry(key) .or_default() - .get_or_init(move || match expr.falsify(&self.session) { + .get_or_init(move || match Arc::clone(&expr).falsify(&self.session) { Ok(predicate) => predicate, Err(error) => { trace!(%expr, %error, "failed to construct stats rewrite predicate"); @@ -191,7 +194,7 @@ impl PruningState { pub(super) struct PruningResult { zone_map: ZoneMap, - predicate: BoundExpression, + predicate: BoundExpressionRef, dynamic_updates: Option, latest_result: RwLock<(u64, Mask)>, session: VortexSession, diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index 52a57c1b92b..0cd6b9fa744 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -12,7 +12,7 @@ use vortex_array::ArrayRef; use vortex_array::MaskFuture; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_buffer::BitBufferMut; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -136,7 +136,7 @@ impl LayoutReader for ZonedReader { fn pruning_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: Mask, ) -> VortexResult { trace!("Stats pruning evaluation: {} - {}", &self.name, expr); @@ -144,7 +144,7 @@ impl LayoutReader for ZonedReader { .data_child()? .pruning_evaluation(row_range, expr, mask.clone())?; - let Some(pruning_mask_future) = self.pruning.pruning_mask_future(expr.clone()) else { + let Some(pruning_mask_future) = self.pruning.pruning_mask_future(Arc::clone(expr)) else { trace!("Stats pruning evaluation: not prune-able {expr}"); return Ok(data_eval); }; @@ -169,7 +169,7 @@ impl LayoutReader for ZonedReader { .try_collect()?; let name = Arc::clone(&self.name); - let expr = expr.clone(); + let expr = Arc::clone(expr); Ok(MaskFuture::new(mask.len(), async move { trace!("Invoking stats pruning evaluation {}: {}", name, expr); @@ -209,7 +209,7 @@ impl LayoutReader for ZonedReader { fn filter_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult { self.data_child()?.filter_evaluation(row_range, expr, mask) @@ -218,7 +218,7 @@ impl LayoutReader for ZonedReader { fn projection_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult>> { // TODO(ngates): there are some projection expressions that we may also be able to diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index c84c0b443dd..3ec31c7f724 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -21,7 +21,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::Expression; use vortex_array::expr::eq; use vortex_array::expr::get_item; @@ -140,12 +140,12 @@ impl ZoneMap { /// only after the predicate has been lowered to the zone-map table. pub fn prune( &self, - predicate: &BoundExpression, + predicate: &BoundExpressionRef, session: &VortexSession, ) -> VortexResult { let mut ctx = session.create_execution_ctx(); let num_zones = self.array.len(); - let predicate = self.lower_stats(predicate.clone())?; + let predicate = self.lower_stats(Arc::clone(predicate))?; let array = self.array.clone().into_array(); let applied = array.apply_bound(&predicate)?; @@ -159,7 +159,7 @@ impl ZoneMap { substituted.null_as_false().execute(&mut ctx) } - fn lower_stats(&self, predicate: BoundExpression) -> VortexResult { + fn lower_stats(&self, predicate: BoundExpressionRef) -> VortexResult { let binder = ZoneMapStatsBinder { zone_map: self }; bind_stats(predicate, &binder) } @@ -172,10 +172,10 @@ struct ZoneMapStatsBinder<'a> { impl StatBinder for ZoneMapStatsBinder<'_> { fn bind_aggregate( &self, - input: &BoundExpression, + input: &BoundExpressionRef, aggregate_fn: &AggregateFnRef, _stat_dtype: &DType, - ) -> VortexResult> { + ) -> VortexResult> { if !input.is_root() { return Ok(None); } @@ -235,7 +235,7 @@ impl StatBinder for ZoneMapStatsBinder<'_> { } impl ZoneMapStatsBinder<'_> { - fn bind_target(&self, expr: Expression) -> VortexResult { + fn bind_target(&self, expr: Expression) -> VortexResult { expr.bind(self.zone_map.array.dtype()) } } @@ -365,7 +365,7 @@ mod tests { use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::BoundExpression; + use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::Expression; use vortex_array::expr::cast; use vortex_array::expr::gt; @@ -389,7 +389,7 @@ mod tests { use crate::layouts::zoned::zone_map::ZoneMap; use crate::test::SESSION; - fn falsify(expr: &Expression, dtype: DType) -> BoundExpression { + fn falsify(expr: &Expression, dtype: DType) -> BoundExpressionRef { expr.bind(&dtype) .unwrap() .falsify(&SESSION) diff --git a/vortex-layout/src/plan/plans/concat.rs b/vortex-layout/src/plan/plans/concat.rs index 2e814156309..0b7d441581c 100644 --- a/vortex-layout/src/plan/plans/concat.rs +++ b/vortex-layout/src/plan/plans/concat.rs @@ -161,7 +161,7 @@ impl PlanParentReduceRule for ExpressionConcatRule { let chunks = child .children() .iter() - .map(|chunk| Ok(EvalPlan::try_new(expression.clone(), chunk?)?.into_plan())) + .map(|chunk| Ok(EvalPlan::try_new(Arc::clone(expression), chunk?)?.into_plan())) .collect::>>()?; Ok(Some( ConcatPlan::try_new(expression.dtype().clone(), chunks)?.into_plan(), diff --git a/vortex-layout/src/plan/plans/eval.rs b/vortex-layout/src/plan/plans/eval.rs index f1edb967e41..9b22f821a67 100644 --- a/vortex-layout/src/plan/plans/eval.rs +++ b/vortex-layout/src/plan/plans/eval.rs @@ -5,7 +5,7 @@ use std::borrow::Cow; use std::fmt; use vortex_array::EmptyMetadata; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; @@ -26,7 +26,7 @@ pub struct Eval; /// The expression evaluated by an [`Eval`]. #[derive(Clone, Debug)] pub struct EvalData { - expression: BoundExpression, + expression: BoundExpressionRef, } /// A plan that applies an expression to its child. @@ -34,7 +34,7 @@ pub type EvalPlan = Plan; impl EvalPlan { /// Creates an evaluation of `expression`, which must be bound to the child's dtype. - pub fn try_new(expression: BoundExpression, child: PlanRef) -> VortexResult { + pub fn try_new(expression: BoundExpressionRef, child: PlanRef) -> VortexResult { validate_expression_child(&expression, &child)?; // SAFETY: The expression root dtype was validated against the child dtype above. @@ -46,7 +46,7 @@ impl EvalPlan { /// # Safety /// /// Every scope root in `expression` must have the same dtype as `child`. - pub unsafe fn new_unchecked(expression: BoundExpression, child: PlanRef) -> Self { + pub unsafe fn new_unchecked(expression: BoundExpressionRef, child: PlanRef) -> Self { PlanParts { vtable: Eval, dtype: expression.dtype().clone(), @@ -58,7 +58,7 @@ impl EvalPlan { } /// Returns the expression evaluated by this plan. - pub fn expression(&self) -> &BoundExpression { + pub fn expression(&self) -> &BoundExpressionRef { &self.data().expression } @@ -115,7 +115,7 @@ impl PlanVTable for Eval { } } -fn validate_expression_child(expression: &BoundExpression, child: &PlanRef) -> VortexResult<()> { +fn validate_expression_child(expression: &BoundExpressionRef, child: &PlanRef) -> VortexResult<()> { if !expression.is_root_bound_to(child.dtype()) { vortex_bail!( "Eval expression is not bound to child dtype {}", diff --git a/vortex-layout/src/plan/plans/pack.rs b/vortex-layout/src/plan/plans/pack.rs index 6497cc63d45..369f61cec44 100644 --- a/vortex-layout/src/plan/plans/pack.rs +++ b/vortex-layout/src/plan/plans/pack.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::borrow::Cow; +use std::sync::Arc; use vortex_array::EmptyMetadata; use vortex_array::dtype::DType; @@ -10,6 +11,7 @@ use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::descendent_bound_annotations; use vortex_array::expr::make_bound_free_field_annotator; @@ -294,13 +296,15 @@ impl PlanParentReduceRule for ExpressionPackRule { let fields = child.fields(); let referenced_fields = descendent_bound_annotations(expression, make_bound_free_field_annotator(fields)) - .get(&ExactBoundExpr(expression.clone())) + .get(&ExactBoundExpr(Arc::clone(expression))) .vortex_expect("Bound expression missing free-field annotations") .clone(); let expanded_root = expanded_struct_root(child.dtype(), fields)?; - let expanded = expand_struct_root(expression.clone(), &expanded_root, fields)?; - let partitioned = - partition_bound(expanded.clone(), make_bound_free_field_annotator(fields))?; + let expanded = expand_struct_root(Arc::clone(expression), &expanded_root, fields)?; + let partitioned = partition_bound( + Arc::clone(&expanded), + make_bound_free_field_annotator(fields), + )?; if partitioned.partition_names.is_empty() { let selected_indices = fields @@ -324,7 +328,7 @@ impl PlanParentReduceRule for ExpressionPackRule { .collect::>>()?; let rewritten = child.with_pruned_fields(pruned_fields)?.into_plan(); return Ok(Some( - EvalPlan::try_new(expression.clone(), rewritten)?.into_plan(), + EvalPlan::try_new(Arc::clone(expression), rewritten)?.into_plan(), )); } @@ -361,9 +365,9 @@ impl PlanParentReduceRule for ExpressionPackRule { .get(0) .ok_or_else(|| vortex_err!("Struct expression partition pack is empty"))?; collapsed.push((name.clone(), value_name.clone())); - partition.children()[0].clone() + Arc::clone(&partition.children()[0]) } else { - partition.clone() + Arc::clone(partition) }; let lowered = step_into_struct_field(lowered, name, field.dtype().clone())?; field_expressions[field_index] = Some(lowered); @@ -408,10 +412,10 @@ impl PlanParentReduceRule for ExpressionPackRule { /// * `collapsed` - `(partition_name, value_name)` pairs whose one-field `Pack` was removed; /// each `$.partition_name.value_name` access is rewritten to `$.partition_name`. pub(super) fn rewrite_partition_root( - expression: BoundExpression, + expression: BoundExpressionRef, root_dtype: DType, collapsed: &[(FieldName, FieldName)], -) -> VortexResult { +) -> VortexResult { Ok(expression .transform_down(|node| { if let Some(value_name) = node.as_opt::() { @@ -461,17 +465,20 @@ fn field_plan(plan: &Plan, index: usize) -> VortexResult { fn expanded_struct_root( root_dtype: &DType, fields: &StructFields, -) -> VortexResult { +) -> VortexResult { let root = BoundExpression::new_root(root_dtype.clone()); let children = fields .names() .iter() - .map(|name| BoundExpression::try_new(GetItem.bind(name.clone()), [root.clone()])) + .map(|name| BoundExpression::try_new(GetItem.bind(name.clone()), [Arc::clone(&root)])) .collect::>>()?; bound_pack(fields.names().clone(), children) } -fn is_identity_expression(expression: &BoundExpression, input_dtype: &DType) -> VortexResult { +fn is_identity_expression( + expression: &BoundExpressionRef, + input_dtype: &DType, +) -> VortexResult { if expression.is_root() { return Ok(expression.dtype() == input_dtype); } @@ -485,15 +492,15 @@ fn is_identity_expression(expression: &BoundExpression, input_dtype: &DType) -> } fn expand_struct_root( - expression: BoundExpression, - expanded_root: &BoundExpression, + expression: BoundExpressionRef, + expanded_root: &BoundExpressionRef, fields: &StructFields, -) -> VortexResult { +) -> VortexResult { Ok(expression .transform_down(|node| { if node.is_root() { return Ok(Transformed { - value: expanded_root.clone(), + value: Arc::clone(expanded_root), changed: true, order: TraversalOrder::Skip, }); @@ -502,11 +509,7 @@ fn expand_struct_root( let Some(scalar_fn) = node.as_scalar() else { return Ok(Transformed::no(node)); }; - if !node - .children() - .first() - .is_some_and(BoundExpression::is_root) - { + if !node.children().first().is_some_and(|child| child.is_root()) { return Ok(Transformed::no(node)); } @@ -515,7 +518,7 @@ fn expand_struct_root( vortex_err!("Field {field_name} not found while expanding struct root") })?; return Ok(Transformed { - value: expanded_root.children()[index].clone(), + value: Arc::clone(&expanded_root.children()[index]), changed: true, order: TraversalOrder::Skip, }); @@ -529,7 +532,7 @@ fn expand_struct_root( let index = fields .find(name) .vortex_expect("normalized selection fields must exist in the root"); - expanded_root.children()[index].clone() + Arc::clone(&expanded_root.children()[index]) }) .collect(); return Ok(Transformed { @@ -545,10 +548,10 @@ fn expand_struct_root( } fn step_into_struct_field( - expression: BoundExpression, + expression: BoundExpressionRef, field_name: &FieldName, field_dtype: DType, -) -> VortexResult { +) -> VortexResult { Ok(expression .transform_down(|node| { let is_field_access = node @@ -570,7 +573,10 @@ fn step_into_struct_field( .into_inner()) } -fn bound_pack(names: FieldNames, children: Vec) -> VortexResult { +fn bound_pack( + names: FieldNames, + children: Vec, +) -> VortexResult { BoundExpression::try_new( PackFn.bind(PackOptions { names, diff --git a/vortex-layout/src/plan/plans/row_idx.rs b/vortex-layout/src/plan/plans/row_idx.rs index 893980825d3..431eb639c28 100644 --- a/vortex-layout/src/plan/plans/row_idx.rs +++ b/vortex-layout/src/plan/plans/row_idx.rs @@ -3,6 +3,7 @@ use std::fmt::Display; use std::fmt::Formatter; +use std::sync::Arc; use vortex_array::EmptyMetadata; use vortex_array::dtype::DType; @@ -11,6 +12,7 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::transform::partition_bound; use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::Transformed; @@ -94,10 +96,10 @@ impl PlanVTable for RowIdx { /// `#row_idx` bypass `child`, and mixed expressions combine independently planned branches with a /// [`PackPlan`]. The file row offset is supplied by the execution row domain. pub fn plan_row_idx_expression( - expression: BoundExpression, + expression: BoundExpressionRef, child: PlanRef, ) -> VortexResult { - let partitioned = partition_bound(expression.clone(), |node| { + let partitioned = partition_bound(Arc::clone(&expression), |node| { if node .as_scalar() .is_some_and(|scalar_fn| scalar_fn.is::()) @@ -173,18 +175,18 @@ pub fn plan_row_idx_expression( return Err(vortex_err!("Row-index expression partition is empty")); }; collapsed.push((row_idx_partition_name, value_name.clone())); - row_idx_partition.children()[0].clone() + Arc::clone(&row_idx_partition.children()[0]) } else { - row_idx_partition.clone() + Arc::clone(row_idx_partition) }; let child_expression = if child_partition.children().len() == 1 { let Some(value_name) = child_pack.names.get(0) else { return Err(vortex_err!("Data expression partition is empty")); }; collapsed.push((child_partition_name, value_name.clone())); - child_partition.children()[0].clone() + Arc::clone(&child_partition.children()[0]) } else { - child_partition.clone() + Arc::clone(child_partition) }; let row_count = child.row_count(); @@ -209,7 +211,7 @@ pub fn plan_row_idx_expression( Ok(EvalPlan::try_new(residual, partitions.into_plan())?.into_plan()) } -fn replace_row_idx(expression: BoundExpression) -> VortexResult { +fn replace_row_idx(expression: BoundExpressionRef) -> VortexResult { Ok(expression .transform_down(|node| { if node diff --git a/vortex-layout/src/plan/plans/take.rs b/vortex-layout/src/plan/plans/take.rs index af9eca55da1..f8b1027291c 100644 --- a/vortex-layout/src/plan/plans/take.rs +++ b/vortex-layout/src/plan/plans/take.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::borrow::Cow; +use std::sync::Arc; use vortex_array::EmptyMetadata; use vortex_array::dtype::DType; @@ -151,14 +152,14 @@ impl PlanParentReduceRule for ExpressionTakeRule { |acc, &child| (acc.0 | child.0, acc.1 & child.1, acc.2 | child.2), ); let (references_root, is_strict, is_fallible) = labels - .get(&ExactBoundExpr(expression.clone())) + .get(&ExactBoundExpr(Arc::clone(expression))) .copied() .unwrap_or((false, false, true)); if !references_root || !is_strict || is_fallible { return Ok(None); } - let values = EvalPlan::try_new(expression.clone(), child.values()?)?.into_plan(); + let values = EvalPlan::try_new(Arc::clone(expression), child.values()?)?.into_plan(); Ok(Some(TakePlan::new(child.codes()?, values).into_plan())) } } diff --git a/vortex-layout/src/reader.rs b/vortex-layout/src/reader.rs index 91dda0e23c5..2c84f770a3a 100644 --- a/vortex-layout/src/reader.rs +++ b/vortex-layout/src/reader.rs @@ -14,7 +14,7 @@ use vortex_array::MaskFuture; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_mask::Mask; @@ -224,7 +224,7 @@ pub trait LayoutReader: 'static + Send + Sync { fn pruning_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: Mask, ) -> VortexResult; @@ -240,7 +240,7 @@ pub trait LayoutReader: 'static + Send + Sync { fn filter_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult; @@ -256,7 +256,7 @@ pub trait LayoutReader: 'static + Send + Sync { fn projection_evaluation( &self, row_range: &Range, - expr: &BoundExpression, + expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult; } diff --git a/vortex-layout/src/scan/filter.rs b/vortex-layout/src/scan/filter.rs index 2cacfbecf89..819f5b97191 100644 --- a/vortex-layout/src/scan/filter.rs +++ b/vortex-layout/src/scan/filter.rs @@ -2,12 +2,13 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::iter; +use std::sync::Arc; use bit_vec::BitVec; use itertools::Itertools; use parking_lot::RwLock; use sketches_ddsketch::DDSketch; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::dynamic::DynamicExprUpdates; use vortex_array::scalar_fn::fns::operators::Operator; @@ -23,7 +24,7 @@ const DEFAULT_SELECTIVITY_QUANTILE: f64 = 0.1; /// conjunctions in an attempt to minimize the work done. pub struct FilterExpr { /// The conjuncts involved in the filter expression. - conjuncts: Vec, + conjuncts: Vec, /// A histogram for the selectivity of each conjunct. conjunct_selectivity: Vec>, /// Dynamic expression trackers for each conjunct, incase they contain dynamic expressions. @@ -34,7 +35,7 @@ pub struct FilterExpr { selectivity_quantile: f64, } -fn bound_conjuncts(expr: &BoundExpression) -> Vec { +fn bound_conjuncts(expr: &BoundExpressionRef) -> Vec { let mut conjuncts = Vec::new(); let mut pending = vec![expr]; @@ -46,7 +47,7 @@ fn bound_conjuncts(expr: &BoundExpression) -> Vec { { pending.extend(expr.children().iter().rev()); } else { - conjuncts.push(expr.clone()); + conjuncts.push(Arc::clone(expr)); } } @@ -54,7 +55,7 @@ fn bound_conjuncts(expr: &BoundExpression) -> Vec { } impl FilterExpr { - pub fn new(expr: BoundExpression) -> Self { + pub fn new(expr: BoundExpressionRef) -> Self { let conjuncts = bound_conjuncts(&expr); let num_conjuncts = conjuncts.len(); @@ -75,7 +76,7 @@ impl FilterExpr { /// The conjuncts that make up this filter expression. #[inline] - pub fn conjuncts(&self) -> &[BoundExpression] { + pub fn conjuncts(&self) -> &[BoundExpressionRef] { &self.conjuncts } diff --git a/vortex-layout/src/scan/layout.rs b/vortex-layout/src/scan/layout.rs index 94ce7acf277..7fecbcd6bf2 100644 --- a/vortex-layout/src/scan/layout.rs +++ b/vortex-layout/src/scan/layout.rs @@ -18,7 +18,7 @@ use vortex_array::arrays::ConstantArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; use vortex_array::dtype::Nullability; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::stats::Precision; use vortex_array::scalar::Scalar; use vortex_array::stats::StatsSet; @@ -193,8 +193,8 @@ struct LayoutReaderScan { reader: LayoutReaderRef, session: VortexSession, dtype: DType, - projection: BoundExpression, - filter: Option, + projection: BoundExpressionRef, + filter: Option, limit: Option, ordered: bool, selection: Selection, @@ -259,7 +259,7 @@ impl Stream for LayoutReaderScan { let split = Box::new(LayoutReaderSplit { reader: Arc::clone(&this.reader), session: this.session.clone(), - projection: this.projection.clone(), + projection: Arc::clone(&this.projection), filter: this.filter.clone(), limit: split_limit, ordered: this.ordered, @@ -286,8 +286,8 @@ impl Stream for LayoutReaderScan { struct LayoutReaderSplit { reader: LayoutReaderRef, session: VortexSession, - projection: BoundExpression, - filter: Option, + projection: BoundExpressionRef, + filter: Option, limit: Option, ordered: bool, row_range: Range, diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index f0b7a506e08..2bc697413a2 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -37,7 +37,7 @@ use itertools::Itertools; use tracing::Instrument; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::stats::Precision; use vortex_array::stats::StatsSet; use vortex_array::stream::ArrayStreamAdapter; @@ -324,8 +324,8 @@ impl DataSource for MultiLayoutDataSource { #[derive(Clone)] struct BoundScanRequest { - projection: BoundExpression, - filter: SharedVortexResult>, + projection: BoundExpressionRef, + filter: SharedVortexResult>, row_range: Option>, selection: Selection, partition_selection: Selection, diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 413761b8103..42cfacadc8d 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -12,7 +12,7 @@ use itertools::Either; use itertools::Itertools; use vortex_array::ArrayRef; use vortex_array::dtype::DType; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorAdapter; use vortex_array::stream::ArrayStream; @@ -38,8 +38,8 @@ use crate::scan::tasks::split_exec; pub struct RepeatedScan { session: VortexSession, layout_reader: LayoutReaderRef, - projection: BoundExpression, - filter: Option, + projection: BoundExpressionRef, + filter: Option, ordered: bool, /// Optionally read a subset of the rows in the file. row_range: Option>, @@ -92,8 +92,8 @@ impl RepeatedScan { pub fn new( session: VortexSession, layout_reader: LayoutReaderRef, - projection: BoundExpression, - filter: Option, + projection: BoundExpressionRef, + filter: Option, ordered: bool, row_range: Option>, selection: Selection, @@ -176,7 +176,7 @@ impl RepeatedScan { let ctx = Arc::new(TaskContext { filter: self.filter.clone().map(|f| Arc::new(FilterExpr::new(f))), reader: Arc::clone(&self.layout_reader), - projection: self.projection.clone(), + projection: Arc::clone(&self.projection), mapper: Arc::clone(&self.map_fn), }); diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index f7abbb21fb2..df248afaafd 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -17,6 +17,7 @@ use vortex_array::ArrayRef; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::analysis::referenced_field_paths; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorAdapter; @@ -58,8 +59,8 @@ use crate::scan::splits::attempt_split_ranges; pub struct ScanBuilder { session: VortexSession, layout_reader: LayoutReaderRef, - projection: BoundExpression, - filter: Option, + projection: BoundExpressionRef, + filter: Option, /// Whether the scan needs to return splits in the order they appear in the file. ordered: bool, /// Optionally read a subset of the rows in the file. @@ -136,19 +137,19 @@ impl ScanBuilder { impl ScanBuilder { /// Add a filter expression bound against the reader dtype. - pub fn with_filter(mut self, filter: BoundExpression) -> Self { + pub fn with_filter(mut self, filter: BoundExpressionRef) -> Self { self.filter = Some(filter); self } /// Add or clear a filter expression bound against the reader dtype. - pub fn with_some_filter(mut self, filter: Option) -> Self { + pub fn with_some_filter(mut self, filter: Option) -> Self { self.filter = filter; self } /// Set a projection expression bound against the reader dtype. - pub fn with_projection(mut self, projection: BoundExpression) -> Self { + pub fn with_projection(mut self, projection: BoundExpressionRef) -> Self { self.projection = projection; self } @@ -473,8 +474,8 @@ impl Stream for LazyScanStream { /// /// Projection and filter must be pre-simplified and bound against the scan dtype. pub fn referenced_field_masks( - projection: &BoundExpression, - filter: Option<&BoundExpression>, + projection: &BoundExpressionRef, + filter: Option<&BoundExpressionRef>, ) -> VortexResult> { let mut field_paths = referenced_field_paths(projection)?; if let Some(filter) = filter { @@ -518,7 +519,7 @@ mod test { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; - use vortex_array::expr::BoundExpression; + use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::eq; use vortex_array::expr::get_item; @@ -564,8 +565,8 @@ mod test { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let projection = eq(root(), lit(1_i32)).bind(&dtype)?; let filter = eq(root(), lit(2_i32)).bind(&dtype)?; - let expected_projection = ExactBoundExpr(projection.clone()); - let expected_filter = ExactBoundExpr(filter.clone()); + let expected_projection = ExactBoundExpr(Arc::clone(&projection)); + let expected_filter = ExactBoundExpr(Arc::clone(&filter)); let reader = Arc::new(CountingLayoutReader::new(Arc::new(AtomicUsize::new(0)))); let builder = ScanBuilder::new(SCAN_SESSION.clone(), reader) @@ -670,7 +671,7 @@ mod test { fn pruning_evaluation( &self, _row_range: &Range, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, _mask: Mask, ) -> VortexResult { unimplemented!("not needed for this test"); @@ -679,7 +680,7 @@ mod test { fn filter_evaluation( &self, _row_range: &Range, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, _mask: MaskFuture, ) -> VortexResult { unimplemented!("not needed for this test"); @@ -688,7 +689,7 @@ mod test { fn projection_evaluation( &self, _row_range: &Range, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, _mask: MaskFuture, ) -> VortexResult { Ok(Box::pin(async move { @@ -761,7 +762,7 @@ mod test { fn pruning_evaluation( &self, _row_range: &Range, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, mask: Mask, ) -> VortexResult { Ok(MaskFuture::ready(mask)) @@ -770,7 +771,7 @@ mod test { fn filter_evaluation( &self, _row_range: &Range, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, mask: MaskFuture, ) -> VortexResult { Ok(mask) @@ -779,7 +780,7 @@ mod test { fn projection_evaluation( &self, row_range: &Range, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, _mask: MaskFuture, ) -> VortexResult { let start = usize::try_from(row_range.start) @@ -914,7 +915,7 @@ mod test { fn pruning_evaluation( &self, _row_range: &Range, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, _mask: Mask, ) -> VortexResult { unimplemented!("not needed for this test"); @@ -923,7 +924,7 @@ mod test { fn filter_evaluation( &self, _row_range: &Range, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, _mask: MaskFuture, ) -> VortexResult { unimplemented!("not needed for this test"); @@ -932,7 +933,7 @@ mod test { fn projection_evaluation( &self, _row_range: &Range, - _expr: &BoundExpression, + _expr: &BoundExpressionRef, _mask: MaskFuture, ) -> VortexResult { Ok(Box::pin(async move { diff --git a/vortex-layout/src/scan/split_by.rs b/vortex-layout/src/scan/split_by.rs index 6106d4f671d..68bdce43ab3 100644 --- a/vortex-layout/src/scan/split_by.rs +++ b/vortex-layout/src/scan/split_by.rs @@ -133,7 +133,7 @@ mod test { use vortex_array::dtype::FieldPath; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::BoundExpression; + use vortex_array::expr::BoundExpressionRef; use vortex_buffer::buffer; use vortex_io::runtime::single::block_on; use vortex_mask::Mask; @@ -242,7 +242,7 @@ mod test { fn pruning_evaluation( &self, _: &Range, - _: &BoundExpression, + _: &BoundExpressionRef, _: Mask, ) -> VortexResult { unimplemented!() @@ -251,7 +251,7 @@ mod test { fn filter_evaluation( &self, _: &Range, - _: &BoundExpression, + _: &BoundExpressionRef, _: MaskFuture, ) -> VortexResult { unimplemented!() @@ -260,7 +260,7 @@ mod test { fn projection_evaluation( &self, _: &Range, - _: &BoundExpression, + _: &BoundExpressionRef, _: MaskFuture, ) -> VortexResult>> { unimplemented!() diff --git a/vortex-layout/src/scan/tasks.rs b/vortex-layout/src/scan/tasks.rs index 218efb64a0d..fa5cd2922ae 100644 --- a/vortex-layout/src/scan/tasks.rs +++ b/vortex-layout/src/scan/tasks.rs @@ -11,7 +11,7 @@ use futures::FutureExt; use futures::future::BoxFuture; use vortex_array::ArrayRef; use vortex_array::MaskFuture; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_scan::row_mask::RowMask; @@ -159,7 +159,7 @@ pub struct TaskContext { /// The layout reader. pub reader: Arc, /// The projection expression to apply to gather the scanned rows. - pub projection: BoundExpression, + pub projection: BoundExpressionRef, /// Function that maps into an A. pub mapper: Arc VortexResult + Send + Sync>, } diff --git a/vortex-spatial/src/prune/distance.rs b/vortex-spatial/src/prune/distance.rs index 9231b33875e..df51d5b2433 100644 --- a/vortex-spatial/src/prune/distance.rs +++ b/vortex-spatial/src/prune/distance.rs @@ -4,7 +4,7 @@ //! `ST_Distance(geom, const) radius` pruning. use geo::Rect as SpatialRect; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::binary::Binary; @@ -42,9 +42,9 @@ impl StatsRewriteRule for SpatialDistancePrune { fn falsify( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { // Only the ordered comparisons prune today. `== r` could prune in the future (a chunk is // provably empty when `r` lies outside its box's [min, max] distance interval), it's just // not implemented. `!= r` cannot: pruning would need every row's distance to equal `r`, @@ -94,11 +94,11 @@ impl StatsRewriteRule for SpatialDistancePrune { /// /// A distance is always `>= 0`, which decides the degenerate radii up front. fn distance_prune_proof( - geom: &BoundExpression, + geom: &BoundExpressionRef, query: SpatialRect, op: Operator, radius: f64, -) -> Option { +) -> Option { // A distance is always non-negative, so degenerate radii resolve without touching the box. match op { // `<= r` / `< r` with a negative radius (or zero, for `<`) match nothing: prune every chunk. @@ -130,7 +130,7 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::BoundExpression; + use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::gt_eq; use vortex_array::expr::lit; use vortex_array::expr::lt_eq; @@ -158,7 +158,7 @@ mod tests { operator: Operator, geom_first: bool, radius: impl Into, - ) -> VortexResult> { + ) -> VortexResult> { let session = spatial_session(); let mut ctx = session.create_execution_ctx(); diff --git a/vortex-spatial/src/prune/intersects.rs b/vortex-spatial/src/prune/intersects.rs index 74103003b2f..049c59c51b9 100644 --- a/vortex-spatial/src/prune/intersects.rs +++ b/vortex-spatial/src/prune/intersects.rs @@ -3,7 +3,7 @@ //! `ST_Intersects(geom, const)` pruning. -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::stats::rewrite::StatsRewriteCtx; @@ -34,9 +34,9 @@ impl StatsRewriteRule for SpatialIntersectsPrune { fn falsify( &self, - expr: &BoundExpression, + expr: &BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { let Some((geom, constant)) = geometry_and_constant(expr, ctx)? else { return Ok(None); }; @@ -56,7 +56,7 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::BoundExpression; + use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::scalar::Scalar; @@ -75,7 +75,7 @@ mod tests { /// Run the intersects rule against `SpatialIntersects(root, point(1.0, 0.5))`, operands swapped /// when `geom_first` is false. - fn falsify_intersects(geom_first: bool) -> VortexResult> { + fn falsify_intersects(geom_first: bool) -> VortexResult> { let session = spatial_session(); let mut ctx = session.create_execution_ctx(); diff --git a/vortex-spatial/src/prune/mod.rs b/vortex-spatial/src/prune/mod.rs index 32c2058b384..dac3b3d1161 100644 --- a/vortex-spatial/src/prune/mod.rs +++ b/vortex-spatial/src/prune/mod.rs @@ -15,6 +15,8 @@ mod intersects; #[cfg(test)] mod test_harness; +use std::sync::Arc; + pub use distance::SpatialDistancePrune; use geo::BoundingRect; use geo::Rect as SpatialRect; @@ -22,7 +24,7 @@ pub use intersects::SpatialIntersectsPrune; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions; -use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionRef; use vortex_array::expr::bound::binary; use vortex_array::expr::bound::case_when; use vortex_array::expr::bound::checked_add; @@ -52,9 +54,9 @@ use crate::extension::single_geometry; /// An asymmetric predicate (e.g. a future contains) must recover which operand is the column /// itself instead of calling this. fn geometry_and_constant<'a>( - expr: &'a BoundExpression, + expr: &'a BoundExpressionRef, ctx: &StatsRewriteCtx<'_>, -) -> VortexResult> { +) -> VortexResult> { // The predicate is symmetric, so the column (scope root) and the constant may be on either // side. let (lhs, rhs) = (expr.child(0), expr.child(1)); @@ -98,22 +100,22 @@ fn query_aabb( /// /// A chunk written without the statistic reads as null here; every proof built on top must let /// that null propagate to its root, where the zone map keeps the chunk. -fn aabb_stat(geom: &BoundExpression) -> BoundExpression { +fn aabb_stat(geom: &BoundExpressionRef) -> BoundExpressionRef { // `ext_storage` unwraps the native `geoarrow.box` stat value to its backing struct, so // proofs can `get_item` the coordinate fields. - ext_storage(stat(geom.clone(), GeometryAabb.bind(EmptyOptions))) + ext_storage(stat(Arc::clone(geom), GeometryAabb.bind(EmptyOptions))) } /// Lower bound on every row's squared distance to the query AABB: zero when the boxes overlap or /// touch, positive iff they are strictly separated. /// /// Prunes "near" predicates: `min_dist_sq > r^2` proves every row is farther than `r`. -fn min_dist_sq(aabb: &BoundExpression, query: SpatialRect) -> BoundExpression { - let field = |name: &str| get_item(name, aabb.clone()); +fn min_dist_sq(aabb: &BoundExpressionRef, query: SpatialRect) -> BoundExpressionRef { + let field = |name: &str| get_item(name, Arc::clone(aabb)); // Per axis: gap = max(0, q_lo - aabb_hi, aabb_lo - q_hi), positive only when the intervals // are separated. The nearest two points of the boxes are one axis-gap apart per axis, so the // squared distance is gap_x^2 + gap_y^2 (squared throughout to avoid a sqrt). - let gap = |q_lo: f64, q_hi: f64, lo: BoundExpression, hi: BoundExpression| { + let gap = |q_lo: f64, q_hi: f64, lo: BoundExpressionRef, hi: BoundExpressionRef| { maximum( lit(0.0), maximum( @@ -130,12 +132,12 @@ fn min_dist_sq(aabb: &BoundExpression, query: SpatialRect) -> BoundExpressi /// Upper bound on every row's squared distance to the query AABB. /// /// Prunes "far" predicates: `max_dist_sq < r^2` proves every row is within `r`. -fn max_dist_sq(aabb: &BoundExpression, query: SpatialRect) -> BoundExpression { - let field = |name: &str| get_item(name, aabb.clone()); +fn max_dist_sq(aabb: &BoundExpressionRef, query: SpatialRect) -> BoundExpressionRef { + let field = |name: &str| get_item(name, Arc::clone(aabb)); // Per axis: span = max(q_hi, aabb_hi) - min(q_lo, aabb_lo), the farthest two points of the // boxes can be apart. The nullable AABB field is the second `maximum`/`minimum` argument so // that `case_when`'s else branch carries the nullability - a missing stat propagates null. - let span = |q_lo: f64, q_hi: f64, lo: BoundExpression, hi: BoundExpression| { + let span = |q_lo: f64, q_hi: f64, lo: BoundExpressionRef, hi: BoundExpressionRef| { binary( Operator::Sub, maximum(lit(q_hi), hi), @@ -148,16 +150,16 @@ fn max_dist_sq(aabb: &BoundExpression, query: SpatialRect) -> BoundExpressi } /// `e * e`. -fn square(e: BoundExpression) -> BoundExpression { - binary(Operator::Mul, e.clone(), e) +fn square(e: BoundExpressionRef) -> BoundExpressionRef { + binary(Operator::Mul, Arc::clone(&e), e) } /// `max(a, b)`. -fn maximum(a: BoundExpression, b: BoundExpression) -> BoundExpression { - case_when(gt(a.clone(), b.clone()), a, b) +fn maximum(a: BoundExpressionRef, b: BoundExpressionRef) -> BoundExpressionRef { + case_when(gt(Arc::clone(&a), Arc::clone(&b)), a, b) } /// `min(a, b)`. -fn minimum(a: BoundExpression, b: BoundExpression) -> BoundExpression { - case_when(lt(a.clone(), b.clone()), a, b) +fn minimum(a: BoundExpressionRef, b: BoundExpressionRef) -> BoundExpressionRef { + case_when(lt(Arc::clone(&a), Arc::clone(&b)), a, b) }