Plan a V2 scan as a chain of planners from one prepared scan

vortex-datafusion/src/persistent/opener.rs

@@ -56,6 +56,8 @@ use vortex::io::InstrumentedReadAt;
 use vortex::layout::LayoutReader;
 use vortex::layout::scan::scan_builder::ScanBuilder;
 use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder;
+use vortex::layout::scan::v2::ScanCache;
+use vortex::layout::scan::v2::WeakScanCache;
 use vortex::layout::scan::v2::partition_row_range;
 use vortex::layout::segments::SegmentSource;
 use vortex::metrics::Label;
@@ -64,6 +66,7 @@ use vortex::session::VortexSession;
 use vortex_arrow::ArrowSessionExt;
 use vortex_utils::aliases::dash_map::DashMap;
 use vortex_utils::aliases::dash_map::Entry;
+use vortex_utils::parallelism::get_available_parallelism;
 
 use crate::VortexAccessPlan;
 use crate::convert::exprs::ExpressionConverter;
@@ -112,6 +115,9 @@ pub(crate) struct VortexOpener {
     /// boundary and the file's metadata segments are read once, as they are for V1 through the
     /// shared layout reader.
     pub segment_sources: Arc<DashMap<Path, Weak<dyn SegmentSource>>>,
+    /// The V2 scan cache each open file's partitions share: its lowered layout, filter and
+    /// projection plans, proved zones and dictionary values, built once per file.
+    pub scan_caches: Arc<DashMap<Path, WeakScanCache>>,
     /// Shared V1 full-file natural splits keyed by file path.
     pub natural_splits: Arc<DashMap<Path, Arc<NaturalSplits>>>,
     /// Shared physical Arrow schemas of each file, keyed by file path.
@@ -125,6 +131,9 @@ pub(crate) struct VortexOpener {
     pub projection_pushdown: bool,
     pub scan_concurrency: Option<usize>,
     pub scan_v2: bool,
+    /// How many partitions the engine runs the scan on, so a V2 scan parallelises inside a
+    /// partition only as far as the engine has not.
+    pub partitions: usize,
 }
 
 impl FileOpener for VortexOpener {
@@ -158,10 +167,12 @@ impl FileOpener for VortexOpener {
         let limit = self.limit;
         let layout_readers = Arc::clone(&self.layout_readers);
         let segment_sources = Arc::clone(&self.segment_sources);
+        let scan_caches = Arc::clone(&self.scan_caches);
         let natural_splits = Arc::clone(&self.natural_splits);
         let file_schemas = Arc::clone(&self.file_schemas);
         let has_output_ordering = self.has_output_ordering;
         let scan_concurrency = self.scan_concurrency;
+        let partitions = self.partitions.max(1);
         let scan_v2 = self.scan_v2;
 
         let expr_converter = Arc::clone(&self.expression_converter);
@@ -423,6 +434,7 @@ impl FileOpener for VortexOpener {
                     &file.object_meta.location,
                     scan_file.segments,
                 );
+                scan_file.cache = Some(shared_scan_cache(&scan_caches, &file.object_meta.location));
                 let mut builder = V2ScanBuilder::new(vxf.session().clone(), scan_file)
                     .with_projection(scan_projection)
                     .with_some_filter(native_filter)
@@ -442,9 +454,15 @@ impl FileOpener for VortexOpener {
                 {
                     builder = builder.with_selection(selection.clone());
                 }
-                if let Some(concurrency) = scan_concurrency {
-                    builder = builder.with_concurrency(concurrency);
-                }
+                // Unless set, chunks run beside each other only as far as the engine's partitions
+                // leave cores idle: at one partition per core they run in the scan's own run.
+                let concurrency = scan_concurrency.unwrap_or_else(|| {
+                    get_available_parallelism()
+                        .unwrap_or(1)
+                        .div_ceil(partitions)
+                        .max(1)
+                });
+                builder = builder.with_concurrency(concurrency);
                 builder.map(map_batch).into_stream()
             } else {
                 // Share V1 layout readers across partitions of the file.
@@ -643,6 +661,28 @@ fn contains_lambda(projection: &ProjectionExprs) -> DFResult<bool> {
     Ok(false)
 }
 
+/// Returns the scan cache other partitions of `path` are planning with, or registers a new one.
+/// A scan holds its cache until its stream ends, so the cache lives as long as one scan.
+fn shared_scan_cache(scan_caches: &DashMap<Path, WeakScanCache>, path: &Path) -> ScanCache {
+    match scan_caches.entry(path.clone()) {
+        Entry::Occupied(mut entry) => {
+            if let Some(shared) = entry.get().upgrade() {
+                tracing::trace!("reusing scan cache for {path}");
+                shared
+            } else {
+                let fresh = ScanCache::default();
+                entry.insert(fresh.downgrade());
+                fresh
+            }
+        }
+        Entry::Vacant(entry) => {
+            let fresh = ScanCache::default();
+            entry.insert(fresh.downgrade());
+            fresh
+        }
+    }
+}
+
 /// Returns the segment source other partitions of `path` are reading through, or registers
 /// `fresh` as that source. Like the layout readers, a source lives only as long as one scan.
 fn shared_segment_source(
@@ -1074,6 +1114,8 @@ mod tests {
             df_metrics: ExecutionPlanMetricsSet::new(),
             layout_readers: Default::default(),
             segment_sources: Default::default(),
+            scan_caches: Default::default(),
+            partitions: 1,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,
@@ -1482,6 +1524,8 @@ mod tests {
             df_metrics: ExecutionPlanMetricsSet::new(),
             layout_readers: Default::default(),
             segment_sources: Default::default(),
+            scan_caches: Default::default(),
+            partitions: 1,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,
@@ -1572,6 +1616,8 @@ mod tests {
             df_metrics: ExecutionPlanMetricsSet::new(),
             layout_readers: Default::default(),
             segment_sources: Default::default(),
+            scan_caches: Default::default(),
+            partitions: 1,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,
@@ -1753,6 +1799,8 @@ mod tests {
             df_metrics: ExecutionPlanMetricsSet::new(),
             layout_readers: Default::default(),
             segment_sources: Default::default(),
+            scan_caches: Default::default(),
+            partitions: 1,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,
@@ -1816,6 +1864,8 @@ mod tests {
             df_metrics: ExecutionPlanMetricsSet::new(),
             layout_readers: Default::default(),
             segment_sources: Default::default(),
+            scan_caches: Default::default(),
+            partitions: 1,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,
@@ -2030,6 +2080,8 @@ mod tests {
             df_metrics: ExecutionPlanMetricsSet::new(),
             layout_readers: Default::default(),
             segment_sources: Default::default(),
+            scan_caches: Default::default(),
+            partitions: 1,
             natural_splits: Default::default(),
             file_schemas: Default::default(),
             has_output_ordering: false,

vortex-datafusion/src/persistent/source.rs

@@ -34,6 +34,7 @@ use object_store::ObjectStore;
 use object_store::path::Path;
 use vortex::file::VORTEX_FILE_EXTENSION;
 use vortex::layout::LayoutReader;
+use vortex::layout::scan::v2::WeakScanCache;
 use vortex::layout::segments::SegmentSource;
 use vortex::metrics::DefaultMetricsRegistry;
 use vortex::metrics::MetricsRegistry;
@@ -203,6 +204,7 @@ pub struct VortexSource {
     layout_readers: Arc<DashMap<Path, Weak<dyn LayoutReader>>>,
     /// Shared V2 segment sources, keyed by path and alive only while a scan reads the file.
     segment_sources: Arc<DashMap<Path, Weak<dyn SegmentSource>>>,
+    scan_caches: Arc<DashMap<Path, WeakScanCache>>,
     /// Shared full-file natural splits keyed by path.
     natural_splits: Arc<DashMap<Path, Arc<NaturalSplits>>>,
     /// Shared physical Arrow schemas of each file, keyed by path.
@@ -240,6 +242,7 @@ impl VortexSource {
             df_metrics: Default::default(),
             layout_readers: Arc::new(DashMap::default()),
             segment_sources: Arc::new(DashMap::default()),
+            scan_caches: Arc::new(DashMap::default()),
             natural_splits: Arc::new(DashMap::default()),
             file_schemas: Arc::new(DashMap::default()),
             expression_converter,
@@ -366,6 +369,7 @@ impl VortexSource {
             df_metrics: self.df_metrics.clone(),
             layout_readers: Arc::clone(&self.layout_readers),
             segment_sources: Arc::clone(&self.segment_sources),
+            scan_caches: Arc::clone(&self.scan_caches),
             natural_splits: Arc::clone(&self.natural_splits),
             file_schemas: Arc::clone(&self.file_schemas),
             has_output_ordering: !base_config.output_ordering.is_empty() || self.ordered,
@@ -374,6 +378,7 @@ impl VortexSource {
             projection_pushdown: self.options.projection_pushdown,
             scan_concurrency: self.options.scan_concurrency,
             scan_v2: self.options.scan_v2,
+            partitions: base_config.file_groups.len(),
         };
 
         Ok(opener)

vortex-layout/src/plan/physical/optimizer/mod.rs

@@ -24,11 +24,12 @@ use super::Take;
 use super::Zoned;
 use super::plans::EvalIdentityRule;
 use super::plans::ExpressionConcatRule;
+use super::plans::ExpressionLegacyStatsRule;
 use super::plans::ExpressionPackRule;
 use super::plans::ExpressionTakeRule;
 use super::plans::ExpressionZonedRule;
 use super::plans::LegacyStatsRule;
-use super::plans::ZonedLegacyStatsRule;
+use super::plans::ZonedDataRule;
 use crate::plan::physical::plans::EvalEvalRule;
 use crate::plan::physical::plans::EvalPackRule;
 
@@ -38,10 +39,14 @@ static EVAL_IDENTITY_RULE: PlanReduceRuleAdapter<Eval, EvalIdentityRule> =
 static LEGACY_STATS_RULE: PlanReduceRuleAdapter<LegacyStats, LegacyStatsRule> =
     PlanReduceRuleAdapter::new(LegacyStatsRule);
 
+static ZONED_DATA_RULE: PlanReduceRuleAdapter<Zoned, ZonedDataRule> =
+    PlanReduceRuleAdapter::new(ZonedDataRule);
+
 static PLAN_RULES: PlanRuleSet = PlanRuleSet::new(&[
     &EVAL_IDENTITY_RULE,
     &PlanReduceRuleAdapter::new(EvalPackRule),
     &LEGACY_STATS_RULE,
+    &ZONED_DATA_RULE,
 ]);
 
 static EXPRESSION_CONCAT_RULE: PlanParentReduceRuleAdapter<Concat, ExpressionConcatRule> =
@@ -52,17 +57,18 @@ static EXPRESSION_PACK_RULE: PlanParentReduceRuleAdapter<Pack, ExpressionPackRul
     PlanParentReduceRuleAdapter::new(ExpressionPackRule);
 static EXPRESSION_ZONED_RULE: PlanParentReduceRuleAdapter<Zoned, ExpressionZonedRule> =
     PlanParentReduceRuleAdapter::new(ExpressionZonedRule);
-
-static ZONED_LEGACY_STATS_RULE: PlanParentReduceRuleAdapter<LegacyStats, ZonedLegacyStatsRule> =
-    PlanParentReduceRuleAdapter::new(ZonedLegacyStatsRule);
+static EXPRESSION_LEGACY_STATS_RULE: PlanParentReduceRuleAdapter<
+    LegacyStats,
+    ExpressionLegacyStatsRule,
+> = PlanParentReduceRuleAdapter::new(ExpressionLegacyStatsRule);
 
 static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[
     &PlanParentReduceRuleAdapter::new(EvalEvalRule),
     &EXPRESSION_CONCAT_RULE,
     &EXPRESSION_TAKE_RULE,
     &EXPRESSION_PACK_RULE,
     &EXPRESSION_ZONED_RULE,
-    &ZONED_LEGACY_STATS_RULE,
+    &EXPRESSION_LEGACY_STATS_RULE,
 ]);
 
 /// Attempts a static rewrite for `plan`.

vortex-layout/src/plan/physical/plans/legacy_stats.rs

@@ -15,6 +15,7 @@ use vortex_mask::Mask;
 use vortex_session::registry::CachedId;
 
 use crate::layouts::zoned::legacy_stats_to_aggregates;
+use crate::plan::physical::Eval;
 use crate::plan::physical::EvalPlan;
 use crate::plan::physical::Plan;
 use crate::plan::physical::PlanChildren;
@@ -29,6 +30,7 @@ use crate::plan::physical::optimizer::PlanParentReduceRule;
 use crate::plan::physical::optimizer::PlanReduceRule;
 use crate::plan::physical::pipeline::Chain;
 use crate::plan::physical::pipeline::Compiler;
+use crate::plan::physical::plans::zoned::is_zone_proof;
 
 const DATA: usize = 0;
 const ZONES: usize = 1;
@@ -37,8 +39,9 @@ const ZONES: usize = 1;
 /// than by aggregate function.
 ///
 /// The optimizer rewrites it into a [`Zoned`] plan over the zone table converted to the aggregate
-/// schema, or into its data alone when a [`Zoned`] plan already summarises the same data. Compiled
-/// unoptimized, it reads its data and prunes nothing.
+/// schema, or into its data alone when a [`Zoned`] plan already summarises the same data; a plan
+/// that reads the data is then optimized to the data alone. Compiled unoptimized, it reads its
+/// data and prunes nothing.
 #[derive(Clone, Debug)]
 pub struct LegacyStats;
 
@@ -173,24 +176,28 @@ impl PlanReduceRule<LegacyStats> for LegacyStatsRule {
     }
 }
 
-/// Drops legacy stats beneath a [`Zoned`] plan: its aggregate zone map already summarises the
-/// same rows.
+/// Converts the legacy stats beneath a zone proof in place, so the proof meets a [`Zoned`] plan
+/// before the data path optimizes the stats away.
 #[derive(Debug)]
-pub(crate) struct ZonedLegacyStatsRule;
+pub(crate) struct ExpressionLegacyStatsRule;
 
-impl PlanParentReduceRule<LegacyStats> for ZonedLegacyStatsRule {
-    type Parent = Zoned;
+impl PlanParentReduceRule<LegacyStats> for ExpressionLegacyStatsRule {
+    type Parent = Eval;
 
     fn reduce_parent(
         &self,
         child: &LegacyStatsPlan,
-        parent: &ZonedPlan,
-        child_idx: usize,
+        parent: &Plan<Eval>,
+        _child_idx: usize,
     ) -> VortexResult<Option<PlanRef>> {
-        if child_idx != DATA || parent.is_pruning() {
+        if !is_zone_proof(parent)? {
             return Ok(None);
         }
-        let children = vec![child.data_plan()?, parent.zones_plan()?];
-        Ok(Some(parent.clone().into_plan().with_children(children)?))
+        let Some(converted) = LegacyStatsRule.reduce(child)? else {
+            return Ok(None);
+        };
+        Ok(Some(
+            EvalPlan::try_new(parent.expression().clone(), converted)?.into_plan(),
+        ))
     }
 }

vortex-layout/src/plan/physical/plans/mod.rs

@@ -26,11 +26,11 @@ pub use eval::EvalPlan;
 pub use filter::Filter;
 pub use filter::FilterData;
 pub use filter::FilterPlan;
+pub(crate) use legacy_stats::ExpressionLegacyStatsRule;
 pub use legacy_stats::LegacyStats;
 pub use legacy_stats::LegacyStatsData;
 pub use legacy_stats::LegacyStatsPlan;
 pub(crate) use legacy_stats::LegacyStatsRule;
-pub(crate) use legacy_stats::ZonedLegacyStatsRule;
 pub use list_pack::ListPack;
 pub use list_pack::ListPackData;
 pub use list_pack::ListPackPlan;
@@ -55,4 +55,5 @@ pub use take::TakePlan;
 pub(crate) use zoned::ExpressionZonedRule;
 pub use zoned::Zoned;
 pub use zoned::ZonedData;
+pub(crate) use zoned::ZonedDataRule;
 pub use zoned::ZonedPlan;