Probe identity-keyed expression caches without cloning the key

vortex-array/src/expr/bound_expression.rs

@@ -1,12 +1,14 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::borrow::Borrow;
 use std::fmt;
 use std::fmt::Display;
 use std::fmt::Formatter;
 use std::hash::Hash;
 use std::hash::Hasher;
 use std::sync::Arc;
+use std::sync::Weak;
 
 use itertools::Itertools;
 use vortex_error::VortexExpect;
@@ -61,53 +63,209 @@ pub enum BoundExpression {
 #[derive(Clone, Debug)]
 pub struct ExactBoundExpr(pub BoundExpression);
 
-impl PartialEq for ExactBoundExpr {
+/// A bound expression compared by shared tree identity, as [`ExactBoundExpr`] compares it.
+///
+/// Maps keyed by [`ExactBoundExpr`] or [`WeakBoundExprKey`] can be probed with
+/// `expr as &dyn ExactKey` for a borrowed [`BoundExpression`], which spares cloning the
+/// expression for every lookup. Clones touch reference counts that every scan thread shares.
+pub trait ExactKey {
+    /// The identity of the expression that this key stands for.
+    fn identity(&self) -> ExactIdentity<'_>;
+}
+
+/// What identifies a bound expression to an [`ExactKey`]: the dtype of a root, or the function,
+/// dtype and shared children of a scalar node.
+pub enum ExactIdentity<'a> {
+    /// A root of the given dtype.
+    Root(&'a DType),
+    /// A scalar node, identified by the address of its shared children.
+    Scalar {
+        /// The address of the shared children of the node.
+        children: *const Vec<BoundExpression>,
+        /// The scalar function of the node.
+        scalar_fn: &'a ScalarFnRef,
+        /// The dtype of the node.
+        dtype: &'a DType,
+    },
+}
+
+impl BoundExpression {
+    fn exact_identity(&self) -> ExactIdentity<'_> {
+        match self {
+            BoundExpression::Root { dtype } => ExactIdentity::Root(dtype),
+            BoundExpression::Scalar {
+                dtype,
+                scalar_fn,
+                children,
+            } => ExactIdentity::Scalar {
+                children: Arc::as_ptr(children),
+                scalar_fn,
+                dtype,
+            },
+        }
+    }
+}
+
+impl ExactKey for ExactBoundExpr {
+    fn identity(&self) -> ExactIdentity<'_> {
+        self.0.exact_identity()
+    }
+}
+
+impl ExactKey for BoundExpression {
+    fn identity(&self) -> ExactIdentity<'_> {
+        self.exact_identity()
+    }
+}
+
+impl<'a> Borrow<dyn ExactKey + 'a> for ExactBoundExpr {
+    fn borrow(&self) -> &(dyn ExactKey + 'a) {
+        self
+    }
+}
+
+impl PartialEq for dyn ExactKey + '_ {
     fn eq(&self, other: &Self) -> bool {
-        match (&self.0, &other.0) {
-            (
-                BoundExpression::Root { dtype: lhs_dtype },
-                BoundExpression::Root { dtype: rhs_dtype },
-            ) => lhs_dtype == rhs_dtype,
+        match (self.identity(), other.identity()) {
+            (ExactIdentity::Root(lhs_dtype), ExactIdentity::Root(rhs_dtype)) => {
+                lhs_dtype == rhs_dtype
+            }
             (
-                BoundExpression::Scalar {
-                    dtype: lhs_dtype,
-                    scalar_fn: lhs_fn,
+                ExactIdentity::Scalar {
                     children: lhs_children,
+                    scalar_fn: lhs_fn,
+                    dtype: lhs_dtype,
                 },
-                BoundExpression::Scalar {
-                    dtype: rhs_dtype,
-                    scalar_fn: rhs_fn,
+                ExactIdentity::Scalar {
                     children: rhs_children,
+                    scalar_fn: rhs_fn,
+                    dtype: rhs_dtype,
                 },
             ) => {
-                lhs_fn == rhs_fn
-                    && Arc::ptr_eq(lhs_children, rhs_children)
+                std::ptr::eq(lhs_children, rhs_children)
+                    && lhs_fn == rhs_fn
                     && lhs_dtype == rhs_dtype
             }
             _ => false,
         }
     }
 }
 
-impl Eq for ExactBoundExpr {}
+impl Eq for dyn ExactKey + '_ {}
 
-impl Hash for ExactBoundExpr {
+impl Hash for dyn ExactKey + '_ {
     fn hash<H: Hasher>(&self, state: &mut H) {
-        // DType differences are resolved by equality. Omitting the potentially lazy dtype keeps
-        // identity-keyed cache lookups from deserializing an entire schema just to compute a hash.
-        match &self.0 {
-            BoundExpression::Root { .. } => state.write_u8(0),
+        // Equal keys share their children, so the identity of the children is enough to hash.
+        // Omitting the function's options and the potentially lazy dtype keeps identity-keyed
+        // cache lookups from hashing literal lists or deserializing an entire schema.
+        match self.identity() {
+            ExactIdentity::Root(_) => state.write_u8(0),
+            ExactIdentity::Scalar { children, .. } => {
+                state.write_u8(1);
+                children.hash(state);
+            }
+        }
+    }
+}
+
+/// A key that identifies a bound expression as [`ExactBoundExpr`] does, without keeping the
+/// expression alive.
+///
+/// Caches that outlive the scans they serve, such as those of layout readers reused across
+/// queries, key their entries by it. A key can only be looked up while something else holds its
+/// expression, so [`Self::is_live`] tells which entries can still be hit, and the caches drop the
+/// others. The weak reference keeps the allocation of the children, so no later expression can
+/// reuse their address and match a stale key.
+#[derive(Clone, Debug)]
+pub enum WeakBoundExprKey {
+    /// A root of the given dtype.
+    Root(DType),
+    /// A scalar node, identified by its shared children.
+    Scalar {
+        /// The shared children of the node.
+        children: Weak<Vec<BoundExpression>>,
+        /// The scalar function of the node.
+        scalar_fn: ScalarFnRef,
+        /// The dtype of the node.
+        dtype: DType,
+    },
+}
+
+impl WeakBoundExprKey {
+    /// The key of `expr`.
+    pub fn new(expr: &BoundExpression) -> Self {
+        match expr {
+            BoundExpression::Root { dtype } => Self::Root(dtype.clone()),
             BoundExpression::Scalar {
+                dtype,
                 scalar_fn,
                 children,
-                ..
-            } => {
-                state.write_u8(1);
-                scalar_fn.hash(state);
-                Arc::as_ptr(children).hash(state);
-            }
+            } => Self::Scalar {
+                children: Arc::downgrade(children),
+                scalar_fn: scalar_fn.clone(),
+                dtype: dtype.clone(),
+            },
         }
     }
+
+    /// Whether the expression of this key is still alive, so that a lookup can still hit it.
+    pub fn is_live(&self) -> bool {
+        match self {
+            Self::Root(_) => true,
+            Self::Scalar { children, .. } => children.strong_count() > 0,
+        }
+    }
+}
+
+impl ExactKey for WeakBoundExprKey {
+    fn identity(&self) -> ExactIdentity<'_> {
+        match self {
+            Self::Root(dtype) => ExactIdentity::Root(dtype),
+            Self::Scalar {
+                children,
+                scalar_fn,
+                dtype,
+            } => ExactIdentity::Scalar {
+                children: children.as_ptr(),
+                scalar_fn,
+                dtype,
+            },
+        }
+    }
+}
+
+impl<'a> Borrow<dyn ExactKey + 'a> for WeakBoundExprKey {
+    fn borrow(&self) -> &(dyn ExactKey + 'a) {
+        self
+    }
+}
+
+impl PartialEq for WeakBoundExprKey {
+    fn eq(&self, other: &Self) -> bool {
+        (self as &dyn ExactKey) == (other as &dyn ExactKey)
+    }
+}
+
+impl Eq for WeakBoundExprKey {}
+
+impl Hash for WeakBoundExprKey {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        (self as &dyn ExactKey).hash(state);
+    }
+}
+
+impl PartialEq for ExactBoundExpr {
+    fn eq(&self, other: &Self) -> bool {
+        (self as &dyn ExactKey) == (other as &dyn ExactKey)
+    }
+}
+
+impl Eq for ExactBoundExpr {}
+
+impl Hash for ExactBoundExpr {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        (self as &dyn ExactKey).hash(state);
+    }
 }
 
 impl BoundExpression {
@@ -485,6 +643,27 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn weak_keys_match_only_their_live_expression() -> VortexResult<()> {
+        let expr = eq(col("a"), lit(1_i32));
+        let bound = expr.bind_scope(&scope())?;
+        let independently_bound = expr.bind_scope(&scope())?;
+        let key = WeakBoundExprKey::new(&bound);
+
+        assert!(key.is_live());
+        assert!((&key as &dyn ExactKey) == (&bound as &dyn ExactKey));
+        assert!((&key as &dyn ExactKey) != (&independently_bound as &dyn ExactKey));
+        // A clone shares the children, so it is the same expression to the key.
+        let clone = bound.clone();
+        assert_eq!(key, WeakBoundExprKey::new(&clone));
+
+        drop(bound);
+        assert!(key.is_live());
+        drop(clone);
+        assert!(!key.is_live());
+        Ok(())
+    }
+
     #[test]
     fn binding_reports_a_type_error() {
         let expr = eq(col("a"), lit("nope"));

vortex-file/src/v2/file_stats_reader.rs

@@ -14,7 +14,8 @@ use vortex_array::MaskFuture;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::FieldMask;
 use vortex_array::expr::BoundExpression;
-use vortex_array::expr::ExactBoundExpr;
+use vortex_array::expr::ExactKey;
+use vortex_array::expr::WeakBoundExprKey;
 use vortex_error::VortexResult;
 use vortex_layout::ArrayFuture;
 use vortex_layout::LayoutReader;
@@ -41,7 +42,7 @@ pub struct FileStatsLayoutReader {
     child: LayoutReaderRef,
     file_stats: FileStatistics,
     session: VortexSession,
-    prune_cache: DashMap<ExactBoundExpr, bool>,
+    prune_cache: DashMap<WeakBoundExprKey, bool>,
 }
 
 impl FileStatsLayoutReader {
@@ -102,18 +103,22 @@ impl LayoutReader for FileStatsLayoutReader {
         expr: &BoundExpression,
         mask: Mask,
     ) -> VortexResult<MaskFuture> {
-        let key = ExactBoundExpr(expr.clone());
-
         // Check cache first with read-only lock.
-        let cached = self.prune_cache.get(&key).map(|pruned| *pruned);
+        let cached = self
+            .prune_cache
+            .get(expr as &dyn ExactKey)
+            .map(|pruned| *pruned);
         let pruned = match cached {
             Some(pruned) => pruned,
-            // Evaluate under the entry lock, so that splits which start together evaluate the
-            // expression once rather than each falsifying it again.
-            None => match self.prune_cache.entry(key) {
-                Entry::Occupied(entry) => *entry.get(),
-                Entry::Vacant(entry) => *entry.insert(self.evaluate_file_stats(expr)?),
-            },
+            None => {
+                self.prune_cache.retain(|key, _| key.is_live());
+                // Evaluate under the entry lock, so that splits which start together evaluate the
+                // expression once rather than each falsifying it again.
+                match self.prune_cache.entry(WeakBoundExprKey::new(expr)) {
+                    Entry::Occupied(entry) => *entry.get(),
+                    Entry::Vacant(entry) => *entry.insert(self.evaluate_file_stats(expr)?),
+                }
+            }
         };
 
         if pruned {
@@ -272,6 +277,48 @@ mod tests {
         })
     }
 
+    #[test]
+    fn prune_cache_drops_expressions_no_scan_holds() -> VortexResult<()> {
+        block_on(|handle| async {
+            let session = SESSION.clone().with_handle(handle);
+            let ctx = ArrayContext::empty();
+            let segments = Arc::new(TestSegments::default());
+            let (ptr, eof) = SequenceId::root().split();
+            let values = PrimitiveArray::from_iter([1i32, 5]);
+            let struct_array = StructArray::from_fields([("col", values.into_array())].as_slice())?;
+            let strategy = TableStrategy::new(
+                Arc::new(FlatLayoutStrategy::default()),
+                Arc::new(FlatLayoutStrategy::default()),
+            );
+            let layout = strategy
+                .write_stream(
+                    ctx.into(),
+                    Arc::<TestSegments>::clone(&segments),
+                    struct_array.into_array().to_array_stream().sequenced(ptr),
+                    eof,
+                    &session,
+                )
+                .await?;
+
+            let child = layout.new_reader("".into(), segments, &SESSION, &Default::default())?;
+            let reader = FileStatsLayoutReader::new(child, test_file_stats(1, 5), SESSION.clone());
+
+            // A reader that outlives its scans sees a new expression from every scan.
+            for _ in 0..3 {
+                let expr = gt(get_item("col", root()), lit(10i32)).bind(reader.dtype())?;
+                for _ in 0..2 {
+                    let result = reader
+                        .pruning_evaluation(&(0..2), &expr, Mask::new_true(2))?
+                        .await?;
+                    assert_eq!(result, Mask::new_false(2));
+                }
+            }
+            assert_eq!(reader.prune_cache.len(), 1);
+
+            Ok(())
+        })
+    }
+
     #[test]
     fn no_pruning_when_filter_in_range() -> VortexResult<()> {
         block_on(|handle| async {