Bound reader expression caches and isolate dynamic query state

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<VortexSession> = 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::<LayoutSession>()
+        .with::<RuntimeSession>();
+    crate::register_default_encodings(&session);
+    let mut ctx = session.create_execution_ctx();
+    let array = StructArray::from_fields(&[
+        (
+            "id",
+            Buffer::from((0..256i32).collect::<Vec<_>>()).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

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<Expression, bool>,
+    prune_cache: ExpressionCache<bool>,
 }
 
 impl FileStatsLayoutReader {
@@ -116,17 +116,9 @@ impl LayoutReader for FileStatsLayoutReader {
         expr: &Expression,
         mask: Mask,
     ) -> VortexResult<MaskFuture> {
-        // 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())))

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<V> {
+    entries: Mutex<HashMap<ExactExpr, V>>,
+}
+
+impl<V> Default for ExpressionCache<V> {
+    fn default() -> Self {
+        Self {
+            entries: Mutex::new(HashMap::default()),
+        }
+    }
+}
+
+impl<V: Clone> ExpressionCache<V> {
+    /// 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<E>(
+        &self,
+        expr: &Expression,
+        init: impl FnOnce() -> Result<V, E>,
+    ) -> Result<V, E> {
+        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::<DynamicComparison>() || 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::<i32, _>("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);
+    }
+}

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<SharedArrayFuture>,
     /// Cache of expression evaluation results on the values array by expression
-    values_evals: DashMap<Expression, SharedArrayFuture>,
+    values_evals: ExpressionCache<SharedArrayFuture>,
 
     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()
+        })
     }
 }