From 0372317a38fdd58e6fc08c8d46424982aa5a6b44 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 19 Oct 2026 16:14:34 +0100 Subject: [PATCH 1/3] Probe identity-keyed expression caches without cloning the key Every split of a scan prunes with each filter conjunct, and each layout reader on the way looks the conjunct up in a cache keyed by tree identity. Each lookup cloned the expression into a key, hashed the options of its function, such as literal lists, and dropped the key again. The clones touch reference counts that every scan thread shares, and on TPC-DS q39 on DuckDB they were more than half of the pruning time. An ExactKey trait lets these maps be probed with a borrowed expression, the hash uses only the identity of the children, and the struct reader lends its cached partitioning instead of cloning it. Pruning in q39 takes a third of the samples it took before. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5.5 --- vortex-array/src/expr/bound_expression.rs | 69 ++++++++++++++++----- vortex-file/src/v2/file_stats_reader.rs | 10 +-- vortex-layout/src/layouts/dict/reader.rs | 12 ++-- vortex-layout/src/layouts/row_idx/mod.rs | 7 +-- vortex-layout/src/layouts/struct_/reader.rs | 41 +++++++----- vortex-layout/src/layouts/zoned/pruning.rs | 10 +-- vortex-layout/src/layouts/zoned/reader.rs | 2 +- 7 files changed, 102 insertions(+), 49 deletions(-) diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 035ca91f404..9d9a360e5fc 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -1,6 +1,7 @@ // 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; @@ -61,9 +62,37 @@ 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`] 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 expression that identifies this key. + fn bound_expr(&self) -> &BoundExpression; +} + +impl ExactKey for ExactBoundExpr { + fn bound_expr(&self) -> &BoundExpression { + &self.0 + } +} + +impl ExactKey for BoundExpression { + fn bound_expr(&self) -> &BoundExpression { + self + } +} + +impl<'a> Borrow 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) { + match (self.bound_expr(), other.bound_expr()) { ( BoundExpression::Root { dtype: lhs_dtype }, BoundExpression::Root { dtype: rhs_dtype }, @@ -80,8 +109,8 @@ impl PartialEq for ExactBoundExpr { children: rhs_children, }, ) => { - lhs_fn == rhs_fn - && Arc::ptr_eq(lhs_children, rhs_children) + Arc::ptr_eq(lhs_children, rhs_children) + && lhs_fn == rhs_fn && lhs_dtype == rhs_dtype } _ => false, @@ -89,27 +118,37 @@ impl PartialEq for ExactBoundExpr { } } -impl Eq for ExactBoundExpr {} +impl Eq for dyn ExactKey + '_ {} -impl Hash for ExactBoundExpr { +impl Hash for dyn ExactKey + '_ { 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 { + // 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.bound_expr() { BoundExpression::Root { .. } => state.write_u8(0), - BoundExpression::Scalar { - scalar_fn, - children, - .. - } => { + BoundExpression::Scalar { children, .. } => { state.write_u8(1); - scalar_fn.hash(state); Arc::as_ptr(children).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(&self, state: &mut H) { + (self as &dyn ExactKey).hash(state); + } +} + impl BoundExpression { /// Create a bound root expression with the given dtype. pub fn new_root(dtype: DType) -> Self { diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index 028228aa275..57563168c6d 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -15,6 +15,7 @@ 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_error::VortexResult; use vortex_layout::ArrayFuture; use vortex_layout::LayoutReader; @@ -102,15 +103,16 @@ impl LayoutReader for FileStatsLayoutReader { expr: &BoundExpression, mask: Mask, ) -> VortexResult { - 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) { + None => match self.prune_cache.entry(ExactBoundExpr(expr.clone())) { Entry::Occupied(entry) => *entry.get(), Entry::Vacant(entry) => *entry.insert(self.evaluate_file_stats(expr)?), }, diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 0beab32f938..24fe8d3e533 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -21,6 +21,7 @@ use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; +use vortex_array::expr::ExactKey; use vortex_array::expr::bound::pack as bound_pack; use vortex_array::expr::root; use vortex_array::optimizer::ArrayOptimizer; @@ -146,22 +147,21 @@ impl DictReader { }) } - fn values_eval(&self, expr: BoundExpression) -> SharedArrayFuture { + fn values_eval(&self, expr: &BoundExpression) -> 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()); - // Check cache first with read-only lock - if let Some(fut) = self.values_evals.get(&key) { + if let Some(fut) = self.values_evals.get(expr as &dyn ExactKey) { return fut.clone(); } self.values_evals - .entry(key) + .entry(ExactBoundExpr(expr.clone())) .or_insert_with(|| { + let expr = expr.clone(); self.values_array_uncanonical() .map(move |array| { let array = array?.apply_bound(&expr)?; @@ -260,7 +260,7 @@ impl LayoutReader for DictReader { 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(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 diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index f9641cbd6cb..a11fedf8894 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -26,6 +26,7 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; +use vortex_array::expr::ExactKey; use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::Transformed; use vortex_array::expr::traversal::TraversalOrder; @@ -66,10 +67,8 @@ impl RowIdxLayoutReader { } fn partition_expr(&self, expr: &BoundExpression) -> VortexResult { - let key = ExactBoundExpr(expr.clone()); - // Check cache first with read-only lock. - if let Some(entry) = self.partition_cache.get(&key) + if let Some(entry) = self.partition_cache.get(expr as &dyn ExactKey) && let Some(partitioning) = entry.value().get() { return Ok(partitioning.clone()); @@ -78,7 +77,7 @@ impl RowIdxLayoutReader { let result = self.compute_partitioning(expr)?; self.partition_cache - .entry(key) + .entry(ExactBoundExpr(expr.clone())) .or_insert_with(|| Arc::new(OnceLock::new())) .get_or_init(|| result.clone()); diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 9d9630473c8..d375f17991f 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Deref; use std::ops::Range; use std::sync::Arc; use std::sync::OnceLock; @@ -21,6 +22,7 @@ use vortex_array::dtype::RESERVED_VALIDITY_FIELD_NAME; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; +use vortex_array::expr::ExactKey; use vortex_array::expr::bound::get_item; use vortex_array::expr::bound::is_not_null; use vortex_array::expr::bound::mask; @@ -200,16 +202,14 @@ 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: &BoundExpression) -> VortexResult { // 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. - let cell = match self.partitioned_expr_cache.get(&key) { + let cell = match self.partitioned_expr_cache.get(expr as &dyn ExactKey) { Some(entry) => Arc::clone(entry.value()), None => Arc::clone( self.partitioned_expr_cache - .entry(key) + .entry(ExactBoundExpr(expr.clone())) .or_insert_with(|| Arc::new(OnceLock::new())) .value(), ), @@ -217,13 +217,13 @@ impl StructReader { // All map guards are dropped here, so partitioning runs outside any shard lock. // Concurrent misses may compute redundantly; `get_or_init` keeps a single winner. - if let Some(value) = cell.get() { - return Ok(value.clone()); + if cell.get().is_none() { + // Cache the bound trees themselves: rebinding would create fresh child identity and + // defeat downstream `ExactBoundExpr` caches on every split. + let result = self.compute_partitioned_expr(expr)?; + cell.get_or_init(|| result); } - // Cache the bound trees themselves: rebinding would create fresh child identity and defeat - // downstream `ExactBoundExpr` caches on every split. - let result = self.compute_partitioned_expr(expr)?; - Ok(cell.get_or_init(|| result).clone()) + Ok(CachedPartitioned(cell)) } fn compute_partitioned_expr(&self, expr: &BoundExpression) -> VortexResult { @@ -540,6 +540,19 @@ enum Partitioned { Direct(Arc), } +/// A cached partitioning, borrowed from the cache rather than cloned for every split. +struct CachedPartitioned(Arc>); + +impl Deref for CachedPartitioned { + type Target = Partitioned; + + fn deref(&self) -> &Partitioned { + self.0 + .get() + .vortex_expect("partition_expr initializes the cell") + } +} + /// The fields read by a pass-through projection of a non-nullable struct. struct DirectProjection { /// The output field names. @@ -591,7 +604,7 @@ impl LayoutReader for StructReader { mask: Mask, ) -> VortexResult { // Partition the expression into expressions that can be evaluated over individual fields - match &self.partition_expr(expr)? { + match &*self.partition_expr(expr)? { Partitioned::Single(name, partition) => { let reader = self.field_reader(name)?; reader @@ -617,7 +630,7 @@ impl LayoutReader for StructReader { mask: MaskFuture, ) -> VortexResult { // Partition the expression into expressions that can be evaluated over individual fields - match &self.partition_expr(expr)? { + match &*self.partition_expr(expr)? { Partitioned::Single(name, partition) => { let reader = self.field_reader(name)?; reader @@ -664,7 +677,7 @@ impl LayoutReader for StructReader { mask_fut: MaskFuture, ) -> VortexResult { // Partition the expression into expressions that can be evaluated over individual fields - match &self.partition_expr(expr)? { + match &*self.partition_expr(expr)? { Partitioned::Single(name, partition) => { let reader = self.field_reader(name)?; reader diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index 6e39823689e..764b00d90e1 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -20,6 +20,7 @@ use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; +use vortex_array::expr::ExactKey; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::dynamic::DynamicExprUpdates; use vortex_error::SharedVortexResult; @@ -83,16 +84,15 @@ impl PruningState { } } - pub(super) fn pruning_mask_future(&self, expr: BoundExpression) -> Option { - let key = ExactBoundExpr(expr.clone()); - - if let Some(result) = self.pruning_result.get(&key) { + pub(super) fn pruning_mask_future(&self, expr: &BoundExpression) -> Option { + if let Some(result) = self.pruning_result.get(expr as &dyn ExactKey) { return result.value().clone(); } self.pruning_result - .entry(key) + .entry(ExactBoundExpr(expr.clone())) .or_insert_with(|| { + let expr = expr.clone(); let dynamic_updates = DynamicExprUpdates::new(&expr); match self.pruning_predicate(expr.clone()) { None => { diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index 46f1d229c46..61e33a5d441 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -158,7 +158,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(expr) else { trace!("Stats pruning evaluation: not prune-able {expr}"); return Ok(data_eval); }; From 1f05d3acdcd406869e9a719ea747cf945fff7ba5 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 19 Oct 2026 17:55:30 +0100 Subject: [PATCH 2/3] Format Co-Authored-By: Claude Opus 5.5 Signed-off-by: Robert Kruszewski --- vortex-layout/src/layouts/zoned/pruning.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index 764b00d90e1..4b63f9ec0cd 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -84,7 +84,10 @@ impl PruningState { } } - pub(super) fn pruning_mask_future(&self, expr: &BoundExpression) -> Option { + pub(super) fn pruning_mask_future( + &self, + expr: &BoundExpression, + ) -> Option { if let Some(result) = self.pruning_result.get(expr as &dyn ExactKey) { return result.value().clone(); } From 1251b6b9df432b76c5f196c11913fbecfbecc1b1 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 20 Oct 2026 10:59:37 +0000 Subject: [PATCH 3/3] Let caches of reused layout readers drop expressions of finished scans DuckDB scans of local files reuse the layout readers of a file across queries. The readers cache pruning results, partitioned expressions and evaluated dictionary values by the identity of the expression, and every query binds new expressions, so each query added entries that no later lookup could hit. The file statistics reader held every filter it pruned, which kept the expressions alive for all the caches below it. On ClickBench q20 the process grew by about 30 MiB per query, and the pages of each new allocation were faulted in and zeroed. Key these caches with WeakBoundExprKey, which identifies an expression as ExactBoundExpr does but holds its children weakly, and drop the entries of dead expressions when inserting a new one. The weak reference keeps the allocation of the children, so a later expression cannot take their address and match a stale key. ExactKey now compares identities so that both kinds of key can be probed with a borrowed expression. Memory of 13 runs of q20 stays at 190 MiB instead of growing to 540 MiB, and q20 and q21 take 16% and 21% less time. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015qFEJyt9gzF3mi2F91ZwML --- vortex-array/src/expr/bound_expression.rs | 190 +++++++++++++++++--- vortex-file/src/v2/file_stats_reader.rs | 61 ++++++- vortex-layout/src/layouts/dict/reader.rs | 6 +- vortex-layout/src/layouts/row_idx/mod.rs | 7 +- vortex-layout/src/layouts/struct_/reader.rs | 21 ++- vortex-layout/src/layouts/zoned/pruning.rs | 12 +- 6 files changed, 245 insertions(+), 52 deletions(-) diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 9d9a360e5fc..832bd6e05c4 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -8,6 +8,7 @@ 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; @@ -64,23 +65,56 @@ pub struct ExactBoundExpr(pub BoundExpression); /// A bound expression compared by shared tree identity, as [`ExactBoundExpr`] compares it. /// -/// Maps keyed by [`ExactBoundExpr`] 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. +/// 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 expression that identifies this key. - fn bound_expr(&self) -> &BoundExpression; + /// 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, + /// 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 bound_expr(&self) -> &BoundExpression { - &self.0 + fn identity(&self) -> ExactIdentity<'_> { + self.0.exact_identity() } } impl ExactKey for BoundExpression { - fn bound_expr(&self) -> &BoundExpression { - self + fn identity(&self) -> ExactIdentity<'_> { + self.exact_identity() } } @@ -92,24 +126,23 @@ impl<'a> Borrow for ExactBoundExpr { impl PartialEq for dyn ExactKey + '_ { fn eq(&self, other: &Self) -> bool { - match (self.bound_expr(), other.bound_expr()) { - ( - 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, }, ) => { - Arc::ptr_eq(lhs_children, rhs_children) + std::ptr::eq(lhs_children, rhs_children) && lhs_fn == rhs_fn && lhs_dtype == rhs_dtype } @@ -125,16 +158,102 @@ impl Hash for dyn ExactKey + '_ { // 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.bound_expr() { - BoundExpression::Root { .. } => state.write_u8(0), - BoundExpression::Scalar { children, .. } => { + match self.identity() { + ExactIdentity::Root(_) => state.write_u8(0), + ExactIdentity::Scalar { children, .. } => { state.write_u8(1); - Arc::as_ptr(children).hash(state); + 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>, + /// 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, + } => 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 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(&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) @@ -524,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")); diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index 57563168c6d..817ca309b8e 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -14,8 +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; @@ -42,7 +42,7 @@ pub struct FileStatsLayoutReader { child: LayoutReaderRef, file_stats: FileStatistics, session: VortexSession, - prune_cache: DashMap, + prune_cache: DashMap, } impl FileStatsLayoutReader { @@ -110,12 +110,15 @@ impl LayoutReader for FileStatsLayoutReader { .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(ExactBoundExpr(expr.clone())) { - 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 { @@ -274,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::::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 { diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 24fe8d3e533..f436619e2c6 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -22,6 +22,7 @@ use vortex_array::dtype::Nullability; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::ExactKey; +use vortex_array::expr::WeakBoundExprKey; use vortex_array::expr::bound::pack as bound_pack; use vortex_array::expr::root; use vortex_array::optimizer::ArrayOptimizer; @@ -55,7 +56,7 @@ pub struct DictReader { /// Cached dict values array values_array: OnceLock, /// Cache of expression evaluation results on the values array by expression - values_evals: DashMap, + values_evals: DashMap, values: LayoutReaderRef, codes: LayoutReaderRef, @@ -158,8 +159,9 @@ impl DictReader { return fut.clone(); } + self.values_evals.retain(|key, _| key.is_live()); self.values_evals - .entry(ExactBoundExpr(expr.clone())) + .entry(WeakBoundExprKey::new(expr)) .or_insert_with(|| { let expr = expr.clone(); self.values_array_uncanonical() diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index a11fedf8894..a3bd9be3615 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -25,8 +25,8 @@ use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::BoundExpression; -use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::ExactKey; +use vortex_array::expr::WeakBoundExprKey; use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::Transformed; use vortex_array::expr::traversal::TraversalOrder; @@ -51,7 +51,7 @@ pub struct RowIdxLayoutReader { name: Arc, row_offset: u64, child: Arc, - partition_cache: DashMap>>, + partition_cache: DashMap>>, session: VortexSession, } @@ -76,8 +76,9 @@ impl RowIdxLayoutReader { let result = self.compute_partitioning(expr)?; + self.partition_cache.retain(|key, _| key.is_live()); self.partition_cache - .entry(ExactBoundExpr(expr.clone())) + .entry(WeakBoundExprKey::new(expr)) .or_insert_with(|| Arc::new(OnceLock::new())) .get_or_init(|| result.clone()); diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index d375f17991f..ec01eb97d77 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -21,8 +21,8 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::RESERVED_VALIDITY_FIELD_NAME; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; -use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::ExactKey; +use vortex_array::expr::WeakBoundExprKey; use vortex_array::expr::bound::get_item; use vortex_array::expr::bound::is_not_null; use vortex_array::expr::bound::mask; @@ -98,7 +98,7 @@ pub struct StructReader { expanded_root_expr: BoundExpression, field_lookup: Option>, - partitioned_expr_cache: DashMap>>, + partitioned_expr_cache: DashMap>>, } impl StructReader { @@ -207,19 +207,22 @@ impl StructReader { // only for as long as it takes to insert an empty cell. let cell = match self.partitioned_expr_cache.get(expr as &dyn ExactKey) { Some(entry) => Arc::clone(entry.value()), - None => Arc::clone( - self.partitioned_expr_cache - .entry(ExactBoundExpr(expr.clone())) - .or_insert_with(|| Arc::new(OnceLock::new())) - .value(), - ), + None => { + self.partitioned_expr_cache.retain(|key, _| key.is_live()); + Arc::clone( + self.partitioned_expr_cache + .entry(WeakBoundExprKey::new(expr)) + .or_insert_with(|| Arc::new(OnceLock::new())) + .value(), + ) + } }; // All map guards are dropped here, so partitioning runs outside any shard lock. // Concurrent misses may compute redundantly; `get_or_init` keeps a single winner. if cell.get().is_none() { // Cache the bound trees themselves: rebinding would create fresh child identity and - // defeat downstream `ExactBoundExpr` caches on every split. + // defeat downstream identity-keyed caches on every split. let result = self.compute_partitioned_expr(expr)?; cell.get_or_init(|| result); } diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index 4b63f9ec0cd..c6fdcccbdae 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -19,8 +19,8 @@ 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::ExactBoundExpr; use vortex_array::expr::ExactKey; +use vortex_array::expr::WeakBoundExprKey; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::dynamic::DynamicExprUpdates; use vortex_error::SharedVortexResult; @@ -50,9 +50,9 @@ pub(super) struct PruningState { aggregate_fns: Arc<[AggregateFnRef]>, lazy_children: Arc, session: VortexSession, - pruning_result: LazyLock>>, + pruning_result: LazyLock>>, zone_map: OnceLock, - pruning_predicates: LazyLock>>, + pruning_predicates: LazyLock>>, } impl PruningState { @@ -92,8 +92,10 @@ impl PruningState { return result.value().clone(); } + self.pruning_result.retain(|key, _| key.is_live()); + self.pruning_predicates.retain(|key, _| key.is_live()); self.pruning_result - .entry(ExactBoundExpr(expr.clone())) + .entry(WeakBoundExprKey::new(expr)) .or_insert_with(|| { let expr = expr.clone(); let dynamic_updates = DynamicExprUpdates::new(&expr); @@ -142,7 +144,7 @@ impl PruningState { } fn pruning_predicate(&self, expr: BoundExpression) -> Option { - let key = ExactBoundExpr(expr.clone()); + let key = WeakBoundExprKey::new(&expr); self.pruning_predicates .entry(key)