Make BoundExpression reference-counted and identity-keyed

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<A>(
-    expr: &BoundExpression,
+    expr: &BoundExpressionRef,
     annotate: A,
 ) -> BoundAnnotations<A::Annotation>
 where
-    A: AnnotationFn<BoundExpression>,
+    A: AnnotationFn<BoundExpressionRef>,
 {
     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<A>(
-    expr: &BoundExpression,
+    expr: &BoundExpressionRef,
     annotate: A,
 ) -> BoundAnnotations<A::Annotation>
 where
-    A: AnnotationFn<BoundExpression>,
+    A: AnnotationFn<BoundExpressionRef>,
 {
     bound_annotations(expr, annotate, false)
 }
 
 fn bound_annotations<A>(
-    expr: &BoundExpression,
+    expr: &BoundExpressionRef,
     annotate: A,
     propagate_up: bool,
 ) -> BoundAnnotations<A::Annotation>
 where
-    A: AnnotationFn<BoundExpression>,
+    A: AnnotationFn<BoundExpressionRef>,
 {
     let mut visitor = BoundAnnotationVisitor {
         annotations: Default::default(),
@@ -176,7 +177,7 @@ where
 
 struct BoundAnnotationVisitor<A>
 where
-    A: AnnotationFn<BoundExpression>,
+    A: AnnotationFn<BoundExpressionRef>,
 {
     annotations: BoundAnnotations<A::Annotation>,
     annotate: A,
@@ -185,9 +186,9 @@ where
 
 impl<'a, A> NodeVisitor<'a> for BoundAnnotationVisitor<A>
 where
-    A: AnnotationFn<BoundExpression>,
+    A: AnnotationFn<BoundExpressionRef>,
 {
-    type NodeTy = BoundExpression;
+    type NodeTy = BoundExpressionRef;
 
     fn visit_down(&mut self, node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {
         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::<Vec<_>>();
         let annotations = self
             .annotations
-            .entry(ExactBoundExpr(node.clone()))
+            .entry(ExactBoundExpr(Arc::clone(node)))
             .or_default();
         child_annotations
             .into_iter()

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<BoundExpression, Annotation = FieldName> {
-    move |expr: &BoundExpression| {
+) -> impl AnnotationFn<BoundExpressionRef, Annotation = FieldName> {
+    move |expr: &BoundExpressionRef| {
         let Some(scalar_fn) = expr.as_scalar() else {
             return scope.names().iter().cloned().collect();
         };

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<L: Clone>(
-    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<L> {
     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<L>,
@@ -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<TraversalOrder> {
         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)
     }
 }

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<FieldPathSet> {
+pub fn referenced_field_paths(expr: &BoundExpressionRef) -> VortexResult<FieldPathSet> {
     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<FieldPath>;
 
     fn visit_down(
         &mut self,
         requested: &Self::Context,
-        node: &BoundExpression,
+        node: &BoundExpressionRef,
     ) -> VortexResult<FoldDownContext<Self::Context, ()>> {
         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<FoldUp<()>> {