From a8f4b028b36ec95b93eefd91629acb6ab3eeb9b1 Mon Sep 17 00:00:00 2001 From: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com> Date: Fri, 18 Sep 2026 03:36:22 +0900 Subject: [PATCH 1/2] test: cover retained reader cache bounds and query-state release Signed-off-by: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com> --- vortex-layout/src/layouts/struct_/reader.rs | 45 +++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index ab03a4624fa..a667d86fd8a 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -410,6 +410,7 @@ mod tests { use vortex_array::dtype::StructFields; use vortex_array::expr::Expression; use vortex_array::expr::col; + use vortex_array::expr::dynamic; use vortex_array::expr::eq; use vortex_array::expr::get_item; use vortex_array::expr::gt; @@ -419,12 +420,16 @@ mod tests { use vortex_array::expr::root; use vortex_array::expr::select; use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::fns::operators::CompareOperator; use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_error::VortexResult; + use vortex_error::vortex_err; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; use vortex_mask::Mask; + use super::StructReader; use crate::LayoutRef; use crate::LayoutStrategy; use crate::layouts::flat::writer::FlatLayoutStrategy; @@ -435,6 +440,46 @@ mod tests { use crate::sequence::SequentialArrayStreamExt; use crate::test::SESSION; + #[rstest] + fn repeated_literal_partitions_are_bounded( + #[from(struct_layout)] (segments, layout): (Arc, LayoutRef), + ) -> VortexResult<()> { + let reader = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; + let reader = reader + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("expected struct reader"))?; + for value in 0..1_000i32 { + reader.partition_expr(eq(col("a"), lit(value)))?; + } + assert!(reader.partitioned_expr_cache.len() <= 32); + Ok(()) + } + + #[rstest] + fn dynamic_partitions_do_not_retain_query_state( + #[from(struct_layout)] (segments, layout): (Arc, LayoutRef), + ) -> VortexResult<()> { + let reader = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; + let reader = reader + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("expected struct reader"))?; + let state = Arc::new(3i32); + let value = Arc::clone(&state); + let expr = dynamic( + CompareOperator::Gt, + move || Some((*value).into()), + DType::Primitive(PType::I32, Nullability::NonNullable), + true, + col("a"), + ); + drop(reader.partition_expr(expr)?); + assert_eq!(Arc::strong_count(&state), 1); + assert_eq!(reader.partitioned_expr_cache.len(), 0); + Ok(()) + } + #[fixture] fn empty_struct() -> (Arc, LayoutRef) { let ctx = ArrayContext::empty(); From 92b0a74f7ead58467eaf4bc9a942ecce29bbf341 Mon Sep 17 00:00:00 2001 From: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com> Date: Fri, 18 Sep 2026 03:37:37 +0900 Subject: [PATCH 2/2] fix: bound reader expression caches and isolate dynamic expressions Signed-off-by: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com> --- vortex-file/src/tests.rs | 113 +++++++++++++++++ vortex-file/src/v2/file_stats_reader.rs | 18 +-- vortex-layout/src/expr_cache.rs | 129 ++++++++++++++++++++ vortex-layout/src/layouts/dict/reader.rs | 31 ++--- vortex-layout/src/layouts/row_idx/mod.rs | 25 +--- vortex-layout/src/layouts/struct_/reader.rs | 20 +-- vortex-layout/src/layouts/zoned/pruning.rs | 31 ++--- vortex-layout/src/lib.rs | 1 + 8 files changed, 280 insertions(+), 88 deletions(-) create mode 100644 vortex-layout/src/expr_cache.rs diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index e665ab18d50..0649630a696 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -5,6 +5,8 @@ use std::iter; use std::sync::Arc; use std::sync::LazyLock; +use std::sync::atomic::AtomicI32; +use std::sync::atomic::Ordering; use bytes::Bytes; use futures::StreamExt; @@ -38,6 +40,7 @@ use vortex_array::dtype::StructFields; use vortex_array::expr::and; use vortex_array::expr::cast; use vortex_array::expr::col; +use vortex_array::expr::dynamic; use vortex_array::expr::eq; use vortex_array::expr::get_item; use vortex_array::expr::gt; @@ -53,6 +56,7 @@ use vortex_array::extension::datetime::Timestamp; use vortex_array::extension::datetime::TimestampOptions; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::fns::operators::CompareOperator; use vortex_array::scalar_fn::fns::pack::Pack; use vortex_array::scalar_fn::fns::pack::PackOptions; use vortex_array::stats::PRUNING_STATS; @@ -92,6 +96,115 @@ static SESSION: LazyLock = LazyLock::new(|| { session }); +#[cfg(not(target_arch = "wasm32"))] +#[tokio::test] +async fn cached_file_reader_preserves_literal_and_dynamic_queries() -> VortexResult<()> { + let session = array_session() + .with::() + .with::(); + crate::register_default_encodings(&session); + let mut ctx = session.create_execution_ctx(); + let array = StructArray::from_fields(&[ + ( + "id", + Buffer::from((0..256i32).collect::>()).into_array(), + ), + ( + "text", + VarBinArray::from_iter_nonnull( + (0..256).map(|i| if i % 2 == 0 { "even" } else { "odd" }), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(), + ), + ( + "nullable", + PrimitiveArray::from_option_iter((0..256i32).map(|i| (i % 7 != 0).then_some(i))) + .into_array(), + ), + ])? + .into_array(); + let mut bytes = ByteBufferMut::empty(); + session + .write_options() + .write(&mut bytes, array.clone().to_array_stream()) + .await?; + let path = std::env::temp_dir().join(format!("vortex-reader-cache-{}.vx", std::process::id())); + tokio::fs::write(&path, bytes.as_slice()).await?; + let file = session + .open_options() + .with_layout_reader_cache() + .open_path(&path) + .await?; + let result = async { + let reader = file.layout_reader()?; + assert!(Arc::ptr_eq(&reader, &file.layout_reader()?)); + for value in 0..1_024i32 { + let actual = file + .scan()? + .with_filter(eq(col("id"), lit(value))) + .into_array_stream()? + .read_all() + .await?; + let start = usize::try_from(value).unwrap().min(256); + let expected = array.slice(start..(start + 1).min(256))?; + assert_arrays_eq!(actual, expected, &mut ctx); + } + + let low = Arc::new(AtomicI32::new(255)); + let high = Arc::new(AtomicI32::new(0)); + let low_value = Arc::clone(&low); + let high_value = Arc::clone(&high); + let dtype = DType::Primitive(I32, Nullability::NonNullable); + let low_expr = dynamic( + CompareOperator::Gt, + move || Some(low_value.load(Ordering::Relaxed).into()), + dtype.clone(), + true, + col("id"), + ); + let high_expr = dynamic( + CompareOperator::Gt, + move || Some(high_value.load(Ordering::Relaxed).into()), + dtype, + true, + col("id"), + ); + for (a, b) in [(255, 0), (0, 255), (127, 63), (63, 127)] { + low.store(a, Ordering::Relaxed); + high.store(b, Ordering::Relaxed); + let scan_a = file + .scan()? + .with_filter(low_expr.clone()) + .into_array_stream()?; + let scan_b = file + .scan()? + .with_filter(high_expr.clone()) + .into_array_stream()?; + let (actual_a, actual_b) = futures::try_join!(scan_a.read_all(), scan_b.read_all())?; + assert_arrays_eq!( + actual_a, + array.slice(usize::try_from(a + 1).unwrap()..256)?, + &mut ctx + ); + assert_arrays_eq!( + actual_b, + array.slice(usize::try_from(b + 1).unwrap()..256)?, + &mut ctx + ); + } + drop(low_expr); + drop(high_expr); + assert_eq!(Arc::strong_count(&low), 1); + assert_eq!(Arc::strong_count(&high), 1); + Ok::<_, vortex_error::VortexError>(()) + } + .await; + drop(file); + tokio::fs::remove_file(&path).await?; + result +} + #[tokio::test] async fn test_eof_values() { // this test exists as a reminder to think about whether we should increment the version diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index fc14eaf57d7..b7093abb7a3 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -21,9 +21,9 @@ use vortex_layout::LayoutReader; use vortex_layout::LayoutReaderRef; use vortex_layout::RowSplits; use vortex_layout::SplitRange; +use vortex_layout::expr_cache::ExpressionCache; use vortex_mask::Mask; use vortex_session::VortexSession; -use vortex_utils::aliases::dash_map::DashMap; use crate::FileStatistics; use crate::pruning::can_prune_file_stats; @@ -41,7 +41,7 @@ pub struct FileStatsLayoutReader { file_stats: FileStatistics, struct_fields: StructFields, session: VortexSession, - prune_cache: DashMap, + prune_cache: ExpressionCache, } impl FileStatsLayoutReader { @@ -116,17 +116,9 @@ impl LayoutReader for FileStatsLayoutReader { expr: &Expression, mask: Mask, ) -> VortexResult { - // Check cache first with read-only lock. - if let Some(pruned) = self.prune_cache.get(expr) { - if *pruned { - return Ok(MaskFuture::ready(Mask::new_false(mask.len()))); - } - return self.child.pruning_evaluation(row_range, expr, mask); - } - - // Evaluate and cache. - let pruned = self.evaluate_file_stats(expr)?; - self.prune_cache.insert(expr.clone(), pruned); + let pruned = self + .prune_cache + .get_or_try_insert_with(expr, || self.evaluate_file_stats(expr))?; if pruned { Ok(MaskFuture::ready(Mask::new_false(mask.len()))) diff --git a/vortex-layout/src/expr_cache.rs b/vortex-layout/src/expr_cache.rs new file mode 100644 index 00000000000..56f22e0208e --- /dev/null +++ b/vortex-layout/src/expr_cache.rs @@ -0,0 +1,129 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Bounded memoization for expressions evaluated by reusable layout readers. + +use std::convert::Infallible; + +use parking_lot::Mutex; +use vortex_array::expr::ExactExpr; +use vortex_array::expr::Expression; +use vortex_array::scalar_fn::fns::dynamic::DynamicComparison; +use vortex_utils::aliases::hash_map::HashMap; + +const MAX_ENTRIES: usize = 32; + +/// An entry-bounded cache that never retains dynamic expression state. +/// +/// A full cache is cleared on insertion. Initializers and evicted-value destructors run +/// outside its lock; concurrent misses may initialize the same key more than once. +/// The entry bound does not account for array buffers retained by cached values. +pub struct ExpressionCache { + entries: Mutex>, +} + +impl Default for ExpressionCache { + fn default() -> Self { + Self { + entries: Mutex::new(HashMap::default()), + } + } +} + +impl ExpressionCache { + /// Return a cached value or initialize it, without caching dynamic expressions. + pub fn get_or_insert_with(&self, expr: &Expression, init: impl FnOnce() -> V) -> V { + match self.get_or_try_insert_with(expr, || Ok::<_, Infallible>(init())) { + Ok(value) => value, + Err(never) => match never {}, + } + } + + /// Return a cached value or initialize it, without caching errors or dynamic expressions. + pub fn get_or_try_insert_with( + &self, + expr: &Expression, + init: impl FnOnce() -> Result, + ) -> Result { + if contains_dynamic(expr) { + return init(); + } + let key = ExactExpr(expr.clone()); + if let Some(value) = self.entries.lock().get(&key).cloned() { + return Ok(value); + } + let value = init()?; + let mut entries = self.entries.lock(); + if let Some(existing) = entries.get(&key) { + return Ok(existing.clone()); + } + let retired = (entries.len() >= MAX_ENTRIES).then(|| std::mem::take(&mut *entries)); + entries.insert(key, value.clone()); + drop(entries); + drop(retired); + Ok(value) + } + + #[cfg(test)] + pub(crate) fn len(&self) -> usize { + self.entries.lock().len() + } +} + +fn contains_dynamic(expr: &Expression) -> bool { + expr.is::() || expr.children().iter().any(contains_dynamic) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use vortex_array::expr::lit; + + use super::*; + + #[test] + fn concurrent_insertions_stay_bounded() { + let cache = ExpressionCache::default(); + std::thread::scope(|scope| { + for worker in 0..8 { + let cache = &cache; + scope.spawn(move || { + for value in 0..1_000i32 { + let key = lit(worker * 1_000 + value); + assert_eq!(cache.get_or_insert_with(&key, || value), value); + assert!(cache.len() <= MAX_ENTRIES); + } + }); + } + }); + } + + #[test] + fn eviction_releases_cached_values() { + let cache = ExpressionCache::default(); + let value = Arc::new(0); + let weak = Arc::downgrade(&value); + cache.get_or_insert_with(&lit(0i32), || value); + assert!(weak.upgrade().is_some()); + for key in 1..=MAX_ENTRIES { + cache.get_or_insert_with(&lit(key as u64), || Arc::new(key)); + } + assert!(weak.upgrade().is_none()); + } + + #[test] + fn failed_initialization_is_retried() { + let cache = ExpressionCache::default(); + let key = lit(1i32); + assert_eq!( + cache.get_or_try_insert_with(&key, || Err::("error")), + Err("error") + ); + assert_eq!( + cache.get_or_try_insert_with(&key, || Ok::<_, &str>(7)), + Ok(7) + ); + assert_eq!(cache.get_or_insert_with(&key, || 8), 7); + } +} diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 30172eb45d2..6a91cfc64e4 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -33,13 +33,13 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; -use vortex_utils::aliases::dash_map::DashMap; use super::DictLayout; use crate::LayoutReader; use crate::LayoutReaderRef; use crate::RowSplits; use crate::SplitRange; +use crate::expr_cache::ExpressionCache; use crate::layouts::SharedArrayFuture; use crate::segments::SegmentSource; @@ -53,7 +53,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: ExpressionCache, values: LayoutReaderRef, codes: LayoutReaderRef, @@ -139,23 +139,16 @@ impl DictReader { // shouldn't. // TODO(joe): fixme - // Check cache first with read-only lock - if let Some(fut) = self.values_evals.get(&expr) { - return fut.clone(); - } - - self.values_evals - .entry(expr.clone()) - .or_insert_with(|| { - self.values_array_uncanonical() - .map(move |array| { - let array = array?.apply(&expr)?; - Ok(SharedArray::new(array).into_array()) - }) - .boxed() - .shared() - }) - .clone() + self.values_evals.get_or_insert_with(&expr, || { + let expr = expr.clone(); + self.values_array_uncanonical() + .map(move |array| { + let array = array?.apply(&expr)?; + Ok(SharedArray::new(array).into_array()) + }) + .boxed() + .shared() + }) } } diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index 4e5e8e494fa..9459ab45630 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -8,7 +8,6 @@ use std::fmt::Formatter; use std::ops::BitAnd; use std::ops::Range; use std::sync::Arc; -use std::sync::OnceLock; use Nullability::NonNullable; pub use expr::*; @@ -24,7 +23,6 @@ use vortex_array::dtype::FieldMask; use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::ExactExpr; use vortex_array::expr::Expression; use vortex_array::expr::is_root; use vortex_array::expr::root; @@ -38,19 +36,19 @@ use vortex_mask::Mask; use vortex_sequence::Sequence; use vortex_sequence::SequenceArray; use vortex_session::VortexSession; -use vortex_utils::aliases::dash_map::DashMap; use crate::ArrayFuture; use crate::LayoutReader; use crate::RowSplits; use crate::SplitRange; +use crate::expr_cache::ExpressionCache; use crate::layouts::partitioned::PartitionedExprEval; pub struct RowIdxLayoutReader { name: Arc, row_offset: u64, child: Arc, - partition_cache: DashMap>>, + partition_cache: ExpressionCache, session: VortexSession, } @@ -60,29 +58,14 @@ impl RowIdxLayoutReader { name: Arc::clone(child.name()), row_offset, child, - partition_cache: DashMap::with_hasher(Default::default()), + partition_cache: Default::default(), session, } } fn partition_expr(&self, expr: &Expression) -> VortexResult { - let key = ExactExpr(expr.clone()); - - // Check cache first with read-only lock. - if let Some(entry) = self.partition_cache.get(&key) - && let Some(partitioning) = entry.value().get() - { - return Ok(partitioning.clone()); - } - - let result = self.compute_partitioning(expr)?; - self.partition_cache - .entry(key) - .or_insert_with(|| Arc::new(OnceLock::new())) - .get_or_init(|| result.clone()); - - Ok(result) + .get_or_try_insert_with(expr, || self.compute_partitioning(expr)) } fn compute_partitioning(&self, expr: &Expression) -> VortexResult { diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index a667d86fd8a..29fdd35f9da 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -3,7 +3,6 @@ use std::ops::Range; use std::sync::Arc; -use std::sync::OnceLock; use futures::try_join; use itertools::Itertools; @@ -19,7 +18,6 @@ use vortex_array::dtype::FieldMask; use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; -use vortex_array::expr::ExactExpr; use vortex_array::expr::Expression; use vortex_array::expr::col; use vortex_array::expr::make_free_field_annotator; @@ -35,7 +33,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::VortexSession; -use vortex_utils::aliases::dash_map::DashMap; use vortex_utils::aliases::hash_map::HashMap; use crate::ArrayFuture; @@ -44,6 +41,7 @@ use crate::LayoutReaderRef; use crate::LazyReaderChildren; use crate::RowSplits; use crate::SplitRange; +use crate::expr_cache::ExpressionCache; use crate::layouts::partitioned::PartitionedExprEval; use crate::layouts::struct_::StructLayout; use crate::segments::SegmentSource; @@ -59,7 +57,7 @@ pub struct StructReader { expanded_root_expr: Expression, field_lookup: Option>, - partitioned_expr_cache: DashMap>>, + partitioned_expr_cache: ExpressionCache, } impl StructReader { @@ -156,18 +154,8 @@ impl StructReader { /// Utility for partitioning an expression over the fields of a struct. fn partition_expr(&self, expr: Expression) -> VortexResult { - let key = ExactExpr(expr.clone()); - let binding = self - .partitioned_expr_cache - .entry(key) - .or_insert_with(|| Arc::new(OnceLock::new())); - let entry = binding.value(); - if let Some(value) = entry.get() { - return Ok(value.clone()); - } - let result = self.compute_partitioned_expr(expr)?; - let result = entry.get_or_init(|| result); - Ok(result.clone()) + self.partitioned_expr_cache + .get_or_try_insert_with(&expr, || self.compute_partitioned_expr(expr.clone())) } fn compute_partitioned_expr(&self, expr: Expression) -> VortexResult { diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index b517df985c9..2710fefd7fa 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -26,16 +26,15 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; -use vortex_utils::aliases::dash_map::DashMap; use crate::LazyReaderChildren; +use crate::expr_cache::ExpressionCache; use crate::layouts::zoned::ZonedLayout; use crate::layouts::zoned::zone_map::ZoneMap; type SharedZoneMap = Shared>>; pub(super) type SharedPruningResult = Shared>>>; -type PredicateCache = Arc>>; pub(super) struct PruningState { zone_count: usize, @@ -46,9 +45,9 @@ pub(super) struct PruningState { lazy_children: Arc, session: VortexSession, - pruning_result: LazyLock>>, + pruning_result: LazyLock>>, zone_map: OnceLock, - pruning_predicates: LazyLock>>, + pruning_predicates: LazyLock>>, } impl PruningState { @@ -73,13 +72,8 @@ impl PruningState { } pub(super) fn pruning_mask_future(&self, expr: Expression) -> Option { - if let Some(result) = self.pruning_result.get(&expr) { - return result.value().clone(); - } - - self.pruning_result - .entry(expr.clone()) - .or_insert_with(|| match self.pruning_predicate(expr.clone()) { + self.pruning_result.get_or_insert_with(&expr, || { + match self.pruning_predicate(expr.clone()) { None => { trace!(%expr, "no pruning predicate"); None @@ -89,6 +83,7 @@ impl PruningState { let zone_map = self.zone_map(); let dynamic_updates = DynamicExprUpdates::new(&expr); let session = self.session.clone(); + let expr = expr.clone(); Some( async move { @@ -112,22 +107,20 @@ impl PruningState { .shared(), ) } - }) - .clone() + } + }) } fn pruning_predicate(&self, expr: Expression) -> Option { - self.pruning_predicates - .entry(expr.clone()) - .or_default() - .get_or_init(move || match expr.falsify(&self.dtype, &self.session) { + self.pruning_predicates.get_or_insert_with(&expr, || { + match expr.falsify(&self.dtype, &self.session) { Ok(predicate) => predicate, Err(error) => { trace!(%expr, %error, "failed to construct stats rewrite predicate"); None } - }) - .clone() + } + }) } fn zone_map(&self) -> SharedZoneMap { diff --git a/vortex-layout/src/lib.rs b/vortex-layout/src/lib.rs index 6612f5e0350..28ea6067fe0 100644 --- a/vortex-layout/src/lib.rs +++ b/vortex-layout/src/lib.rs @@ -29,6 +29,7 @@ pub mod aliases; mod children; pub mod display; mod encoding; +pub mod expr_cache; mod flatbuffers; mod layout; mod reader;