diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 035ca91f404..832bd6e05c4 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/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,27 +63,87 @@ 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, + /// 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 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, @@ -89,25 +151,121 @@ 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 { - 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>, + /// 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 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) + } +} + +impl Eq for ExactBoundExpr {} + +impl Hash for ExactBoundExpr { + fn hash(&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")); diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index 028228aa275..817ca309b8e 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/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, + prune_cache: DashMap, } impl FileStatsLayoutReader { @@ -102,18 +103,22 @@ 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) { - 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::::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 0beab32f938..f436619e2c6 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -21,6 +21,8 @@ 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::WeakBoundExprKey; use vortex_array::expr::bound::pack as bound_pack; use vortex_array::expr::root; use vortex_array::optimizer::ArrayOptimizer; @@ -54,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, @@ -146,22 +148,22 @@ 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.retain(|key, _| key.is_live()); self.values_evals - .entry(key) + .entry(WeakBoundExprKey::new(expr)) .or_insert_with(|| { + let expr = expr.clone(); self.values_array_uncanonical() .map(move |array| { let array = array?.apply_bound(&expr)?; @@ -260,7 +262,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..a3bd9be3615 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -25,7 +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; @@ -50,7 +51,7 @@ pub struct RowIdxLayoutReader { name: Arc, row_offset: u64, child: Arc, - partition_cache: DashMap>>, + partition_cache: DashMap>>, session: VortexSession, } @@ -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()); @@ -77,8 +76,9 @@ impl RowIdxLayoutReader { let result = self.compute_partitioning(expr)?; + self.partition_cache.retain(|key, _| key.is_live()); self.partition_cache - .entry(key) + .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 9d9630473c8..ec01eb97d77 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; @@ -20,7 +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; @@ -96,7 +98,7 @@ pub struct StructReader { expanded_root_expr: BoundExpression, field_lookup: Option>, - partitioned_expr_cache: DashMap>>, + partitioned_expr_cache: DashMap>>, } impl StructReader { @@ -200,30 +202,31 @@ 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) - .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 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 identity-keyed 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 +543,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 +607,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 +633,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 +680,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..c6fdcccbdae 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -19,7 +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; @@ -49,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 { @@ -83,16 +84,20 @@ 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.retain(|key, _| key.is_live()); + self.pruning_predicates.retain(|key, _| key.is_live()); self.pruning_result - .entry(key) + .entry(WeakBoundExprKey::new(expr)) .or_insert_with(|| { + let expr = expr.clone(); let dynamic_updates = DynamicExprUpdates::new(&expr); match self.pruning_predicate(expr.clone()) { None => { @@ -139,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) 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); };