Speed up V2 scans and default datafusion-bench to V2 for TPC-DS and ClickBench

benchmarks/datafusion-bench/src/main.rs

@@ -120,7 +120,7 @@ struct Args {
     #[arg(long, value_delimiter = ',', value_parser = value_parser!(Format))]
     formats: Vec<Format>,
 
-    /// Use the planner-driven V2 Vortex scan (defaults to true for TPC-H).
+    /// Use the planner-driven V2 Vortex scan (defaults to true for TPC-H, TPC-DS and ClickBench).
     #[arg(long, action = clap::ArgAction::Set)]
     scan_v2: Option<bool>,
 
@@ -137,8 +137,10 @@ async fn main() -> anyhow::Result<()> {
     let args = Args::parse();
     let opts = Opts::from(args.options);
     let scan_v2 = args.scan_v2.unwrap_or_else(|| {
-        matches!(args.benchmark, BenchmarkArg::TpcH)
-            && !std::env::var("VORTEX_SCAN_V2").is_ok_and(|v| v == "0")
+        matches!(
+            args.benchmark,
+            BenchmarkArg::TpcH | BenchmarkArg::TpcDS | BenchmarkArg::ClickBench
+        ) && !std::env::var("VORTEX_SCAN_V2").is_ok_and(|v| v == "0")
     });
 
     set_join_set_tracer(get_static_tracer())?;

vortex-layout/src/scan/execution/chunk.rs

@@ -7,6 +7,7 @@ use std::ops::Range;
 use std::sync::Arc;
 use std::vec::IntoIter;
 
+use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_io::request::IoConsumer;
 use vortex_io::request::IoRequestId;
@@ -58,6 +59,24 @@ pub(crate) fn row_chunks(rows: Range<u64>, rows_per_chunk: usize) -> Vec<Range<u
         .collect()
 }
 
+/// The chunks a scan of `rows` plans: [`row_chunks`], with a chunk at either edge of `rows` that
+/// is under half of `rows_per_chunk` joined to its neighbour. An engine cuts a file into
+/// partitions whose edges fall off the chunk grid, and a short edge chunk would cost a whole
+/// chunk's planning and conjunct setup for a few rows.
+pub(crate) fn scan_chunks(rows: Range<u64>, rows_per_chunk: usize) -> Vec<Range<u64>> {
+    let mut chunks = row_chunks(rows, rows_per_chunk);
+    let short = |chunk: &Range<u64>| chunk.end - chunk.start < rows_per_chunk as u64 / 2;
+    if chunks.len() > 1 && chunks.last().is_some_and(short) {
+        let last = chunks.pop().vortex_expect("more than one chunk");
+        chunks.last_mut().vortex_expect("more than one chunk").end = last.end;
+    }
+    if chunks.len() > 1 && short(&chunks[0]) {
+        let first = chunks.remove(0);
+        chunks[0].start = first.start;
+    }
+    chunks
+}
+
 /// Emits one planner per chunk of a scan's rows, in row order, each spawned as independent so
 /// the driver may run it beside the others. A chunk none of whose rows is selected and kept is
 /// never planned.
@@ -79,7 +98,7 @@ impl ChunkPlanner {
         selection: Selection,
         kept: Option<Mask>,
     ) -> Self {
-        let chunks = row_chunks(rows.clone(), scan.chunk_rows).into_iter();
+        let chunks = scan_chunks(rows.clone(), scan.chunk_rows).into_iter();
         Self {
             scan,
             rows,

vortex-layout/src/scan/execution/mod.rs

@@ -37,6 +37,8 @@ pub(crate) use self::chunk::ROW_BLOCK_SIZE;
 #[cfg(test)]
 pub(crate) use self::chunk::row_chunks;
 pub(crate) use self::chunk::rows_per_chunk;
+#[cfg(test)]
+pub(crate) use self::chunk::scan_chunks;
 use self::filter::FilterPlanner;
 pub use self::filter::FilterPlans;
 use self::projection::ProjectionPlanner;

vortex-layout/src/scan/v2/cache.rs

@@ -3,18 +3,22 @@
 
 //! State derived from a [`ScanFile`](super::ScanFile) that repeated V2 scans can share.
 
+use std::hash::Hash;
+use std::hash::Hasher;
 use std::sync::Arc;
 use std::sync::Weak;
 
 use parking_lot::Mutex;
 use rustc_hash::FxHashMap;
 use vortex_array::expr::BoundExpression;
-use vortex_array::expr::ExactBoundExpr;
 
 use super::io::SegmentRanges;
+use super::retain::RetainedSegments;
+use super::retain::RetainingSegmentSource;
 use crate::plan::physical::PlanRef;
 use crate::plan::physical::pipeline::ValueCache;
 use crate::scan::execution::FilterPlans;
+use crate::segments::SegmentSource;
 
 /// Reusable state for repeated V2 scans of one file, the V2 counterpart of a cached
 /// [`LayoutReader`](crate::LayoutReader) tree.
@@ -25,7 +29,8 @@ use crate::scan::execution::FilterPlans;
 /// scans. The lowered plan keeps the children it has initialized, each distinct projection keeps
 /// its optimized plan, each distinct filter keeps its conjunct and pruning plans along with the
 /// zones those plans have already excluded, and dictionary values stay keyed by the cached plan
-/// that produced them.
+/// that produced them. Recently read segments are retained up to a byte budget, so the chunks of
+/// the file's scans that share a stored segment read it once.
 ///
 /// A cache describes one file's layout, segment locations, and segment source. Use a new cache
 /// when any of them changes. Cached plans and values live as long as the cache, one set per
@@ -38,15 +43,26 @@ struct Inner {
     root: Mutex<Option<PlanRef>>,
     /// Optimized plans of non-identity projections. Values are cached by the plan that produced
     /// them, so a projection must reuse its plan for its values to be found again.
-    projections: Mutex<FxHashMap<ExactBoundExpr, PlanRef>>,
+    projections: Mutex<FxHashMap<ExprKey, PlanRef>>,
     ranges: Mutex<Option<SegmentRanges>>,
-    /// Plans of filters, keyed by the exact filter expression. A pruning plan remembers the zones it
+    /// Plans of filters, keyed by the filter expression. A pruning plan remembers the zones it
     /// excluded, so sharing the plans prunes a file's zones once for every scan of that filter.
-    filters: Mutex<FxHashMap<ExactBoundExpr, Arc<FilterPlans>>>,
+    filters: Mutex<FxHashMap<ExprKey, Arc<FilterPlans>>>,
     values: ValueCache,
+    /// Recently read segments, shared by the chunks of every scan of the file.
+    retained: Arc<RetainedSegments>,
 }
 
 impl ScanCache {
+    /// A cache whose scans retain up to `bytes` of recently read segments, so chunks that share
+    /// a stored segment read it once. The default retains 64 MiB; zero retains nothing.
+    pub fn with_retained_bytes(bytes: u64) -> Self {
+        Self(Arc::new(Inner {
+            retained: Arc::new(RetainedSegments::new(bytes)),
+            ..Inner::default()
+        }))
+    }
+
     /// The lowered root plan an earlier scan kept.
     pub(super) fn root(&self) -> Option<PlanRef> {
         self.0.root.lock().clone()
@@ -62,7 +78,7 @@ impl ScanCache {
         self.0
             .projections
             .lock()
-            .get(&ExactBoundExpr(projection.clone()))
+            .get(&ExprKey(projection.clone()))
             .cloned()
     }
 
@@ -72,7 +88,7 @@ impl ScanCache {
         self.0
             .projections
             .lock()
-            .entry(ExactBoundExpr(projection))
+            .entry(ExprKey(projection))
             .or_insert(plan)
             .clone()
     }
@@ -89,11 +105,7 @@ impl ScanCache {
 
     /// The plans of `filter` an earlier scan kept.
     pub(super) fn filter_plans(&self, filter: &BoundExpression) -> Option<Arc<FilterPlans>> {
-        self.0
-            .filters
-            .lock()
-            .get(&ExactBoundExpr(filter.clone()))
-            .cloned()
+        self.0.filters.lock().get(&ExprKey(filter.clone())).cloned()
     }
 
     /// Keeps the plans of `filter`, returning the plans already kept if another scan got there
@@ -107,11 +119,19 @@ impl ScanCache {
             self.0
                 .filters
                 .lock()
-                .entry(ExactBoundExpr(filter))
+                .entry(ExprKey(filter))
                 .or_insert(plans),
         )
     }
 
+    /// Serves `segments` through the segments the file's scans retain.
+    pub(super) fn retaining(&self, segments: Arc<dyn SegmentSource>) -> Arc<dyn SegmentSource> {
+        Arc::new(RetainingSegmentSource {
+            inner: segments,
+            retained: Arc::clone(&self.0.retained),
+        })
+    }
+
     pub(super) fn values(&self) -> ValueCache {
         self.0.values.clone()
     }
@@ -122,6 +142,47 @@ impl ScanCache {
     }
 }
 
+/// A bound expression keyed by structure, so equal expressions bound separately share an entry:
+/// an engine that opens a file once per partition binds its filter and projection once per
+/// partition.
+///
+/// Hashing skips dtypes, which equality still compares, so a lookup does not walk a wide schema.
+/// Both walk the tree iteratively, however deep it is.
+#[derive(Clone)]
+struct ExprKey(BoundExpression);
+
+impl PartialEq for ExprKey {
+    fn eq(&self, other: &Self) -> bool {
+        let mut pending = vec![(&self.0, &other.0)];
+        while let Some((lhs, rhs)) = pending.pop() {
+            if std::ptr::eq(lhs, rhs) {
+                continue;
+            }
+            if lhs.as_scalar() != rhs.as_scalar()
+                || lhs.children().len() != rhs.children().len()
+                || lhs.dtype() != rhs.dtype()
+            {
+                return false;
+            }
+            pending.extend(lhs.children().iter().zip(rhs.children()));
+        }
+        true
+    }
+}
+
+impl Eq for ExprKey {}
+
+impl Hash for ExprKey {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        let mut pending = vec![&self.0];
+        while let Some(expr) = pending.pop() {
+            expr.as_scalar().hash(state);
+            expr.children().len().hash(state);
+            pending.extend(expr.children());
+        }
+    }
+}
+
 /// A [`ScanCache`] handle that does not keep the cache alive. Engines that open a file once per
 /// partition keep one per file, so concurrent partitions share a cache that is freed with the
 /// last scan holding it.
@@ -134,3 +195,61 @@ impl WeakScanCache {
         self.0.upgrade().map(ScanCache)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::hash::BuildHasher;
+
+    use rstest::rstest;
+    use vortex_array::dtype::DType;
+    use vortex_array::dtype::Nullability;
+    use vortex_array::dtype::PType;
+    use vortex_array::expr::Expression;
+    use vortex_array::expr::get_item;
+    use vortex_array::expr::gt;
+    use vortex_array::expr::lit;
+    use vortex_array::expr::lt;
+    use vortex_array::expr::root;
+    use vortex_error::VortexResult;
+
+    use super::ExprKey;
+
+    fn ints_gt(value: i32) -> Expression {
+        gt(get_item("ints", root()), lit(value))
+    }
+
+    fn table(ints: PType) -> DType {
+        DType::struct_(
+            [("ints", DType::Primitive(ints, Nullability::NonNullable))],
+            Nullability::NonNullable,
+        )
+    }
+
+    #[test]
+    fn separately_bound_equal_expressions_share_a_key() -> VortexResult<()> {
+        let dtype = table(PType::I32);
+        let lhs = ExprKey(ints_gt(5).bind(&dtype)?);
+        let rhs = ExprKey(ints_gt(5).bind(&dtype)?);
+        assert!(lhs == rhs);
+        assert_eq!(
+            rustc_hash::FxBuildHasher.hash_one(&lhs),
+            rustc_hash::FxBuildHasher.hash_one(&rhs)
+        );
+        Ok(())
+    }
+
+    #[rstest]
+    #[case::literal(ints_gt(6), table(PType::I32))]
+    #[case::function(lt(get_item("ints", root()), lit(5i32)), table(PType::I32))]
+    #[case::dtype(get_item("ints", root()), table(PType::I64))]
+    fn different_expressions_have_different_keys(
+        #[case] other: Expression,
+        #[case] dtype: DType,
+    ) -> VortexResult<()> {
+        let i32s = table(PType::I32);
+        for key in [ints_gt(5), get_item("ints", root())] {
+            assert!(ExprKey(key.bind(&i32s)?) != ExprKey(other.bind(&dtype)?));
+        }
+        Ok(())
+    }
+}

vortex-layout/src/scan/v2/mod.rs

@@ -18,11 +18,12 @@
 //! serialized bytes per row are high.
 //!
 //! A [`ScanFile`] carrying a [`ScanCache`] shares its lowered plan, filter plans and pruned
-//! zones, segment index, and dictionary values with every other scan of that file, like a cached
-//! layout reader does for V1 scans.
+//! zones, segment index, dictionary values, and recently read segments with every other scan of
+//! that file, like a cached layout reader does for V1 scans.
 
 mod cache;
 pub(crate) mod io;
+mod retain;
 #[cfg(test)]
 mod tests;
 
@@ -206,7 +207,7 @@ impl ScanBuilder {
             session: self.session,
             rows: self.row_range.unwrap_or(0..row_count),
             selection: self.selection,
-            segments: file.segments,
+            segments: cache.retaining(file.segments),
             ranges,
             cache,
             locations: file.locations,