diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 71dfb22ffc1..26620ea2522 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/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>>, + /// 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>, /// Shared V1 full-file natural splits keyed by file path. pub natural_splits: Arc>>, /// 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, 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 { 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: &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, diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 5a86abee50c..644fea4e71c 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/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>>, /// Shared V2 segment sources, keyed by path and alive only while a scan reads the file. segment_sources: Arc>>, + scan_caches: Arc>, /// Shared full-file natural splits keyed by path. natural_splits: Arc>>, /// 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) diff --git a/vortex-layout/src/plan/physical/optimizer/mod.rs b/vortex-layout/src/plan/physical/optimizer/mod.rs index 9b928c80daf..073bc712d36 100644 --- a/vortex-layout/src/plan/physical/optimizer/mod.rs +++ b/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 = static LEGACY_STATS_RULE: PlanReduceRuleAdapter = PlanReduceRuleAdapter::new(LegacyStatsRule); +static ZONED_DATA_RULE: PlanReduceRuleAdapter = + 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 = @@ -52,9 +57,10 @@ static EXPRESSION_PACK_RULE: PlanParentReduceRuleAdapter = PlanParentReduceRuleAdapter::new(ExpressionZonedRule); - -static ZONED_LEGACY_STATS_RULE: PlanParentReduceRuleAdapter = - PlanParentReduceRuleAdapter::new(ZonedLegacyStatsRule); +static EXPRESSION_LEGACY_STATS_RULE: PlanParentReduceRuleAdapter< + LegacyStats, + ExpressionLegacyStatsRule, +> = PlanParentReduceRuleAdapter::new(ExpressionLegacyStatsRule); static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[ &PlanParentReduceRuleAdapter::new(EvalEvalRule), @@ -62,7 +68,7 @@ static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[ &EXPRESSION_TAKE_RULE, &EXPRESSION_PACK_RULE, &EXPRESSION_ZONED_RULE, - &ZONED_LEGACY_STATS_RULE, + &EXPRESSION_LEGACY_STATS_RULE, ]); /// Attempts a static rewrite for `plan`. diff --git a/vortex-layout/src/plan/physical/plans/legacy_stats.rs b/vortex-layout/src/plan/physical/plans/legacy_stats.rs index 7eb6bec85ff..4abd2932509 100644 --- a/vortex-layout/src/plan/physical/plans/legacy_stats.rs +++ b/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 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 for ZonedLegacyStatsRule { - type Parent = Zoned; +impl PlanParentReduceRule for ExpressionLegacyStatsRule { + type Parent = Eval; fn reduce_parent( &self, child: &LegacyStatsPlan, - parent: &ZonedPlan, - child_idx: usize, + parent: &Plan, + _child_idx: usize, ) -> VortexResult> { - 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(), + )) } } diff --git a/vortex-layout/src/plan/physical/plans/mod.rs b/vortex-layout/src/plan/physical/plans/mod.rs index 8dba359d635..3d37390cfcb 100644 --- a/vortex-layout/src/plan/physical/plans/mod.rs +++ b/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; diff --git a/vortex-layout/src/plan/physical/plans/zoned.rs b/vortex-layout/src/plan/physical/plans/zoned.rs index d1a17702ebf..8b1729e953e 100644 --- a/vortex-layout/src/plan/physical/plans/zoned.rs +++ b/vortex-layout/src/plan/physical/plans/zoned.rs @@ -46,6 +46,7 @@ use crate::plan::physical::operators::OnceSource; use crate::plan::physical::operators::UNBOUNDED; use crate::plan::physical::operators::WholeSource; use crate::plan::physical::optimizer::PlanParentReduceRule; +use crate::plan::physical::optimizer::PlanReduceRule; use crate::plan::physical::pipeline::Chain; use crate::plan::physical::pipeline::Compiler; @@ -151,7 +152,8 @@ pub struct ZonedData { /// Its zone table has the aggregate schema of `vortex.zoned` layouts; legacy `vortex.stats` /// layouts reach it through [`LegacyStats`](crate::plan::physical::LegacyStats), which converts /// their tables. An expression containing abstract statistic functions can rewrite it into a -/// pruning plan that retains only the zone-statistics child. +/// pruning plan that retains only the zone-statistics child; a plan that reads the data is +/// optimized to its data child alone. #[derive(Clone, Debug)] pub struct Zoned; @@ -239,6 +241,15 @@ impl ZonedPlan { self.child_required(index) } + /// Whether an earlier execution already proved this pruning plan's expression over its zones, + /// so compiling it emits the kept result without reading the zone table. + pub(crate) fn is_proved(&self, session: &VortexSession) -> VortexResult { + match (self.pruned_zones(), self.pruning_expression()) { + (Some(pruned), Some(expression)) => Ok(pruned.get(expression, session)?.is_some()), + _ => Ok(false), + } + } + /// Returns whether this plan represents a zone-backed pruning proof. pub fn is_pruning(&self) -> bool { self.data().pruning.is_some() @@ -453,29 +464,46 @@ impl PlanParentReduceRule for ExpressionZonedRule { parent: &Plan, _child_idx: usize, ) -> VortexResult> { - let mut contains_stat = false; - let mut contains_root = false; - parent.expression().clone().transform_down(|expression| { - if expression - .as_scalar() - .is_some_and(|scalar_fn| scalar_fn.is::()) - { - contains_stat = true; - return Ok(Transformed { - value: expression, - order: TraversalOrder::Skip, - changed: false, - }); - } - contains_root |= expression.is_root(); - Ok(Transformed::no(expression)) - })?; - if !parent.dtype().is_boolean() || !contains_stat || contains_root { + if !is_zone_proof(parent)? { return Ok(None); } - Ok(child .with_pruning(parent.expression().clone())? .map(Plan::into_plan)) } } + +/// Whether `eval` is a boolean expression over abstract statistics alone, which a zone table can +/// prove without reading the data it summarises. +pub(crate) fn is_zone_proof(eval: &Plan) -> VortexResult { + let mut contains_stat = false; + let mut contains_root = false; + eval.expression().clone().transform_down(|expression| { + if expression + .as_scalar() + .is_some_and(|scalar_fn| scalar_fn.is::()) + { + contains_stat = true; + return Ok(Transformed { + value: expression, + order: TraversalOrder::Skip, + changed: false, + }); + } + contains_root |= expression.is_root(); + Ok(Transformed::no(expression)) + })?; + Ok(eval.dtype().is_boolean() && contains_stat && !contains_root) +} + +/// Replaces a zoned data plan by its data: a scan reads the data through it and never its zone +/// table. Pruning proofs are derived from the filter over the lowered plan before this runs, so +/// an optimized plan carries a zone table only where a proof reads it. +#[derive(Debug)] +pub(crate) struct ZonedDataRule; + +impl PlanReduceRule for ZonedDataRule { + fn reduce(&self, plan: &Plan) -> VortexResult> { + plan.data_plan() + } +} diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index 7ee6647bd47..6698b6a0fe3 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -58,6 +58,9 @@ use crate::layouts::zoned::MAX_IS_TRUNCATED; use crate::layouts::zoned::ZonedLayout; use crate::layouts::zoned::legacy_stats_to_aggregates; use crate::layouts::zoned::zone_map::ZoneMap; +use crate::plan::physical::Filter; +use crate::plan::physical::LegacyStats; +use crate::plan::physical::pipeline::Scan; use crate::segments::SegmentId; use crate::test::SESSION; @@ -1306,14 +1309,37 @@ fn stats_expression_rewrites_to_zoned_pruning_plan() -> VortexResult<()> { let mixed = mixed .as_opt::() .ok_or_else(|| vortex_err!("expression with a data reference was pushed into zones"))?; - let mixed_child = mixed.child_plan()?; - assert!(mixed_child.is::()); - assert!( - !mixed_child - .as_opt::() - .ok_or_else(|| vortex_err!("mixed expression child is not zoned"))? - .is_pruning() - ); + assert!(mixed.child_plan()?.is::()); + Ok(()) +} + +/// A plan that reads a zoned column's data is optimized to the data alone: the zone table is +/// neither a child of the plan nor among the segments it reads. +#[test] +fn data_expression_strips_zoned_to_its_data() -> VortexResult<()> { + let dtype = primitive(PType::I32, Nullability::NonNullable); + let (max, zones_dtype) = max_zones_dtype(&dtype)?; + let zone_len = NonZeroUsize::new(3).ok_or_else(|| vortex_err!("zone length is zero"))?; + let layout = ZonedLayout::try_new( + flat(5, dtype.clone(), 0), + flat(2, zones_dtype, 1), + zone_len, + vec![max].into(), + )? + .into_layout(); + let source = make_plan(layout)?; + + let plan = EvalPlan::try_new(gt(root(), lit(4_i32)).bind(&dtype)?, source.clone())?.into_plan(); + insta::assert_snapshot!(optimize(plan)?.display_tree(), @" + root: vortex.plan.eval(bool, rows=5) expr=($ > 4i32) + child: vortex.plan.filter(i32, rows=5) + child: vortex.plan.segment_scan(i32, rows=5) + "); + + let identity = optimize(source)?; + assert!(identity.is::()); + let segments = Scan::segments(&SESSION, &identity, 0..5, &Mask::new_true(5))?; + assert_eq!(segments, vec![SegmentId::from(0)]); Ok(()) } @@ -1483,23 +1509,8 @@ fn legacy_stats_rewrite_to_zoned_over_converted_zones() -> VortexResult<()> { let (max, zones_dtype) = max_zones_dtype(&dtype)?; let legacy = legacy_max_plan(&dtype)?; - let optimized = optimize(legacy.clone())?; - let zoned = optimized - .as_opt::() - .ok_or_else(|| vortex_err!("legacy stats did not become zoned"))?; - assert!(!zoned.is_pruning()); - assert_eq!(zoned.zone_len(), 3); - assert_eq!( - zoned - .aggregate_fns() - .iter() - .map(ToString::to_string) - .collect::>(), - vec![max.to_string()] - ); - let zones = zoned.zones_plan()?; - assert!(zones.is::()); - assert_eq!(zones.dtype(), &zones_dtype); + // Read as data, the legacy stats are dropped along with the zoned plan they convert to. + assert!(optimize(legacy.clone())?.is::()); let falsifier = gt(root(), lit(5_i32)) .bind(&dtype)? @@ -1510,7 +1521,18 @@ fn legacy_stats_rewrite_to_zoned_over_converted_zones() -> VortexResult<()> { .as_opt::() .ok_or_else(|| vortex_err!("legacy stats predicate did not become a pruning plan"))?; assert!(pruning.is_pruning()); - assert!(pruning.zones_plan()?.is::()); + assert_eq!(pruning.zone_len(), 3); + assert_eq!( + pruning + .aggregate_fns() + .iter() + .map(ToString::to_string) + .collect::>(), + vec![max.to_string()] + ); + let zones = pruning.zones_plan()?; + assert!(zones.is::()); + assert_eq!(zones.dtype(), &zones_dtype); Ok(()) } @@ -1549,43 +1571,52 @@ fn converted_legacy_zones_prune() -> VortexResult<()> { Ok(()) } +/// A zoned plan over legacy stats read as data keeps neither zone table. #[test] -fn legacy_stats_under_zoned_are_dropped() -> VortexResult<()> { +fn zoned_over_legacy_stats_read_as_data_keeps_only_the_data() -> VortexResult<()> { let dtype = primitive(PType::I32, Nullability::NonNullable); let (max, zones_dtype) = max_zones_dtype(&dtype)?; let zones = make_plan(flat(2, zones_dtype, 2))?; - let plan = - ZonedPlan::new(legacy_max_plan(&dtype)?, zones.clone(), 3, vec![max].into()).into_plan(); + let legacy = legacy_max_plan(&dtype)?; + let data = legacy + .as_opt::() + .ok_or_else(|| vortex_err!("legacy plan is not legacy stats"))? + .data_plan()?; + let plan = ZonedPlan::new(legacy, zones, 3, vec![max].into()).into_plan(); - let optimized = optimize(plan)?; - let zoned = optimized - .as_opt::() - .ok_or_else(|| vortex_err!("zoned plan was rewritten away"))?; - assert!(zoned.data_plan()?.is_some_and(|data| data.is::())); - assert!(PlanRef::ptr_eq(&zoned.zones_plan()?, &zones)); + assert!(PlanRef::ptr_eq(&optimize(plan)?, &data)); Ok(()) } +/// Legacy stats over a zoned plan are dropped: read as data only the data remains, and a zone +/// proof is proved over the aggregate zones, never the legacy table. #[test] fn legacy_stats_over_zoned_are_dropped() -> VortexResult<()> { let dtype = primitive(PType::I32, Nullability::NonNullable); let (max, zones_dtype) = max_zones_dtype(&dtype)?; - let zoned = ZonedPlan::new( - make_plan(flat(5, dtype.clone(), 0))?, - make_plan(flat(2, zones_dtype, 2))?, - 3, - vec![max].into(), - ) - .into_plan(); + let data = make_plan(flat(5, dtype.clone(), 0))?; + let zones = make_plan(flat(2, zones_dtype, 2))?; + let zoned = ZonedPlan::new(data.clone(), zones.clone(), 3, vec![max].into()).into_plan(); let legacy = LegacyStatsPlan::new( - zoned.clone(), + zoned, make_plan(flat(2, legacy_max_zones_dtype(&dtype), 1))?, 3, vec![Stat::Max].into(), ) .into_plan(); - assert!(PlanRef::ptr_eq(&optimize(legacy)?, &zoned)); + assert!(PlanRef::ptr_eq(&optimize(legacy.clone())?, &data)); + + let falsifier = gt(root(), lit(5_i32)) + .bind(&dtype)? + .falsify(&SESSION)? + .ok_or_else(|| vortex_err!("filter has no falsifier"))?; + let pruning = optimize(EvalPlan::try_new(falsifier, legacy)?.into_plan())?; + let pruning = pruning + .as_opt::() + .ok_or_else(|| vortex_err!("predicate over legacy stats did not become a pruning plan"))?; + assert!(pruning.is_pruning()); + assert!(PlanRef::ptr_eq(&pruning.zones_plan()?, &zones)); Ok(()) } diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs new file mode 100644 index 00000000000..6808c707c48 --- /dev/null +++ b/vortex-layout/src/scan/execution/chunk.rs @@ -0,0 +1,118 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Cuts a scan's rows into chunks and plans each one. + +use std::ops::Range; +use std::sync::Arc; +use std::vec::IntoIter; + +use vortex_error::VortexResult; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_mask::Mask; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::PlannerOutput; +use vortex_scan::planning::planner::State; +use vortex_scan::planning::planner::WorkScope; +use vortex_scan::selection::Selection; + +use super::PreparedScan; +use super::SegmentLocation; +use super::selection::plan_selection; + +/// The row grid every chunk boundary lies on. +pub(crate) const ROW_BLOCK_SIZE: usize = 8192; +/// The rows of a projection chunk. +pub(crate) const MAX_ROWS_PER_CHUNK: usize = 8 * ROW_BLOCK_SIZE; +/// The rows of a filtered chunk: large, so a conjunct's reads coalesce and its mask is cheap. +pub(crate) const FILTERED_ROWS_PER_CHUNK: usize = 16 * MAX_ROWS_PER_CHUNK; +const TARGET_BYTES_PER_CHUNK: u64 = 8 * 1024 * 1024; + +/// Limit estimated bytes per chunk without changing the writer-aligned row granularity. +pub(crate) fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { + if row_count <= MAX_ROWS_PER_CHUNK as u64 { + return MAX_ROWS_PER_CHUNK; + } + let bytes: u64 = locations + .iter() + .map(|location| u64::from(location.length)) + .sum(); + let bytes_per_row = bytes.div_ceil(row_count).max(1); + let blocks = TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64; + blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE +} + +/// Fixed boundaries in the file's row domain, clipped to the requested range. +pub(crate) fn row_chunks(rows: Range, rows_per_chunk: usize) -> Vec> { + if rows.is_empty() { + return Vec::new(); + } + let first = rows.start / rows_per_chunk as u64 * rows_per_chunk as u64; + (first..rows.end) + .step_by(rows_per_chunk) + .map(|start| { + start.max(rows.start)..start.saturating_add(rows_per_chunk as u64).min(rows.end) + }) + .collect() +} + +/// Emits one planner per chunk of a scan's rows, in row order, each spawned as independent so +/// its owner may run it beside the others. A chunk none of whose rows is selected and kept is +/// never planned. +pub(in crate::scan) struct ChunkPlanner { + scan: Arc, + rows: Range, + chunks: IntoIter>, + selection: Selection, + /// The rows zone pruning kept, over `rows`, or `None` when nothing was proved. + kept: Option, +} + +impl ChunkPlanner { + pub(super) fn new( + scan: Arc, + rows: Range, + selection: Selection, + kept: Option, + ) -> Self { + let chunks = row_chunks(rows.clone(), scan.chunk_rows).into_iter(); + Self { + scan, + rows, + chunks, + selection, + kept, + } + } +} + +impl IoConsumer for ChunkPlanner { + fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {} +} + +impl Planner for ChunkPlanner { + fn state(&self) -> State { + if self.chunks.as_slice().is_empty() { + State::Done + } else { + State::NeedsCompute + } + } + + fn compute(&mut self) -> VortexResult { + let kept = self.kept.as_ref().map(|kept| (kept, self.rows.start)); + for chunk in self.chunks.by_ref() { + let scope = WorkScope { + file_ordinal: 0, + rows: chunk, + }; + if let Some((scope, planner)) = plan_selection(&self.scan, scope, &self.selection, kept) + { + return Ok(PlannerOutput::Spawn(scope, planner)); + } + } + Ok(PlannerOutput::Done) + } +} diff --git a/vortex-layout/src/scan/execution/filter.rs b/vortex-layout/src/scan/execution/filter.rs index 2ce5f588f02..48718307577 100644 --- a/vortex-layout/src/scan/execution/filter.rs +++ b/vortex-layout/src/scan/execution/filter.rs @@ -1,9 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; use std::sync::Arc; use bit_vec::BitVec; +use rustc_hash::FxHashSet; use vortex_array::VortexSessionExecute; use vortex_array::builtins::ArrayBuiltins; use vortex_array::expr::BoundExpression; @@ -11,6 +13,7 @@ use vortex_array::expr::transform::find_between; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; +use vortex_io::request::IoBatch; use vortex_io::request::IoConsumer; use vortex_io::request::IoIntent; use vortex_io::request::IoRequest; @@ -22,9 +25,11 @@ use vortex_scan::planning::planner::PlannerOutput; use vortex_scan::planning::planner::State; use vortex_session::VortexSession; -use super::ScanPlans; +use super::PreparedScan; +use super::SegmentLocation; use super::SelectedRows; use super::adaptor::AdaptorStep; +use super::adaptor::Piece; use super::adaptor::ScanAdaptor; use super::projection::ProjectionPlanner; use crate::plan::physical::FilterPlan; @@ -36,11 +41,13 @@ use crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Split; use crate::plan::physical::plan_row_idx_expression; use crate::scan::filter::FilterExpr; +use crate::segments::SegmentId; -/// Plans and adaptive conjunct ordering shared by all splits of a scan. +/// Plans and adaptive conjunct ordering shared by every chunk of a scan. pub(in crate::scan) struct FilterPlans { expression: FilterExpr, predicates: Vec, + /// A metadata-only proof per conjunct, where the conjunct has one. pruning: Vec>, } @@ -81,6 +88,11 @@ impl FilterPlans { pruning, }) } + + /// The pruning proofs, run once over a scan's rows before its chunks are planned. + pub(in crate::scan) fn proofs(&self) -> impl Iterator { + self.pruning.iter().flatten() + } } /// A pruning proof must not fetch data when its requested statistics are unavailable. @@ -102,37 +114,87 @@ fn metadata_only(plan: &PlanRef) -> VortexResult { Ok(true) } -/// Proves zone exclusions, evaluates conjuncts over survivors, then hands the mask to projection. +/// Evaluates a chunk's conjuncts over its selected rows, then hands the mask to projection. +/// +/// The first compute starts the reads of every conjunct's segments for the selected rows and +/// announces the projection's, so the conjuncts run back to back and the projection's reads +/// coalesce with them. Conjuncts run in the order the shared [`FilterExpr`] picks. pub(super) struct FilterPlanner { - plans: ScanPlans, + scan: Arc, selected: Option, filter: Arc, remaining: BitVec, - prune_index: usize, - active: Option<(usize, bool, ScanAdaptor)>, + active: Option<(usize, ScanAdaptor)>, pieces: Vec, announced: bool, + /// Ids for hints, counting down from `u32::MAX` so they never meet an adaptor's fetch ids. + next_hint: u32, +} + +/// The segments `plans` read over `rows` under `mask`, each once, in id order. Compiling finds +/// them without IO. +pub(super) fn segments_read<'a>( + session: &VortexSession, + plans: impl Iterator, + rows: Range, + mask: &Mask, +) -> VortexResult> { + let mut segments = Vec::new(); + for plan in plans { + segments.extend(Scan::segments(session, plan, rows.clone(), mask)?); + } + segments.sort_unstable(); + segments.dedup(); + Ok(segments) } impl FilterPlanner { - pub(super) fn new(plans: ScanPlans, selected: SelectedRows) -> Self { - let filter = Arc::clone(plans.filter.as_ref().vortex_expect("filter plans exist")); + pub(super) fn new(scan: Arc, selected: SelectedRows) -> Self { + let filter = Arc::clone(scan.filter.as_ref().vortex_expect("filter plans exist")); Self { remaining: BitVec::from_elem(filter.predicates.len(), true), - plans, + scan, selected: Some(selected), filter, - prune_index: 0, active: None, pieces: Vec::new(), announced: false, + next_hint: u32::MAX, } } } +/// Publishes a hint for each of `segments` not in `seen`, and adds it there. +pub(super) fn hint( + locations: &[SegmentLocation], + next_hint: &mut u32, + intent: IoIntent, + segments: &[SegmentId], + seen: &mut FxHashSet, + batch: &mut IoBatch, +) { + for &segment in segments { + if !seen.insert(segment) { + continue; + } + batch.push(IoRequest { + intent, + request: IoRequestId(*next_hint), + target: locations[*segment as usize].target(), + }); + *next_hint -= 1; + } +} + +/// Executes a scan's boolean piece to a mask, nulls as false. +pub(super) fn piece_mask(session: &VortexSession, piece: Piece) -> VortexResult { + let mut ctx = session.create_execution_ctx(); + piece.array.fill_null(false)?.execute::(&mut ctx) +} + impl IoConsumer for FilterPlanner { fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { - if let Some((_, _, adaptor)) = &mut self.active { + if let Some((_, adaptor)) = &mut self.active { adaptor.set_io_result(request, result); } } @@ -144,7 +206,7 @@ impl Planner for FilterPlanner { return State::Done; } match &self.active { - Some((_, _, adaptor)) if adaptor.state() != State::Done => adaptor.state(), + Some((_, adaptor)) if adaptor.state() != State::Done => adaptor.state(), _ => State::NeedsCompute, } } @@ -155,42 +217,50 @@ impl Planner for FilterPlanner { }; if !self.announced { self.announced = true; + let session = &self.scan.session; + let rows = selected.scope.rows.clone(); + let predicates = segments_read( + session, + self.filter.predicates.iter(), + rows.clone(), + &selected.mask, + )?; + let projection = segments_read( + session, + std::iter::once(&self.scan.projection), + rows, + &selected.mask, + )?; let mut batch = Vec::new(); - for plan in self - .filter - .pruning - .iter() - .flatten() - .chain(self.filter.predicates.iter()) - .chain(std::iter::once(&self.plans.projection)) - { - let segments = Scan::segments( - &self.plans.session, - plan, - selected.scope.rows.clone(), - &selected.mask, - )?; - for segment in segments { - batch.push(IoRequest { - intent: IoIntent::Announce, - request: IoRequestId(u32::MAX - u32::try_from(batch.len())?), - target: self.plans.locations[*segment as usize].target(), - }); - } - } + let mut seen = FxHashSet::default(); + let locations = &self.scan.locations; + hint( + locations, + &mut self.next_hint, + IoIntent::Prefetch, + &predicates, + &mut seen, + &mut batch, + ); + hint( + locations, + &mut self.next_hint, + IoIntent::Announce, + &projection, + &mut seen, + &mut batch, + ); if !batch.is_empty() { return Ok(PlannerOutput::NeedsIO(batch)); } } - if let Some((index, pruning, adaptor)) = &mut self.active { + if let Some((index, adaptor)) = &mut self.active { if adaptor.state() != State::Done { return Ok(match adaptor.compute()? { AdaptorStep::Yield => PlannerOutput::Continue, AdaptorStep::NeedsIO(batch) => PlannerOutput::NeedsIO(batch), AdaptorStep::Piece(piece) => { - let mut ctx = self.plans.session.create_execution_ctx(); - let mask = piece.array.fill_null(false)?.execute::(&mut ctx)?; - self.pieces.push(mask); + self.pieces.push(piece_mask(&self.scan.session, piece)?); PlannerOutput::Continue } }); @@ -198,22 +268,18 @@ impl Planner for FilterPlanner { let input = selected.mask.true_count(); let result = Mask::concat(self.pieces.iter())?; vortex_ensure_eq!(result.len(), input); - let result = if *pruning { !result } else { result }; selected.mask = selected.mask.intersect_by_rank(&result); tracing::debug!( row_start = selected.scope.rows.start, row_end = selected.scope.rows.end, input_rows = input, output_rows = selected.mask.true_count(), - pruning = *pruning, conjunct = *index, "v2 filter conjunct evaluated" ); - if !*pruning { - self.filter - .expression - .report_selectivity(*index, selected.mask.true_count() as f64 / input as f64); - } + self.filter + .expression + .report_selectivity(*index, selected.mask.true_count() as f64 / input as f64); self.active = None; self.pieces.clear(); } @@ -221,44 +287,28 @@ impl Planner for FilterPlanner { self.selected = None; return Ok(PlannerOutput::Done); } - - let mut next = None; - while self.prune_index < self.filter.pruning.len() { - let index = self.prune_index; - self.prune_index += 1; - if let Some(plan) = &self.filter.pruning[index] { - next = Some((index, true, plan.clone())); - break; - } - } - if next.is_none() - && let Some(index) = self.filter.expression.next_conjunct(&self.remaining) - { + if let Some(index) = self.filter.expression.next_conjunct(&self.remaining) { self.remaining.set(index, false); - next = Some((index, false, self.filter.predicates[index].clone())); - } - if let Some((index, pruning, plan)) = next { let scan = Scan::try_new( - self.plans.session.clone(), - plan, + self.scan.session.clone(), + self.filter.predicates[index].clone(), vec![Split { rows: selected.scope.rows.clone(), mask: selected.mask.clone(), }], )? - .with_row_offset(self.plans.row_offset) - .with_value_cache(self.plans.values.clone()); + .with_row_offset(self.scan.row_offset) + .with_value_cache(self.scan.values.clone()); self.active = Some(( index, - pruning, - ScanAdaptor::new(scan, Arc::clone(&self.plans.locations)), + ScanAdaptor::new(scan, Arc::clone(&self.scan.locations)), )); return Ok(PlannerOutput::Continue); } let selected = self.selected.take().vortex_expect("selected rows exist"); Ok(PlannerOutput::Planner( selected.scope.clone(), - Box::new(ProjectionPlanner::new(self.plans.clone(), selected)), + Box::new(ProjectionPlanner::new(Arc::clone(&self.scan), selected)), )) } } diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 53cc685bccb..e371fdb539c 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -1,32 +1,49 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Scan execution over physical plans through the planner I/O and compute protocol. +//! The planners of a V2 scan, and the prepared plans they run. //! -//! A `ProjectionPlanner` emits bounded `ProjectionMorsel`s for a work chunk's selected rows. Each -//! morsel runs the physical projection plan through -//! the pipeline [`Scan`](crate::plan::physical::pipeline::Scan). Every segment read is published -//! through a `ScanAdaptor`, so the planning driver can advance other work while reads wait. +//! A [`PreparedScan`] holds a scan's plans over one file and makes its planners. The planners +//! form a chain, each emitting the next through the planning protocol: a [`PrunePlanner`] +//! proves the filter's zone maps over the scan's rows and emits one [`ChunkPlanner`] for the rows +//! they keep; the chunk planner cuts those rows into chunks and emits a [`FilterPlanner`] or a +//! [`ProjectionPlanner`] per chunk; a filter planner runs its chunk's conjuncts and emits a +//! projection planner for the rows that survive; a projection planner emits one morsel per piece +//! of its rows. A stage with nothing to do is never made. mod adaptor; +mod chunk; mod filter; mod projection; +mod prune; mod selection; +use std::ops::Range; use std::sync::Arc; use vortex_buffer::Alignment; use vortex_io::request::IoTarget; use vortex_mask::Mask; +use vortex_scan::planning::planner::Planner; use vortex_scan::planning::planner::WorkScope; +use vortex_scan::selection::Selection; use vortex_session::VortexSession; +pub(in crate::scan) use self::chunk::ChunkPlanner; +pub(crate) use self::chunk::FILTERED_ROWS_PER_CHUNK; +pub(crate) use self::chunk::MAX_ROWS_PER_CHUNK; +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; +use self::filter::FilterPlanner; pub(super) use self::filter::FilterPlans; -pub(super) use self::selection::plan_selection; +use self::projection::ProjectionPlanner; +use self::prune::PrunePlanner; use crate::plan::physical::PlanRef; use crate::plan::physical::pipeline::ValueCache; -/// Where a segment's bytes live in its source. +/// Where a segment's bytes live in the file. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct SegmentLocation { /// Byte offset of the segment from the start of the source. @@ -48,26 +65,77 @@ impl SegmentLocation { } } -/// What every split of a scan shares: the projection plan and where its segments live. -#[derive(Clone)] -pub(super) struct ScanPlans { +/// A scan's plans over one file, built once from the file's footer and read by every planner of +/// the scan. Nothing here touches a segment, and no planner plans again. +pub(crate) struct PreparedScan { /// The session used for decoding and expression evaluation. pub(super) session: VortexSession, /// The location of every segment, indexed by segment id. pub(super) locations: Arc<[SegmentLocation]>, /// The optimized projection plan, over the scan's root row domain. pub(super) projection: PlanRef, + /// The filter's conjuncts, their plans and their proofs. pub(super) filter: Option>, /// The global row index of the root row domain's first row. pub(super) row_offset: u64, - /// Whole values, such as dictionary values, shared by the splits of this scan. Their plans - /// always read the same complete row domain, independently of a split's selection. + /// Whole values, such as dictionary values, shared by the chunks of this scan. Their plans + /// always read the same complete row domain, independently of a chunk's selection. pub(super) values: ValueCache, + /// The rows of one chunk, the unit of work a chunk planner emits. + pub(super) chunk_rows: usize, } -/// The rows of a split selected for projection. -struct SelectedRows { - /// The split the rows belong to. +impl PreparedScan { + /// The first planner a scan of `rows` needs. Every later planner is emitted by the one before + /// it: the zone proofs run first when the filter has any, then one planner per chunk of the + /// rows they keep. + pub(crate) fn planner( + self: &Arc, + rows: Range, + selection: Selection, + ) -> Box { + let proofs = self + .filter + .as_ref() + .is_some_and(|filter| filter.proofs().next().is_some()); + if proofs && !rows.is_empty() { + let scan = Arc::clone(self); + let chunks = rows.clone(); + Box::new(PrunePlanner::new( + Arc::clone(self), + rows, + Box::new(move |kept| Box::new(scan.chunk_planner(chunks, selection, Some(kept)))), + )) + } else { + Box::new(self.chunk_planner(rows, selection, None)) + } + } + + /// The planner cutting `rows` into chunks, planning only the chunks `selection` and `kept` + /// leave something of. + pub(super) fn chunk_planner( + self: &Arc, + rows: Range, + selection: Selection, + kept: Option, + ) -> ChunkPlanner { + ChunkPlanner::new(Arc::clone(self), rows, selection, kept) + } + + /// The planner of one chunk's selected rows: its conjuncts when the scan has a filter, else + /// its projection. + pub(super) fn split_planner(self: &Arc, selected: SelectedRows) -> Box { + if self.filter.is_some() { + Box::new(FilterPlanner::new(Arc::clone(self), selected)) + } else { + Box::new(ProjectionPlanner::new(Arc::clone(self), selected)) + } + } +} + +/// The rows a planner works on: a chunk, and which of its rows are selected. +pub(super) struct SelectedRows { + /// The chunk the rows belong to. scope: WorkScope, /// Which rows of `scope.rows` are selected. mask: Mask, diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index 00c0a0614f6..141b7377c26 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; use std::vec::IntoIter; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_io::request::IoConsumer; @@ -23,7 +23,7 @@ use vortex_scan::planning::planner::WorkScope; use crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Split; use crate::plan::physical::selection::selected; -use crate::scan::execution::ScanPlans; +use crate::scan::execution::PreparedScan; use crate::scan::execution::SelectedRows; use crate::scan::execution::adaptor::AdaptorStep; use crate::scan::execution::adaptor::Piece; @@ -34,13 +34,13 @@ const MAX_OUTPUT_ROWS: usize = 65_536; /// Keeps sparse selections together and subdivides dense selections to avoid large copies when /// converting chunked projection results to Arrow. pub(super) struct ProjectionPlanner { - plans: ScanPlans, + scan: Arc, selected: IntoIter, } impl ProjectionPlanner { /// Creates a planner for `selected`. - pub(super) fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + pub(super) fn new(scan: Arc, selected: SelectedRows) -> Self { let selected = if selected.mask.true_count() <= MAX_OUTPUT_ROWS { vec![selected] } else { @@ -64,7 +64,7 @@ impl ProjectionPlanner { .collect() }; Self { - plans, + scan, selected: selected.into_iter(), } } @@ -92,7 +92,7 @@ impl Planner for ProjectionPlanner { } Ok(PlannerOutput::Morsel( selected.scope.clone(), - Box::new(ProjectionMorsel::new(self.plans.clone(), selected)), + Box::new(ProjectionMorsel::new(Arc::clone(&self.scan), selected)), )) } } @@ -100,7 +100,7 @@ impl Planner for ProjectionPlanner { /// Runs the projection plan over a split's selected rows and returns them as one batch, in row /// order. A selection that projects to no rows finishes without a batch. struct ProjectionMorsel { - plans: Option, + scan: Arc, dtype: DType, selected: SelectedRows, adaptor: Option, @@ -110,10 +110,10 @@ struct ProjectionMorsel { impl ProjectionMorsel { /// Creates a morsel projecting `selected`. - fn new(plans: ScanPlans, selected: SelectedRows) -> Self { + fn new(scan: Arc, selected: SelectedRows) -> Self { Self { - dtype: plans.projection.dtype().clone(), - plans: Some(plans), + dtype: scan.projection.dtype().clone(), + scan, selected, adaptor: None, pieces: Vec::new(), @@ -163,14 +163,10 @@ impl Morsel for ProjectionMorsel { vortex_bail!("ProjectionMorsel: compute called after Done"); } let Some(adaptor) = self.adaptor.as_mut() else { - let plans = self - .plans - .take() - .vortex_expect("an uncompiled projection still has its plans"); let scan = Scan::try_new( - plans.session, + self.scan.session.clone(), selected( - plans.projection, + self.scan.projection.clone(), self.selected.scope.rows.clone(), self.selected.mask.clone(), ), @@ -179,9 +175,9 @@ impl Morsel for ProjectionMorsel { mask: Mask::new_true(self.selected.mask.len()), }], )? - .with_row_offset(plans.row_offset) - .with_value_cache(plans.values); - self.adaptor = Some(ScanAdaptor::new(scan, plans.locations)); + .with_row_offset(self.scan.row_offset) + .with_value_cache(self.scan.values.clone()); + self.adaptor = Some(ScanAdaptor::new(scan, Arc::clone(&self.scan.locations))); return Ok(MorselOutput::Continue); }; if adaptor.state() == State::Done { diff --git a/vortex-layout/src/scan/execution/prune.rs b/vortex-layout/src/scan/execution/prune.rs new file mode 100644 index 00000000000..1fd153f989c --- /dev/null +++ b/vortex-layout/src/scan/execution/prune.rs @@ -0,0 +1,213 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Proves a filter's zone maps once over a scan's rows, before any chunk is planned. + +use std::ops::Range; +use std::sync::Arc; +use std::vec::IntoIter; + +use rustc_hash::FxHashSet; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoIntent; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_mask::Mask; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::PlannerOutput; +use vortex_scan::planning::planner::State; +use vortex_scan::planning::planner::WorkScope; +use vortex_session::VortexSession; + +use super::PreparedScan; +use super::adaptor::AdaptorStep; +use super::adaptor::ScanAdaptor; +use super::filter::hint; +use super::filter::piece_mask; +use super::filter::segments_read; +use crate::plan::physical::PlanRef; +use crate::plan::physical::Zoned; +use crate::plan::physical::pipeline::Scan; +use crate::plan::physical::pipeline::Split; + +/// Whether every zone proof in `plan` has already been proved, so compiling it reads no zones. +fn is_proved(plan: &PlanRef, session: &VortexSession) -> VortexResult { + if let Some(zoned) = plan.as_opt::() + && zoned.is_pruning() + { + return zoned.is_proved(session); + } + for child in plan.children().iter() { + if !is_proved(&child?, session)? { + return Ok(false); + } + } + Ok(true) +} + +/// The planner a prune planner continues into, given the rows the proofs kept. +pub(super) type Next = Box Box + Send>; + +/// Runs the filter's zone-map proofs over the scan's rows and emits the planner it was built +/// to continue into, with the rows no proof excludes. The zone maps are the only reads, +/// started together. Rows every proof excludes end the scan here: nothing is emitted. +pub(super) struct PrunePlanner { + scan: Arc, + rows: Range, + /// What to emit once every proof has run, taken when it is. + next: Option, + /// The rows no proof has excluded yet, over `rows`. + mask: Mask, + proofs: IntoIter, + active: Option, + pieces: Vec, + announced: bool, + /// Ids for hints, counting down from `u32::MAX` so they never meet an adaptor's fetch ids. + next_hint: u32, + done: bool, +} + +impl PrunePlanner { + pub(super) fn new(scan: Arc, rows: Range, next: Next) -> Self { + let proofs: Vec<_> = scan + .filter + .iter() + .flat_map(|filter| filter.proofs().cloned()) + .collect(); + let len = usize::try_from(rows.end - rows.start).unwrap_or(usize::MAX); + Self { + scan, + rows, + next: Some(next), + mask: Mask::new_true(len), + proofs: proofs.into_iter(), + active: None, + pieces: Vec::new(), + announced: false, + next_hint: u32::MAX, + done: false, + } + } +} + +impl IoConsumer for PrunePlanner { + fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { + if let Some(adaptor) = &mut self.active { + adaptor.set_io_result(request, result); + } + } +} + +impl Planner for PrunePlanner { + fn state(&self) -> State { + if self.done { + return State::Done; + } + match &self.active { + Some(adaptor) if adaptor.state() != State::Done => adaptor.state(), + _ => State::NeedsCompute, + } + } + + fn compute(&mut self) -> VortexResult { + if self.done { + return Ok(PlannerOutput::Done); + } + if !self.announced { + self.announced = true; + // A proof whose zones an earlier scan of the file proved reads nothing: a shared + // plan keeps its proved zones, so only unproved proofs need their zone maps. + let mut unproved = Vec::new(); + for proof in self.proofs.as_slice() { + if !is_proved(proof, &self.scan.session)? { + unproved.push(proof); + } + } + let segments = segments_read( + &self.scan.session, + unproved.into_iter(), + self.rows.clone(), + &self.mask, + )?; + let mut batch = Vec::new(); + let mut seen = FxHashSet::default(); + hint( + &self.scan.locations, + &mut self.next_hint, + IoIntent::Prefetch, + &segments, + &mut seen, + &mut batch, + ); + if !batch.is_empty() { + return Ok(PlannerOutput::NeedsIO(batch)); + } + } + if let Some(adaptor) = &mut self.active { + if adaptor.state() != State::Done { + return Ok(match adaptor.compute()? { + AdaptorStep::Yield => PlannerOutput::Continue, + AdaptorStep::NeedsIO(batch) => PlannerOutput::NeedsIO(batch), + AdaptorStep::Piece(piece) => { + self.pieces.push(piece_mask(&self.scan.session, piece)?); + PlannerOutput::Continue + } + }); + } + let input = self.mask.true_count(); + let excluded = Mask::concat(self.pieces.iter())?; + vortex_ensure_eq!(excluded.len(), input); + self.mask = self.mask.intersect_by_rank(&!excluded); + tracing::debug!( + row_start = self.rows.start, + row_end = self.rows.end, + input_rows = input, + output_rows = self.mask.true_count(), + "v2 zone proof evaluated" + ); + self.active = None; + self.pieces.clear(); + } + if self.mask.all_false() { + self.done = true; + return Ok(PlannerOutput::Done); + } + if let Some(proof) = self.proofs.next() { + let scan = Scan::try_new( + self.scan.session.clone(), + proof, + vec![Split { + rows: self.rows.clone(), + mask: self.mask.clone(), + }], + )? + .with_row_offset(self.scan.row_offset) + .with_value_cache(self.scan.values.clone()); + self.active = Some(ScanAdaptor::new(scan, Arc::clone(&self.scan.locations))); + return Ok(PlannerOutput::Continue); + } + // Every proof has run: the planner this one continues into gets the kept rows. + self.done = true; + tracing::debug!( + row_start = self.rows.start, + row_end = self.rows.end, + kept_rows = self.mask.true_count(), + "v2 zone proofs finished" + ); + let next = self + .next + .take() + .ok_or_else(|| vortex_err!("the prune planner has nothing to continue into"))?; + let kept = std::mem::replace(&mut self.mask, Mask::new_true(0)); + Ok(PlannerOutput::Planner( + WorkScope { + file_ordinal: 0, + rows: self.rows.clone(), + }, + next(kept), + )) + } +} diff --git a/vortex-layout/src/scan/execution/selection.rs b/vortex-layout/src/scan/execution/selection.rs index fe751b48887..d4bd8f0f41e 100644 --- a/vortex-layout/src/scan/execution/selection.rs +++ b/vortex-layout/src/scan/execution/selection.rs @@ -1,7 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Resolves known selections before allocating projection planners or their masks. +//! Resolves known selections before allocating a chunk's planner or its masks. + +use std::sync::Arc; use vortex_error::VortexExpect; use vortex_mask::Mask; @@ -9,25 +11,33 @@ use vortex_scan::planning::planner::Planner; use vortex_scan::planning::planner::WorkScope; use vortex_scan::selection::Selection; -use super::ScanPlans; +use super::PreparedScan; use super::SelectedRows; -use super::filter::FilterPlanner; -use super::projection::ProjectionPlanner; -/// Plans only the selected span, retaining one projection planner and IO session per work chunk. -pub(crate) fn plan_selection( - plans: &ScanPlans, +/// Plans only the selected span of a chunk, with one planner and one mask per chunk. +/// +/// `kept` is the mask zone pruning left over the scan's rows, with the row its first bit +/// stands for. A chunk none of whose rows were kept is not planned. +pub(super) fn plan_selection( + scan: &Arc, scope: WorkScope, selection: &Selection, + kept: Option<(&Mask, u64)>, ) -> Option<(WorkScope, Box)> { - let selected = select_rows(scope, selection)?; + let mut selected = select_rows(scope, selection)?; + if let Some((kept, first_row)) = kept { + let offset = usize::try_from(selected.scope.rows.start - first_row) + .vortex_expect("a work chunk fits usize"); + let kept = kept.slice(offset..offset + selected.mask.len()); + let mask = &selected.mask & &kept; + let first = mask.first()?; + let end = mask.last()? + 1; + selected.scope.rows = + selected.scope.rows.start + first as u64..selected.scope.rows.start + end as u64; + selected.mask = mask.slice(first..end); + } let scope = selected.scope.clone(); - let planner: Box = if plans.filter.is_some() { - Box::new(FilterPlanner::new(plans.clone(), selected)) - } else { - Box::new(ProjectionPlanner::new(plans.clone(), selected)) - }; - Some((scope, planner)) + Some((scope, scan.split_planner(selected))) } fn select_rows(mut scope: WorkScope, selection: &Selection) -> Option { diff --git a/vortex-layout/src/scan/v2/cache.rs b/vortex-layout/src/scan/v2/cache.rs index 3daa0975a0e..5a26e28ec97 100644 --- a/vortex-layout/src/scan/v2/cache.rs +++ b/vortex-layout/src/scan/v2/cache.rs @@ -4,6 +4,7 @@ //! State derived from a [`ScanFile`](super::ScanFile) that repeated V2 scans can share. use std::sync::Arc; +use std::sync::Weak; use parking_lot::Mutex; use rustc_hash::FxHashMap; @@ -13,6 +14,7 @@ use vortex_array::expr::ExactBoundExpr; use super::io::SegmentRanges; use crate::plan::physical::PlanRef; use crate::plan::physical::pipeline::ValueCache; +use crate::scan::execution::FilterPlans; /// Reusable state for repeated V2 scans of one file, the V2 counterpart of a cached /// [`LayoutReader`](crate::LayoutReader) tree. @@ -21,7 +23,9 @@ use crate::plan::physical::pipeline::ValueCache; /// over it, indexes every segment's byte range, and decodes the dictionary values it needs. A scan /// that finds an entry here reuses it, and one that misses builds the entry and keeps it for later /// scans. The lowered plan keeps the children it has initialized, each distinct projection keeps -/// its optimized plan, and dictionary values stay keyed by the cached plan that produced them. +/// its optimized plan, each distinct filter keeps its conjunct and zone-proof plans along with the +/// zones those proofs have already excluded, and dictionary values stay keyed by the cached plan +/// that produced them. /// /// 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 @@ -36,6 +40,9 @@ struct Inner { /// them, so a projection must reuse its plan for its values to be found again. projections: Mutex>, ranges: Mutex>, + /// Plans of filters, keyed by the exact filter expression. A proof plan remembers the zones it + /// excluded, so sharing the plans proves a file's zones once for every scan of that filter. + filters: Mutex>>, values: ValueCache, } @@ -80,7 +87,50 @@ impl ScanCache { Arc::clone(self.0.ranges.lock().get_or_insert(ranges)) } + /// The plans of `filter` an earlier scan kept. + pub(super) fn filter_plans(&self, filter: &BoundExpression) -> Option> { + self.0 + .filters + .lock() + .get(&ExactBoundExpr(filter.clone())) + .cloned() + } + + /// Keeps the plans of `filter`, returning the plans already kept if another scan got there + /// first. + pub(super) fn insert_filter_plans( + &self, + filter: BoundExpression, + plans: Arc, + ) -> Arc { + Arc::clone( + self.0 + .filters + .lock() + .entry(ExactBoundExpr(filter)) + .or_insert(plans), + ) + } + pub(super) fn values(&self) -> ValueCache { self.0.values.clone() } + + /// A handle that finds this cache while a scan holds it, without keeping it alive itself. + pub fn downgrade(&self) -> WeakScanCache { + WeakScanCache(Arc::downgrade(&self.0)) + } +} + +/// 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. +#[derive(Clone)] +pub struct WeakScanCache(Weak); + +impl WeakScanCache { + /// The cache, while a scan still holds it. + pub fn upgrade(&self) -> Option { + self.0.upgrade().map(ScanCache) + } } diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index ce565846d7e..2779739b281 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -80,7 +80,8 @@ pub(super) struct SegmentIoSource { ranges: SegmentRanges, reads: Mutex>>, ready: Mutex>, - /// Reads started by prefetches, consumed by the next fetch of their segment. + /// Every read the session has started, by segment, shared by later fetches of the segment + /// until the session is cleared. prefetched: Mutex>, /// Announced segments, requested and never polled. announced: Mutex>>, @@ -146,19 +147,24 @@ impl IoSource for SegmentIoSource { }); continue; } - let mut bytes = match prefetched.remove(&segment) { - Some(read) => read - .map_err(move |err| { - vortex_err!("prefetched read of segment {segment} failed: {err}") - }) - .boxed(), - None => self - .announced - .lock() - .get_mut(&segment) - .and_then(VecDeque::pop_front) - .unwrap_or_else(|| self.segments.request(segment)), - }; + // A segment stays with the session once read, so a later stage or morsel of the + // split that reads it again shares the read instead of starting another once the + // first has completed. + let mut bytes = prefetched + .entry(segment) + .or_insert_with(|| { + self.announced + .lock() + .get_mut(&segment) + .and_then(VecDeque::pop_front) + .unwrap_or_else(|| self.segments.request(segment)) + .map_err(Arc::new) + .boxed() + .shared() + }) + .clone() + .map_err(move |err| vortex_err!("read of segment {segment} failed: {err}")) + .boxed(); let id = request.request; // In-memory and cached segments need no allocated asynchronous task. A pending // read is polled again by FuturesUnordered before the driver can go to sleep. @@ -325,7 +331,7 @@ mod tests { } /// A fetch of a prefetched segment waits on the prefetch's read rather than reading again, - /// and only the fetch is delivered. + /// only the fetch is delivered, and a later fetch of the segment shares the same read. #[tokio::test] async fn fetch_reuses_a_prefetched_read() -> VortexResult<()> { let segments = Arc::new(CountingSegments::default()); @@ -350,6 +356,16 @@ mod tests { let completion = std::future::poll_fn(|cx| io.poll_completion(cx)).await?; assert_eq!(completion.request, IoRequestId(2)); assert!(completion.result.is_ok()); + assert_eq!(segments.0.load(Ordering::Relaxed), 1); + + io.clear(); + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 3)])?; + assert!( + std::future::poll_fn(|cx| io.poll_completion(cx)) + .await? + .result + .is_ok() + ); assert_eq!(segments.0.load(Ordering::Relaxed), 2); Ok(()) } @@ -379,10 +395,9 @@ mod tests { .result .is_ok() ); - assert_eq!( - segments.0.load(Ordering::Relaxed), - if id < 4 { 2 } else { 3 } - ); + // The first fetch takes one announcement and the session keeps its read, so the + // later fetches share it and the other announcement is never consumed. + assert_eq!(segments.0.load(Ordering::Relaxed), 2); } Ok(()) } diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index c6b359fa328..d60f6b76289 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -4,7 +4,7 @@ //! The planner-driven scan executor. //! //! A V2 scan lowers the file's layout once into a physical plan and runs every split of it -//! through the planning protocol: the [`Run`](vortex_scan::planning::driver::Run) driver advances +//! through the planning protocol: the [`Run`] driver advances //! the split's planners and morsels, and every segment read is a protocol request served from the //! file's [`SegmentSource`]. The layout reader is never asked to evaluate anything. //! @@ -15,22 +15,25 @@ //! layout's chunk boundaries. //! Full-file projections use smaller 8,192-row multiples when serialized bytes per row are high. //! -//! A [`ScanFile`] carrying a [`ScanCache`] shares its lowered plan, segment index, and dictionary -//! values with every other scan of that file, like a cached layout reader does for V1 scans. +//! A [`ScanFile`] carrying a [`ScanCache`] shares its lowered plan, filter plans and proved +//! zones, segment index, and dictionary values with every other scan of that file, like a cached +//! layout reader does for V1 scans. mod cache; mod io; -mod split; #[cfg(test)] mod tests; +use std::future; use std::ops::Range; use std::sync::Arc; +use std::task::Poll; use arrow_array::RecordBatch; use arrow_schema::ArrowError; use arrow_schema::Field; use arrow_schema::SchemaRef; +use futures::SinkExt; use futures::Stream; use futures::StreamExt; use futures::TryStreamExt; @@ -43,8 +46,14 @@ use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_io::request::IoService; +use vortex_io::request::IoSource; +use vortex_io::runtime::Handle; use vortex_io::session::RuntimeSessionExt; +use vortex_scan::planning::driver::Progress; +use vortex_scan::planning::driver::Run; +use vortex_scan::planning::planner::Planner; use vortex_scan::planning::planner::WorkScope; use vortex_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; @@ -52,53 +61,23 @@ use vortex_session::VortexSession; use vortex_utils::parallelism::get_available_parallelism; pub use self::cache::ScanCache; +pub use self::cache::WeakScanCache; use crate::LayoutRef; use crate::plan::physical::lower; use crate::plan::physical::optimize; use crate::plan::physical::plan_row_idx_expression; use crate::scan::arrow::to_record_batch; +use crate::scan::execution::FILTERED_ROWS_PER_CHUNK; use crate::scan::execution::FilterPlans; -use crate::scan::execution::ScanPlans; +use crate::scan::execution::MAX_ROWS_PER_CHUNK; +use crate::scan::execution::PreparedScan; +use crate::scan::execution::ROW_BLOCK_SIZE; pub use crate::scan::execution::SegmentLocation; -use crate::scan::execution::plan_selection; +use crate::scan::execution::rows_per_chunk; use crate::scan::v2::io::SegmentScanIo; use crate::scan::v2::io::segment_ranges; -use crate::scan::v2::split::SplitPlan; use crate::segments::SegmentSource; -const ROW_BLOCK_SIZE: usize = 8192; -const MAX_ROWS_PER_CHUNK: usize = 8 * ROW_BLOCK_SIZE; -const FILTERED_ROWS_PER_CHUNK: usize = 16 * MAX_ROWS_PER_CHUNK; -const TARGET_BYTES_PER_CHUNK: u64 = 8 * 1024 * 1024; - -/// Limit estimated bytes per task without changing the writer-aligned row granularity. -fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { - if row_count <= MAX_ROWS_PER_CHUNK as u64 { - return MAX_ROWS_PER_CHUNK; - } - let bytes: u64 = locations - .iter() - .map(|location| u64::from(location.length)) - .sum(); - let bytes_per_row = bytes.div_ceil(row_count).max(1); - let blocks = TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64; - blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE -} - -/// Fixed boundaries in the file's row domain, clipped to the requested range. -fn row_chunks(rows: Range, rows_per_chunk: usize) -> Vec> { - if rows.is_empty() { - return Vec::new(); - } - let first = rows.start / rows_per_chunk as u64 * rows_per_chunk as u64; - (first..rows.end) - .step_by(rows_per_chunk) - .map(|start| { - start.max(rows.start)..start.saturating_add(rows_per_chunk as u64).min(rows.end) - }) - .collect() -} - /// The rows a byte range of the file stands for when an engine splits the file by bytes before /// opening it. Consecutive byte ranges give consecutive row ranges that tile the file, with /// boundaries on the grid's 8,192-row blocks so the executor's chunks stay block aligned. @@ -226,7 +205,7 @@ impl ScanBuilder { self } - /// Sets the number of active splits per worker thread. Must be greater than zero. + /// Sets the number of chunks run at once per worker thread. Must be greater than zero. pub fn with_concurrency(mut self, concurrency: usize) -> Self { assert!(concurrency > 0); self.concurrency = concurrency; @@ -259,7 +238,7 @@ impl ScanBuilder { let handle = self.session.handle(); Ok(stream::once(async move { handle - .spawn_cpu(move || PreparedScan::try_new(self)?.execute_stream()) + .spawn_cpu(move || ScanRun::try_new(self)?.execute_stream()) .await }) .try_flatten() @@ -267,23 +246,25 @@ impl ScanBuilder { } } -/// A scan planned over its file, ready to run its splits. -struct PreparedScan { +/// A scan prepared over its file, ready to run. +struct ScanRun { session: VortexSession, - plans: ScanPlans, - /// Serves the splits' reads. + scan: Arc, + /// Serves the scan's reads. io: Arc, - /// The scan's splits, over the root row domain, in row order. - splits: Vec>, + /// The rows of the root row domain the scan covers. + rows: Range, selection: Selection, map_fn: Arc VortexResult + Send + Sync>, ordered: bool, concurrency: usize, + /// Held until the scan ends, so scans of the file that start meanwhile share it. + cache: ScanCache, } -impl PreparedScan { - /// Lowers the file's layout, plans the projection, and partitions the selected range into - /// fixed-size row chunks. +impl ScanRun { + /// Lowers the file's layout and plans the filter and the projection, from the file's cache + /// where an earlier scan planned them already. Touches no segment. fn try_new(builder: ScanBuilder) -> VortexResult { let file = builder.file; let cache = file.cache.unwrap_or_default(); @@ -294,21 +275,25 @@ impl PreparedScan { let row_count = root.row_count(); let filter = builder .filter - .map(|filter| FilterPlans::new(filter, &root, &builder.session).map(Arc::new)) + .map(|filter| -> VortexResult<_> { + match cache.filter_plans(&filter) { + Some(plans) => Ok(plans), + None => { + let plans = + Arc::new(FilterPlans::new(filter.clone(), &root, &builder.session)?); + Ok(cache.insert_filter_plans(filter, plans)) + } + } + }) .transpose()?; - // An identity projection needs no rewrites. Keep the lowered children lazy, including - // statistics that a projection-only scan never reads. + // Optimizing drops the zone tables a projection never reads and keeps the chunks of a + // chunked column lazy, so an identity projection is optimized like any other. let identity = builder.projection.is_root() && builder.projection.dtype() == root.dtype(); - let projection = if identity { - root - } else { - match cache.projection(&builder.projection) { - Some(plan) => plan, - None => { - let plan = - optimize(plan_row_idx_expression(builder.projection.clone(), root)?)?; - cache.insert_projection(builder.projection, plan) - } + let projection = match cache.projection(&builder.projection) { + Some(plan) => plan, + None => { + let plan = optimize(plan_row_idx_expression(builder.projection.clone(), root)?)?; + cache.insert_projection(builder.projection, plan) } }; let chunk_rows = if filter.is_some() { @@ -318,7 +303,6 @@ impl PreparedScan { } else { MAX_ROWS_PER_CHUNK }; - let splits = row_chunks(builder.row_range.unwrap_or(0..row_count), chunk_rows); let ranges = match cache.ranges() { Some(ranges) => ranges, None => cache.insert_ranges(segment_ranges(&file.locations)), @@ -326,87 +310,212 @@ impl PreparedScan { let io = Arc::new(SegmentScanIo::new(file.segments, ranges)); Ok(Self { session: builder.session.clone(), - plans: ScanPlans { + scan: Arc::new(PreparedScan { session: builder.session, locations: file.locations, projection, filter, row_offset: builder.row_offset, values: cache.values(), - }, + chunk_rows, + }), io, - splits, + rows: builder.row_range.unwrap_or(0..row_count), selection: builder.selection, map_fn: builder.map_fn, ordered: builder.ordered, concurrency: builder.concurrency, + cache, }) } - /// The splits with selected rows, each ready to be admitted to a planning run. - fn split_plans(&self) -> Vec { - self.splits - .iter() - .filter_map(|rows| { - let (scope, root) = plan_selection( - &self.plans, - WorkScope { - file_ordinal: 0, - rows: rows.clone(), - }, - &self.selection, - )?; - Some(SplitPlan { scope, root }) - }) - .collect() - } - - /// Runs every split as a task on the session's runtime and yields their mapped batches. + /// Runs the scan on planning runs the returned stream owns. /// - /// Like the default executor, up to `concurrency` splits per thread run at once, and their - /// batches are yielded in row order when the scan is ordered. + /// The scan's first planner runs on one run the stream drives as it is polled, each of its + /// CPU steps a [`Handle::spawn_cpu`] task. Every chunk it spawns runs on a run and a session + /// of its own, on one task of the runtime that drives the run to its end and awaits only its + /// reads, up to `concurrency` chunks per worker thread at once; a chunk is started only when + /// the consumer has room for its batches. Every read is a `spawn_io` task of the segment + /// source. The engine thread polling the stream does no scan work itself. Batches of an + /// unordered scan are yielded as morsels finish; an ordered scan yields them in row order + /// once every morsel has. fn execute_stream(self) -> VortexResult>> { - let concurrency = self.concurrency * get_available_parallelism().unwrap_or(1); let handle = self.session.handle(); - let split_plans = self.split_plans(); + let in_flight = self.concurrency * get_available_parallelism().unwrap_or(1); + let root = self.scan.planner(self.rows.clone(), self.selection); + let scope = WorkScope { + file_ordinal: 0, + rows: self.rows.clone(), + }; let io = self.io; let map_fn = self.map_fn; - if !self.ordered - && split_plans.iter().any(|split| { - split.scope.rows.end - split.scope.rows.start > MAX_ROWS_PER_CHUNK as u64 + let (chunk_handle, chunk_io, chunk_map) = + (handle.clone(), Arc::clone(&io), Arc::clone(&map_fn)); + let inline = in_flight <= 1; + let batches = run_stream(handle, root, scope, io.session(), map_fn, inline) + .map_ok(move |yielded| match yielded { + Yielded::Batch(start, batch) => { + stream::once(future::ready(Ok((start, batch)))).boxed() + } + Yielded::Spawn(scope, chunk) => chunk_stream( + &chunk_handle, + chunk, + scope, + chunk_io.session(), + Arc::clone(&chunk_map), + ), }) - { - let streams = stream::iter(split_plans).map(move |split| { - let (send, recv) = mpsc::channel(1); - let task = - handle.spawn(split.run_unordered(Arc::clone(&io), Arc::clone(&map_fn), send)); - let errors = stream::once(task) - .filter_map(|result| std::future::ready(result.err().map(Err))); - stream::select(recv, errors) - }); - return Ok(streams.flatten_unordered(concurrency).boxed()); + .try_flatten_unordered(in_flight); + let cache = self.cache; + let batches = batches.chain(stream::poll_fn(move |_| { + let _ = &cache; + Poll::Ready(None) + })); + if !self.ordered { + return Ok(batches.map_ok(|(_, batch)| batch).boxed()); } - // Spawned as the buffer pulls them, so at most `concurrency` splits read at once. - let tasks = stream::iter(split_plans).map(move |split| { - let io = Arc::clone(&io); - let map_fn = Arc::clone(&map_fn); - handle.spawn(async move { - split - .run(io) - .await? - .into_iter() - .map(|array| map_fn(array)) - .collect::>>() - }) - }); - let batches = if self.ordered { - tasks.buffered(concurrency).boxed() - } else { - tasks.buffer_unordered(concurrency).boxed() - }; - Ok(batches - .map_ok(|batches| stream::iter(batches.into_iter().map(Ok))) - .try_flatten() - .boxed()) + Ok(stream::once(async move { + let mut held: Vec<(u64, A)> = batches.try_collect().await?; + // Morsels finish in whatever order their reads arrive. + held.sort_by_key(|(start, _)| *start); + VortexResult::Ok(stream::iter(held.into_iter().map(|(_, batch)| Ok(batch)))) + }) + .try_flatten() + .boxed()) } } + +/// What the root run's stream yields: a mapped batch with the row it starts at, or a chunk a +/// planner of the run spawned. +enum Yielded { + Batch(u64, A), + Spawn(WorkScope, Box), +} + +/// One CPU step of a run, taken on a `spawn_cpu` task with the run. +enum Step { + Batch(u64, A), + Spawn(WorkScope, Box), + RootDone, + Idle, + Waiting, +} + +/// The state of the root run's stream. +struct RunState { + handle: Handle, + /// Taken while a CPU step runs on its task and put back with its result. + run: Option, + io: Arc, + map_fn: Arc VortexResult + Send + Sync>, + /// Whether a spawned chunk runs on this run, as another root, instead of being yielded. + inline: bool, +} + +/// Runs `root` over `scope` on a planning run of its own, reading through `io`, each CPU step +/// on a `spawn_cpu` task, and yields its mapped batches and, unless `inline`, the chunks it +/// spawns. Inline, a spawned chunk runs on this run beside the others: the run's steps already +/// interleave every chunk's work, and a run per chunk would only add a hand-off per chunk for +/// parallelism the engine provides. +fn run_stream( + handle: Handle, + root: Box, + scope: WorkScope, + io: Arc, + map_fn: Arc VortexResult + Send + Sync>, + inline: bool, +) -> BoxStream<'static, VortexResult>> { + let mut run = Run::new(); + run.admit(root, scope, Arc::clone(&io)); + let state = RunState { + handle, + run: Some(run), + io, + map_fn, + inline, + }; + stream::try_unfold(state, |mut state| async move { + loop { + let run = state + .run + .take() + .ok_or_else(|| vortex_err!("the scan's run is missing"))?; + let map_fn = Arc::clone(&state.map_fn); + let (run, step) = state + .handle + .spawn_cpu(move || { + let mut run = run; + let step = run.advance().and_then(|progress| { + Ok(match progress { + Progress::Batch(batch) => { + Step::Batch(batch.scope.rows.start, map_fn(batch.array)?) + } + Progress::Spawn(scope, child) => Step::Spawn(scope, child), + Progress::RootDone(_) => Step::RootDone, + Progress::Idle => Step::Idle, + Progress::Waiting => Step::Waiting, + }) + }); + (run, step) + }) + .await; + let run = state.run.insert(run); + match step? { + Step::Batch(start, batch) => { + return Ok(Some((Yielded::Batch(start, batch), state))); + } + Step::Spawn(scope, child) if state.inline => { + run.admit(child, scope, Arc::clone(&state.io)); + } + Step::Spawn(scope, child) => { + return Ok(Some((Yielded::Spawn(scope, child), state))); + } + Step::RootDone if run.live_roots() > 0 => {} + Step::RootDone | Step::Idle => return Ok(None), + Step::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + }) + .boxed() +} + +/// Runs a spawned chunk's `root` over `scope` on a planning run of its own, on one task that +/// drives the run to its end and awaits only its reads, and yields its mapped batches with the +/// row each starts at. A planner the chunk spawns in turn runs as another root of the same run. +/// +/// One task per chunk, not one per step: a chunk's run waits for reads many times, and a task +/// hop at each wait is latency on the chunk's critical path that no other chunk's work fills. +fn chunk_stream( + handle: &Handle, + root: Box, + scope: WorkScope, + io: Arc, + map_fn: Arc VortexResult + Send + Sync>, +) -> BoxStream<'static, VortexResult<(u64, A)>> { + let (mut send, recv) = mpsc::channel(1); + let task = handle.spawn(async move { + let mut run = Run::new(); + run.admit(root, scope, Arc::clone(&io)); + loop { + match run.advance()? { + Progress::Batch(batch) => { + let start = batch.scope.rows.start; + let batch = map_fn(batch.array); + let failed = batch.is_err(); + if send.send(batch.map(|batch| (start, batch))).await.is_err() || failed { + return Ok(()); + } + } + Progress::Spawn(scope, child) => { + run.admit(child, scope, Arc::clone(&io)); + } + Progress::RootDone(_) if run.live_roots() > 0 => {} + Progress::RootDone(_) | Progress::Idle => return Ok(()), + Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + }); + let errors = stream::once(task) + .filter_map(|result: VortexResult<()>| future::ready(result.err().map(Err))); + stream::select(recv, errors).boxed() +} diff --git a/vortex-layout/src/scan/v2/split.rs b/vortex-layout/src/scan/v2/split.rs deleted file mode 100644 index 624a775ad61..00000000000 --- a/vortex-layout/src/scan/v2/split.rs +++ /dev/null @@ -1,74 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::future; -use std::sync::Arc; - -use futures::SinkExt; -use futures::channel::mpsc; -use vortex_array::ArrayRef; -use vortex_error::VortexResult; -use vortex_io::request::IoService; -use vortex_scan::planning::driver::Progress; -use vortex_scan::planning::driver::Run; -use vortex_scan::planning::planner::Planner; -use vortex_scan::planning::planner::WorkScope; - -/// One split of a scan, ready to be admitted to a planning run. -/// -/// The split's planner projects the rows of `scope`, and its morsels emit one batch per -/// projection split with selected rows, in whatever order their reads arrive. -pub(super) struct SplitPlan { - /// The rows of the scan's root row domain the split covers. - pub(super) scope: WorkScope, - /// The split's first planner. - pub(super) root: Box, -} - -impl SplitPlan { - /// Publishes mapped batches as they finish, bounding how far projection runs ahead of its - /// consumer. Unordered scans do not need to retain the whole split to sort its output. - pub(super) async fn run_unordered( - self, - io: Arc, - map: Arc VortexResult + Send + Sync>, - mut send: mpsc::Sender>, - ) -> VortexResult<()> { - let mut run = Run::new(); - run.admit(self.root, self.scope, io.session()); - loop { - match run.advance()? { - Progress::Batch(batch) => { - let batch = map(batch.array); - let failed = batch.is_err(); - if send.send(batch).await.is_err() || failed { - return Ok(()); - } - } - Progress::RootDone(_) | Progress::Idle => return Ok(()), - Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, - } - } - } - - /// Runs the split on a planning run of its own, reading through a session of `io`, and - /// returns its batches in row order. - /// - /// The run advances inside this future, on whichever thread polls it, and the future awaits - /// its reads between steps. Nothing is handed to another thread. - pub(super) async fn run(self, io: Arc) -> VortexResult> { - let mut run = Run::new(); - run.admit(self.root, self.scope, io.session()); - let mut batches = Vec::new(); - loop { - match run.advance()? { - Progress::Batch(batch) => batches.push(batch), - Progress::RootDone(_) | Progress::Idle => break, - Progress::Waiting => future::poll_fn(|cx| run.poll_completion(cx)).await?, - } - } - // Morsels finish in whatever order their reads arrive. - batches.sort_by_key(|batch| batch.scope.rows.start); - Ok(batches.into_iter().map(|batch| batch.array).collect()) - } -} diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index ba29376bafb..eabf863eb47 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -43,18 +43,24 @@ use vortex_arrow::ArrowSessionExt; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_io::session::RuntimeSessionExt; +use vortex_scan::planning::driver::Batch; +use vortex_scan::planning::driver::Progress; +use vortex_scan::planning::driver::Run; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::PlannerOutput; +use vortex_scan::planning::planner::WorkScope; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; use super::MAX_ROWS_PER_CHUNK; -use super::PreparedScan; use super::ROW_BLOCK_SIZE; use super::ScanBuilder; use super::ScanCache; use super::ScanFile; -use super::row_chunks; +use super::ScanRun; use super::rows_per_chunk; use crate::LayoutRef; use crate::LayoutStrategy; @@ -66,6 +72,7 @@ use crate::layouts::table::TableStrategy; use crate::layouts::zoned::writer::ZonedLayoutOptions; use crate::layouts::zoned::writer::ZonedStrategy; use crate::scan::execution::SegmentLocation; +use crate::scan::execution::row_chunks; use crate::scan::scan_builder::ScanBuilder as V1ScanBuilder; use crate::segments::SegmentFuture; use crate::segments::SegmentId; @@ -188,6 +195,49 @@ fn builder( )) } +/// Runs `rows` of a prepared scan on a planning run of its own and returns its batches, in row +/// order, each with the rows it stands for. +async fn run_batches( + scan: &ScanRun, + rows: Range, +) -> VortexResult> { + let mut run = Run::new(); + run.admit( + scan.scan.planner(rows.clone(), scan.selection.clone()), + WorkScope { + file_ordinal: 0, + rows, + }, + scan.io.session(), + ); + let mut batches = Vec::new(); + loop { + match run.advance()? { + Progress::Batch(batch) => batches.push(batch), + Progress::Spawn(scope, chunk) => { + run.admit(chunk, scope, scan.io.session()); + } + Progress::RootDone(_) if run.live_roots() > 0 => {} + Progress::RootDone(_) | Progress::Idle => break, + Progress::Waiting => std::future::poll_fn(|cx| run.poll_completion(cx)).await?, + } + } + batches.sort_by_key(|batch| batch.scope.rows.start); + Ok(batches) +} + +/// Runs `rows` of a prepared scan on a planning run of its own and returns its arrays. +async fn run_rows( + scan: &ScanRun, + rows: Range, +) -> VortexResult> { + Ok(run_batches(scan, rows) + .await? + .into_iter() + .map(|batch| batch.array) + .collect()) +} + async fn collect( dtype: DType, stream: impl futures::Stream>, @@ -214,8 +264,7 @@ async fn assert_same_as_v1( .with_some_filter(builder.filter.clone()) .with_selection(builder.selection.clone()) .with_row_offset(builder.row_offset) - .with_ordered(builder.ordered) - .with_concurrency(builder.concurrency); + .with_ordered(builder.ordered); if let Some(rows) = builder.row_range.clone() { v1 = v1.with_row_range(rows); } @@ -469,27 +518,26 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { reads: AtomicUsize::new(0), }); let segments = Arc::clone(&counted) as Arc; - let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; - let splits = scan.split_plans(); + let scan = ScanRun::try_new(builder(&session, &segments, &layout)?)?; + let chunks = row_chunks(scan.rows.clone(), scan.scan.chunk_rows); assert_eq!( - splits.len(), + chunks.len(), (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(MAX_ROWS_PER_CHUNK) ); - for (index, split) in splits.into_iter().enumerate() { - let rows = split.scope.rows.clone(); + for (index, rows) in chunks.into_iter().enumerate() { let before = counted.reads.load(Ordering::Relaxed); - let arrays = split.run(Arc::clone(&scan.io)).await?; + let arrays = run_rows(&scan, rows.clone()).await?; let reads = counted.reads.load(Ordering::Relaxed) - before; let code_chunks = (rows.end - 1) / STORED_CHUNK_ROWS as u64 - rows.start / STORED_CHUNK_ROWS as u64 + 1; - // The first split reads every values chunk once for the whole scan, plus its codes. - let code_chunks = usize::try_from(code_chunks)?; + // The first chunk reads every values chunk once for the whole scan; its codes are one of + // them, and the session shares that read. Later chunks read only their codes. assert_eq!( reads, if index == 0 { - STORED_CHUNKS + code_chunks + STORED_CHUNKS } else { - code_chunks + usize::try_from(code_chunks)? } ); let expected = @@ -535,7 +583,7 @@ async fn projection_skips_unsupported_unselected_layouts() -> VortexResult<()> { async fn unordered_returns_every_row() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; - let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?.with_ordered(false))?; + let scan = ScanRun::try_new(builder(&session, &segments, &layout)?.with_ordered(false))?; let chunks: Vec = scan.execute_stream()?.try_collect().await?; let rows: usize = chunks.iter().map(|chunk| chunk.len()).sum(); assert_eq!(rows, CHUNKS * CHUNK_ROWS); @@ -603,10 +651,11 @@ async fn unordered_yields_before_the_split_finishes() -> VortexResult<()> { blocked, release: gate.shared(), }); - let mut scan = - PreparedScan::try_new(builder(&session, &segments, &layout)?.with_ordered(false))?; - // Put both projection morsels in one split, as a dense filtered scan would. - scan.splits = std::iter::once(0..layout.row_count()).collect(); + let mut scan = ScanRun::try_new(builder(&session, &segments, &layout)?.with_ordered(false))?; + // Put both projection morsels in one chunk, as a dense filtered scan would. + Arc::get_mut(&mut scan.scan) + .ok_or_else(|| vortex_err!("the prepared scan is not shared yet"))? + .chunk_rows = usize::try_from(layout.row_count())?; let mut batches = scan.execute_stream()?; let first = batches.try_next().await; let _ = release.send(()); @@ -629,7 +678,7 @@ async fn primitive_column() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; let dtype = layout.dtype().clone(); - let scan = PreparedScan::try_new( + let scan = ScanRun::try_new( builder(&session, &segments, &layout)? .with_projection(get_item("ints", root()).bind(&dtype)?) .with_row_range(990..1010), @@ -655,8 +704,11 @@ async fn work_chunks_ignore_layout_boundaries( ) -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, strategy.get().as_ref()).await?; - let scan = PreparedScan::try_new(builder(&session, &segments, &layout)?)?; - assert_eq!(scan.splits, vec![0..(CHUNKS * CHUNK_ROWS) as u64]); + let scan = ScanRun::try_new(builder(&session, &segments, &layout)?)?; + assert_eq!( + row_chunks(scan.rows.clone(), scan.scan.chunk_rows), + vec![0..(CHUNKS * CHUNK_ROWS) as u64] + ); let arrays: Vec<_> = scan.execute_stream()?.try_collect().await?; assert_eq!(arrays.len(), 1); assert_eq!(arrays[0].len(), CHUNKS * CHUNK_ROWS); @@ -734,13 +786,18 @@ async fn scan_cache_reuses_values_between_scans( async fn selected_split_scope_is_bounded_before_projection() -> VortexResult<()> { let session = new_session().with_tokio(); let (segments, layout) = write(&session, table().as_ref()).await?; - let scan = PreparedScan::try_new( + let scan = ScanRun::try_new( builder(&session, &segments, &layout)? .with_row_indices(StrictSortedBuffer::try_new(Buffer::from(vec![7, 2001]))?), )?; - let splits = scan.split_plans(); - assert_eq!(splits.len(), 1); - assert_eq!(splits[0].scope.rows, 7..2002); + let mut chunks = scan + .scan + .chunk_planner(scan.rows.clone(), scan.selection.clone(), None); + let PlannerOutput::Spawn(scope, _) = chunks.compute()? else { + vortex_bail!("the selected rows were not planned"); + }; + assert_eq!(scope.rows, 7..2002); + assert!(matches!(chunks.compute()?, PlannerOutput::Done)); Ok(()) } @@ -801,3 +858,68 @@ fn partition_rows_tile_the_file_on_row_blocks() { "a range inside one block owns no rows" ); } + +/// Zone proofs run once before any chunk is planned: rows every proof excludes are never +/// given a chunk, and a scan whose zones all fail plans nothing. +#[rstest] +#[case::tail(2_999, Some(3_000..4_000))] +#[case::nothing(10_000, None)] +#[tokio::test] +async fn pruned_rows_get_no_chunk( + #[case] above: i32, + #[case] expected: Option>, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, zoned_table().as_ref()).await?; + let scan = ScanRun::try_new( + builder(&session, &segments, &layout)?.with_some_filter(Some( + gt(get_item("ints", root()), lit(above)).bind(layout.dtype())?, + )), + )?; + let batches = run_batches(&scan, scan.rows.clone()).await?; + assert_eq!( + batches + .iter() + .map(|batch| batch.scope.rows.clone()) + .collect::>(), + expected.into_iter().collect::>() + ); + Ok(()) +} + +/// Scans of one file that share a cache share their filter plans, so the zone proofs run once: +/// the second scan with the same filter re-reads the data it needs but never the zone maps. +#[tokio::test] +async fn shared_cache_proves_zones_once() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, zoned_table().as_ref()).await?; + let counted = Arc::new(CountingSegments { + inner: segments, + reads: AtomicUsize::new(0), + }); + let segments = Arc::clone(&counted) as Arc; + let mut file = scan_file(&segments, &layout)?; + let cache = ScanCache::default(); + file.cache = Some(cache.clone()); + let filter = gt(get_item("ints", root()), lit(2_999i32)).bind(layout.dtype())?; + let projection = get_item("ints", root()).bind(layout.dtype())?; + assert!(cache.filter_plans(&filter).is_none()); + + let mut reads = Vec::new(); + for _ in 0..2 { + let before = counted.reads.load(Ordering::Relaxed); + let scan = ScanBuilder::new(session.clone(), file.clone()) + .with_projection(projection.clone()) + .with_some_filter(Some(filter.clone())); + let actual = collect(projection.dtype().clone(), scan.into_stream()?).await?; + assert_eq!(actual.len(), 1_000); + reads.push(counted.reads.load(Ordering::Relaxed) - before); + } + let plans = cache + .filter_plans(&filter) + .expect("the first scan kept its filter plans"); + assert_eq!(plans.proofs().count(), 1); + // Each scan reads the one chunk the proof keeps; only the first reads the zone map. + assert_eq!(reads, vec![2, 1]); + Ok(()) +} diff --git a/vortex-scan/src/planning/driver/mod.rs b/vortex-scan/src/planning/driver/mod.rs index 8cb78099272..e9cf27da3c2 100644 --- a/vortex-scan/src/planning/driver/mod.rs +++ b/vortex-scan/src/planning/driver/mod.rs @@ -65,6 +65,9 @@ pub enum Progress { Batch(Batch), /// Every item descended from the root has retired, and its IO session was cleared. RootDone(RootId), + /// A planner declared this child independent. The run has not queued it: the owner runs it, + /// on this run as another root through [`Run::admit`] or on a run of its own elsewhere. + Spawn(WorkScope, Box), /// Every live item waits for IO. Await [`Run::poll_completion`], or pass the next completion /// to [`Run::complete`]. Waiting, @@ -462,6 +465,10 @@ impl Run { self.spawn(work.root, scope, Item::Morsel(morsel), path); self.queue.insert(work); } + Output::Planner(PlannerOutput::Spawn(scope, child)) => { + self.events.push_back(Progress::Spawn(scope, child)); + self.queue.insert(work); + } Output::Morsel(MorselOutput::Batch(array)) => { if array.is_empty() { vortex_bail!("morsel for scope {:?} returned an empty batch", work.scope); @@ -527,7 +534,8 @@ impl Driver { } } - /// Runs `root` and everything it spawns, returning batches in emission order. + /// Runs `root` and everything it spawns, returning batches in emission order. A planner + /// spawned as independent runs as another root of the same run. pub fn run(&self, root: Box) -> VortexResult> { let mut run = Run::new(); run.step_limit = self.step_limit; @@ -540,6 +548,10 @@ impl Driver { loop { match run.advance()? { Progress::Batch(batch) => batches.push(batch), + Progress::Spawn(scope, child) => { + run.admit(child, scope, Arc::clone(&self.io)); + } + Progress::RootDone(_) if run.live_roots() > 0 => {} Progress::RootDone(_) | Progress::Idle => return Ok(batches), Progress::Waiting => run.complete(self.io.wait()?)?, } diff --git a/vortex-scan/src/planning/driver/tests.rs b/vortex-scan/src/planning/driver/tests.rs index 389ad24372d..9461cd4cd36 100644 --- a/vortex-scan/src/planning/driver/tests.rs +++ b/vortex-scan/src/planning/driver/tests.rs @@ -503,6 +503,7 @@ fn drain(run: &mut Run) -> VortexResult> { events.push(match run.advance()? { Progress::Batch(batch) => format!("batch {:?} {:?}", batch.root, batch.scope.rows), Progress::RootDone(root) => format!("done {root:?}"), + Progress::Spawn(scope, _) => format!("spawn {:?}", scope.rows), Progress::Waiting => { events.push("waiting".to_string()); return Ok(events); diff --git a/vortex-scan/src/planning/planner.rs b/vortex-scan/src/planning/planner.rs index 91d875ff317..cf795c99d7b 100644 --- a/vortex-scan/src/planning/planner.rs +++ b/vortex-scan/src/planning/planner.rs @@ -51,6 +51,11 @@ pub enum PlannerOutput { NeedsIO(IoBatch), /// A child planner, with the scope it is responsible for. Planner(WorkScope, Box), + /// A child planner independent of this one, with the scope it is responsible for. The run + /// hands it to its owner as [`Progress::Spawn`](super::driver::Progress::Spawn) instead of + /// running it: the owner may run it anywhere, including on another thread, and this planner + /// continues at once. + Spawn(WorkScope, Box), /// A morsel authorised to produce arrays for the scope. Morsel(WorkScope, Box), } @@ -75,6 +80,7 @@ impl fmt::Debug for PlannerOutput { Self::Continue => f.write_str("Continue"), Self::NeedsIO(batch) => f.debug_tuple("NeedsIO").field(batch).finish(), Self::Planner(scope, _) => f.debug_tuple("Planner").field(scope).finish(), + Self::Spawn(scope, _) => f.debug_tuple("Spawn").field(scope).finish(), Self::Morsel(scope, _) => f.debug_tuple("Morsel").field(scope).finish(), } }