diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 208691f0b0e..a5d3b881ed5 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -137,9 +137,6 @@ 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 { @@ -178,7 +175,6 @@ impl FileOpener for VortexOpener { 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); @@ -444,6 +440,11 @@ impl FileOpener for VortexOpener { let mut builder = V2ScanBuilder::new(vxf.session().clone(), scan_file) .with_projection(scan_projection) .with_some_filter(native_filter); + // How many partitions split the file, assuming they cut it evenly. + let readers = byte_range.as_ref().map_or(1, |bytes| { + let len = (bytes.end - bytes.start).max(1); + usize::try_from(file.object_meta.size.div_ceil(len)).unwrap_or(usize::MAX) + }); if let Some(bytes) = byte_range { // Every partition maps the same byte boundary to the same row boundary. let rows = partition_row_range(bytes, file.object_meta.size, vxf.row_count()); @@ -459,14 +460,16 @@ impl FileOpener for VortexOpener { { builder = builder.with_selection(selection.clone()); } - // 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(|| { + // Unless set, chunks run beside each other only as far as the cores outnumber the + // partitions reading this file: a file split a partition per core is read in each + // scan's own run, a file read whole by one partition on every core. + let concurrency = scan_concurrency.filter(|&c| c > 0).unwrap_or_else(|| { get_available_parallelism() .unwrap_or(1) - .div_ceil(partitions) + .div_ceil(readers) .max(1) }); + builder = builder.with_concurrency(concurrency); let handle = vxf.session().handle(); // Lowering and planning a wide file is real CPU: it runs where compute runs. let scan = handle @@ -1147,7 +1150,6 @@ mod tests { 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, @@ -1557,7 +1559,6 @@ mod tests { 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, @@ -1649,7 +1650,6 @@ mod tests { 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, @@ -1832,7 +1832,6 @@ mod tests { 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, @@ -1897,7 +1896,6 @@ mod tests { 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, @@ -2113,7 +2111,6 @@ mod tests { 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/scan_v2.rs b/vortex-datafusion/src/persistent/scan_v2.rs index 10ea7b969ea..05a010d0405 100644 --- a/vortex-datafusion/src/persistent/scan_v2.rs +++ b/vortex-datafusion/src/persistent/scan_v2.rs @@ -25,14 +25,13 @@ use futures::stream::BoxStream; use vortex::array::ArrayRef; use vortex::error::VortexResult; use vortex::error::vortex_err; -use vortex::io::request::IoSource; +use vortex::io::request::IoService; use vortex::io::runtime::Handle; use vortex::layout::scan::v2::PreparedScan; 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_utils::parallelism::get_available_parallelism; /// Maps a scan's batches to Arrow on the task that produced them. type MapBatch = Arc VortexResult + Send + Sync>; @@ -40,7 +39,7 @@ type MapBatch = Arc VortexResult + Send + Sync> /// Runs one partition's V2 scan on DataFusion's runtime. pub(crate) struct ScanDriver { handle: Handle, - /// Chunks run beside each other per worker thread; at 1 they run on the scan's own run. + /// The scan's chunks run at once; at 1 they run on the scan's own run. concurrency: usize, /// Whether batches are yielded in row order, which holds them until the scan is done. ordered: bool, @@ -64,8 +63,8 @@ impl ScanDriver { map: impl Fn(ArrayRef) -> VortexResult + Send + Sync + 'static, ) -> BoxStream<'static, VortexResult> { let map: MapBatch = Arc::new(map); - let in_flight = self.concurrency * get_available_parallelism().unwrap_or(1); - let inline = in_flight <= 1; + let in_flight = self.concurrency; + let inline = self.concurrency <= 1; let scope = WorkScope { file_ordinal: 0, rows: scan.rows().clone(), @@ -73,7 +72,7 @@ impl ScanDriver { let io = scan.io(); let (chunk_handle, chunk_io, chunk_map) = (self.handle.clone(), Arc::clone(&io), Arc::clone(&map)); - let batches = run_stream(self.handle, root, scope, io.session(), map, inline) + let batches = run_stream(self.handle, root, scope, Arc::clone(&io), map, inline) .map_ok(move |yielded| match yielded { Yielded::Batch(start, batch) => { stream::once(future::ready(Ok((start, batch)))).boxed() @@ -82,7 +81,7 @@ impl ScanDriver { &chunk_handle, chunk, scope, - chunk_io.session(), + Arc::clone(&chunk_io), Arc::clone(&chunk_map), ), }) @@ -128,7 +127,8 @@ struct RunState { handle: Handle, /// Taken while a CPU step runs on its task and put back with its result. run: Option, - io: Arc, + /// Gives every root of the run a session of its own: retiring a root clears its session. + io: Arc, map: MapBatch, /// Whether a spawned chunk runs on this run, as another root, instead of being yielded. inline: bool, @@ -143,12 +143,12 @@ fn run_stream( handle: Handle, root: Box, scope: WorkScope, - io: Arc, + io: Arc, map: MapBatch, inline: bool, ) -> BoxStream<'static, VortexResult> { let mut run = Run::new(); - run.admit(root, scope, Arc::clone(&io)); + run.admit(root, scope, io.session()); let state = RunState { handle, run: Some(run), @@ -187,7 +187,7 @@ fn run_stream( return Ok(Some((Yielded::Batch(start, batch), state))); } Step::Spawn(scope, child) if state.inline => { - run.admit(child, scope, Arc::clone(&state.io)); + run.admit(child, scope, state.io.session()); } Step::Spawn(scope, child) => { return Ok(Some((Yielded::Spawn(scope, child), state))); @@ -211,13 +211,13 @@ fn chunk_stream( handle: &Handle, root: Box, scope: WorkScope, - io: Arc, + io: Arc, map: MapBatch, ) -> BoxStream<'static, VortexResult<(u64, RecordBatch)>> { 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)); + run.admit(root, scope, io.session()); loop { match run.advance()? { Progress::Batch(batch) => { @@ -229,7 +229,7 @@ fn chunk_stream( } } Progress::Spawn(scope, child) => { - run.admit(child, scope, Arc::clone(&io)); + run.admit(child, scope, io.session()); } Progress::RootDone(_) if run.live_roots() > 0 => {} Progress::RootDone(_) | Progress::Idle => return Ok(()), diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 644fea4e71c..c738df95550 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -49,6 +49,31 @@ use crate::convert::exprs::ExpressionConverter; use crate::persistent::reader::DefaultVortexReaderFactory; use crate::persistent::reader::VortexReaderFactory; +/// The fewest rows a filtered V2 scan partition reads when the engine splits files. Each +/// partition proves zone maps, plans its filter and compiles its plans once, so a partition much +/// smaller than a few output windows spends more on that than on reading rows. An unfiltered scan +/// keeps the engine's partitions: every row it reads flows into the operators above it, which +/// run a task per partition. +const V2_MIN_ROWS_PER_PARTITION: usize = 256 * 1024; + +/// Caps `target_partitions` so no V2 partition reads fewer than [`V2_MIN_ROWS_PER_PARTITION`] +/// rows, when the files' row counts are known. +fn v2_target_partitions(target_partitions: usize, config: &FileScanConfig) -> usize { + let rows = config + .file_groups + .iter() + .flat_map(|group| group.iter()) + .map(|file| { + file.statistics + .as_ref() + .and_then(|statistics| statistics.num_rows.get_value().copied()) + }) + .sum::>(); + rows.map_or(target_partitions, |rows| { + target_partitions.min(rows.div_ceil(V2_MIN_ROWS_PER_PARTITION).max(1)) + }) +} + /// File scan implementation for reading one or more `.vortex` files. /// /// `VortexSource` is the lower-level read component underneath @@ -378,7 +403,6 @@ 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) @@ -468,6 +492,11 @@ impl FileSource for VortexSource { if config.file_compression_type.is_compressed() { return Ok(None); } + let target_partitions = if self.options.scan_v2 && self.vortex_predicate.is_some() { + v2_target_partitions(target_partitions, config) + } else { + target_partitions + }; let Some(file_groups) = FileGroupPartitioner::new() .with_target_partitions(target_partitions) .with_repartition_file_min_size(repartition_file_min_size) diff --git a/vortex-datafusion/src/persistent/tests.rs b/vortex-datafusion/src/persistent/tests.rs index 5d91242c91d..15509f5f941 100644 --- a/vortex-datafusion/src/persistent/tests.rs +++ b/vortex-datafusion/src/persistent/tests.rs @@ -486,9 +486,10 @@ async fn test_repartitioned_scan_matches_non_repartitioned_for_uneven_splits( let session = VortexSession::default(); let path = object_store::path::Path::parse("/split-aligned-repartition.vortex")?; - let chunk_1_len = 2_000; - let chunk_2_len = 5_000; - let chunk_3_len = 6_000; + // Enough rows that a V2 scan, which keeps a minimum of rows per partition, still splits. + let chunk_1_len = 200_000; + let chunk_2_len = 250_000; + let chunk_3_len = 300_000; let row_count = chunk_1_len + chunk_2_len + chunk_3_len; let chunk_1 = StructArray::try_new( @@ -582,6 +583,38 @@ async fn test_repartitioned_scan_matches_non_repartitioned_for_uneven_splits( Ok(()) } +/// A filtered V2 scan of a file smaller than the minimum rows per partition reads it in one +/// partition. An unfiltered scan keeps the engine's partitions: every row it reads flows into the +/// operators above it. +#[tokio::test] +async fn v2_scan_keeps_partitions_unless_filtered() -> anyhow::Result<()> { + let store = Arc::new(InMemory::new()) as _; + let session = VortexSession::default(); + let path = object_store::path::Path::parse("/small.vortex")?; + let table = StructArray::try_new( + ["value"].into(), + vec![Buffer::from_iter(0_i32..100_000).into_array()], + 100_000, + Validity::NonNullable, + )? + .into_array(); + let mut writer = ObjectStoreWrite::new(Arc::clone(&store), &path).await?; + session + .write_options() + .write(&mut writer, table.to_array_stream()) + .await?; + writer.shutdown().await?; + + let ctx = make_session(store, true, true); + let unfiltered = count_query_partitions(&ctx, "SELECT value FROM '/small.vortex'").await?; + let filtered = + count_query_partitions(&ctx, "SELECT value FROM '/small.vortex' WHERE value > 10").await?; + + assert_eq!(unfiltered, 4); + assert_eq!(filtered, 1); + Ok(()) +} + /// Roundtrip an `arrow.uuid` extension column through a Vortex file: write the column directly /// via the session-aware Arrow→Vortex conversion, then `SELECT *` and assert both the field /// metadata and the underlying values survive the trip. diff --git a/vortex-file/src/v2/scan.rs b/vortex-file/src/v2/scan.rs index b5a3f104d9c..8e6ff1df736 100644 --- a/vortex-file/src/v2/scan.rs +++ b/vortex-file/src/v2/scan.rs @@ -30,7 +30,7 @@ use vortex_array::VortexSessionExecute; use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_err; -use vortex_io::request::IoSource; +use vortex_io::request::IoService; use vortex_io::runtime::Handle; use vortex_layout::scan::v2::ChunkPlanner; use vortex_layout::scan::v2::PreparedScan; @@ -108,6 +108,7 @@ impl ScanDriver { map_fn: MapFn, ) -> BoxStream<'static, VortexResult> { let handle = self.handle.clone(); + let builder = builder.with_concurrency(self.in_flight()); stream::once(async move { handle .spawn_cpu(move || { @@ -120,13 +121,18 @@ impl ScanDriver { .boxed() } + /// The chunks running at once: `concurrency` per worker thread. + fn in_flight(&self) -> usize { + self.concurrency * get_available_parallelism().unwrap_or(1) + } + fn execute( self, scan: Arc, map_fn: MapFn, ) -> BoxStream<'static, VortexResult> { - let in_flight = self.concurrency * get_available_parallelism().unwrap_or(1); - let inline = in_flight <= 1; + let in_flight = self.in_flight(); + let inline = self.concurrency <= 1; let scope = WorkScope { file_ordinal: 0, rows: scan.rows().clone(), @@ -135,7 +141,7 @@ impl ScanDriver { let root = compose(&scan); let (chunk_handle, chunk_io, chunk_map) = (self.handle.clone(), Arc::clone(&io), Arc::clone(&map_fn)); - let batches = run_stream(self.handle, root, scope, io.session(), map_fn, inline) + let batches = run_stream(self.handle, root, scope, Arc::clone(&io), map_fn, inline) .map_ok(move |yielded| match yielded { Yielded::Batch(start, batch) => { stream::once(future::ready(Ok((start, batch)))).boxed() @@ -144,7 +150,7 @@ impl ScanDriver { &chunk_handle, chunk, scope, - chunk_io.session(), + Arc::clone(&chunk_io), Arc::clone(&chunk_map), ), }) @@ -212,7 +218,8 @@ struct RunState { handle: Handle, /// Taken while a CPU step runs on its task and put back with its result. run: Option, - io: Arc, + /// Gives every root of the run a session of its own: retiring a root clears its session. + io: Arc, map_fn: MapFn, /// Whether a spawned chunk runs on this run, as another root, instead of being yielded. inline: bool, @@ -225,12 +232,12 @@ fn run_stream( handle: Handle, root: Box, scope: WorkScope, - io: Arc, + io: Arc, map_fn: MapFn, inline: bool, ) -> BoxStream<'static, VortexResult>> { let mut run = Run::new(); - run.admit(root, scope, Arc::clone(&io)); + run.admit(root, scope, io.session()); let state = RunState { handle, run: Some(run), @@ -269,7 +276,7 @@ fn run_stream( return Ok(Some((Yielded::Batch(start, batch), state))); } Step::Spawn(scope, child) if state.inline => { - run.admit(child, scope, Arc::clone(&state.io)); + run.admit(child, scope, state.io.session()); } Step::Spawn(scope, child) => { return Ok(Some((Yielded::Spawn(scope, child), state))); @@ -290,13 +297,13 @@ fn chunk_stream( handle: &Handle, root: Box, scope: WorkScope, - io: Arc, + io: Arc, map_fn: MapFn, ) -> 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)); + run.admit(root, scope, io.session()); loop { match run.advance()? { Progress::Batch(batch) => { @@ -308,7 +315,7 @@ fn chunk_stream( } } Progress::Spawn(scope, child) => { - run.admit(child, scope, Arc::clone(&io)); + run.admit(child, scope, io.session()); } Progress::RootDone(_) if run.live_roots() > 0 => {} Progress::RootDone(_) | Progress::Idle => return Ok(()), diff --git a/vortex-layout/src/plan/physical/optimizer/mod.rs b/vortex-layout/src/plan/physical/optimizer/mod.rs index adf38185b6b..a74ee24acc9 100644 --- a/vortex-layout/src/plan/physical/optimizer/mod.rs +++ b/vortex-layout/src/plan/physical/optimizer/mod.rs @@ -29,7 +29,7 @@ use super::plans::ExpressionLegacyStatsRule; use super::plans::ExpressionPackRule; use super::plans::ExpressionTakeRule; use super::plans::ExpressionZonedRule; -use super::plans::LegacyStatsRule; +use super::plans::LegacyStatsDataRule; use super::plans::ZonedDataRule; use crate::plan::physical::plans::EvalEvalRule; use crate::plan::physical::plans::EvalPackRule; @@ -37,8 +37,8 @@ use crate::plan::physical::plans::EvalPackRule; static EVAL_IDENTITY_RULE: PlanReduceRuleAdapter = PlanReduceRuleAdapter::new(EvalIdentityRule); -static LEGACY_STATS_RULE: PlanReduceRuleAdapter = - PlanReduceRuleAdapter::new(LegacyStatsRule); +static LEGACY_STATS_DATA_RULE: PlanReduceRuleAdapter = + PlanReduceRuleAdapter::new(LegacyStatsDataRule); static ZONED_DATA_RULE: PlanReduceRuleAdapter = PlanReduceRuleAdapter::new(ZonedDataRule); @@ -46,7 +46,7 @@ static ZONED_DATA_RULE: PlanReduceRuleAdapter = static PLAN_RULES: PlanRuleSet = PlanRuleSet::new(&[ &EVAL_IDENTITY_RULE, &PlanReduceRuleAdapter::new(EvalPackRule), - &LEGACY_STATS_RULE, + &LEGACY_STATS_DATA_RULE, &ZONED_DATA_RULE, ]); diff --git a/vortex-layout/src/plan/physical/plans/eval.rs b/vortex-layout/src/plan/physical/plans/eval.rs index 043f0676198..f163a548f4b 100644 --- a/vortex-layout/src/plan/physical/plans/eval.rs +++ b/vortex-layout/src/plan/physical/plans/eval.rs @@ -33,6 +33,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::pack::is_identity_expression; /// Applies an expression to the output of its child. #[derive(Clone, Debug)] @@ -150,17 +151,20 @@ fn validate_expression_child(expression: &BoundExpression, child: &PlanRef) -> V Ok(()) } -/// Removes an [`Eval`] whose expression is the identity expression. +/// Removes an [`Eval`] whose expression is the identity: the root, or a non-nullable struct +/// rebuilt from every field of the root in order, as pushing a projection into a +/// [`PackPlan`] leaves. #[derive(Debug)] pub(crate) struct EvalIdentityRule; impl PlanReduceRule for EvalIdentityRule { fn reduce(&self, plan: &Plan, _session: &VortexSession) -> VortexResult> { - if plan.expression().is_root() { - Ok(Some(plan.child_plan()?)) - } else { - Ok(None) + let expression = plan.expression(); + if !expression.is_root() && !expression.dtype().is_struct() { + return Ok(None); } + let child = plan.child_plan()?; + Ok(is_identity_expression(expression, child.dtype())?.then_some(child)) } } diff --git a/vortex-layout/src/plan/physical/plans/legacy_stats.rs b/vortex-layout/src/plan/physical/plans/legacy_stats.rs index 8aa0e897da7..e7e050258bd 100644 --- a/vortex-layout/src/plan/physical/plans/legacy_stats.rs +++ b/vortex-layout/src/plan/physical/plans/legacy_stats.rs @@ -39,10 +39,10 @@ const ZONES: usize = 1; /// Data alongside a legacy `vortex.stats` zone table, whose fields are named by [`Stat`] rather /// 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; a plan -/// that reads the data is then optimized to the data alone. Compiled unoptimized, it reads its -/// data and prunes nothing. +/// Beneath a pruning plan 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; a plan that reads the data is optimized to the data alone. Compiled +/// unoptimized, it reads its data and prunes nothing. #[derive(Clone, Debug)] pub struct LegacyStats; @@ -156,28 +156,35 @@ impl PlanVTable for LegacyStats { /// A plan whose data is already [`Zoned`] keeps only that data: legacy stats are used only when no /// aggregate zone map summarises the same rows. A zero zone length, as older files may record, or /// a table with no stat an aggregate can stand for, leaves only the data too. +fn to_zoned(plan: &LegacyStatsPlan) -> VortexResult { + let data = plan.data_plan()?; + if plan.zone_len() == 0 || data.is::() { + return Ok(data); + } + let zones = plan.zones_plan()?; + let (aggregate_fns, conversion) = + legacy_stats_to_aggregates(plan.dtype(), zones.dtype(), plan.stats()); + if aggregate_fns.is_empty() { + return Ok(data); + } + let zones = EvalPlan::try_new(conversion.bind(zones.dtype())?, zones)?.into_plan(); + Ok(ZonedPlan::new(data, zones, plan.zone_len(), aggregate_fns).into_plan()) +} + +/// Replaces a legacy stats plan by its data, as [`ZonedDataRule`](super::ZonedDataRule) does for +/// the [`Zoned`] plan it would convert to. A pruning plan over it is converted by +/// [`ExpressionLegacyStatsRule`] before this runs, so the zone table is never lowered or +/// converted only to be dropped. #[derive(Debug)] -pub(crate) struct LegacyStatsRule; +pub(crate) struct LegacyStatsDataRule; -impl PlanReduceRule for LegacyStatsRule { +impl PlanReduceRule for LegacyStatsDataRule { fn reduce( &self, plan: &Plan, _session: &VortexSession, ) -> VortexResult> { - let data = plan.data_plan()?; - if plan.zone_len() == 0 || data.is::() { - return Ok(Some(data)); - } - let zones = plan.zones_plan()?; - let (aggregate_fns, conversion) = - legacy_stats_to_aggregates(plan.dtype(), zones.dtype(), plan.stats()); - if aggregate_fns.is_empty() { - return Ok(Some(data)); - } - let zones = EvalPlan::try_new(conversion.bind(zones.dtype())?, zones)?.into_plan(); - let zoned = ZonedPlan::new(data, zones, plan.zone_len(), aggregate_fns); - Ok(Some(zoned.into_plan())) + plan.data_plan().map(Some) } } @@ -194,16 +201,13 @@ impl PlanParentReduceRule for ExpressionLegacyStatsRule { child: &LegacyStatsPlan, parent: &Plan, _child_idx: usize, - session: &VortexSession, + _session: &VortexSession, ) -> VortexResult> { if !is_zone_pruning(parent)? { return Ok(None); } - let Some(converted) = LegacyStatsRule.reduce(child, session)? else { - return Ok(None); - }; Ok(Some( - EvalPlan::try_new(parent.expression().clone(), converted)?.into_plan(), + EvalPlan::try_new(parent.expression().clone(), to_zoned(child)?)?.into_plan(), )) } } diff --git a/vortex-layout/src/plan/physical/plans/mod.rs b/vortex-layout/src/plan/physical/plans/mod.rs index 3d37390cfcb..4dd6e655999 100644 --- a/vortex-layout/src/plan/physical/plans/mod.rs +++ b/vortex-layout/src/plan/physical/plans/mod.rs @@ -29,8 +29,8 @@ pub use filter::FilterPlan; pub(crate) use legacy_stats::ExpressionLegacyStatsRule; pub use legacy_stats::LegacyStats; pub use legacy_stats::LegacyStatsData; +pub(crate) use legacy_stats::LegacyStatsDataRule; pub use legacy_stats::LegacyStatsPlan; -pub(crate) use legacy_stats::LegacyStatsRule; pub use list_pack::ListPack; pub use list_pack::ListPackData; pub use list_pack::ListPackPlan; diff --git a/vortex-layout/src/plan/physical/plans/pack.rs b/vortex-layout/src/plan/physical/plans/pack.rs index d3886de19ad..ff5855e0943 100644 --- a/vortex-layout/src/plan/physical/plans/pack.rs +++ b/vortex-layout/src/plan/physical/plans/pack.rs @@ -518,7 +518,10 @@ fn expanded_struct_root( bound_pack(fields.names().clone(), children) } -fn is_identity_expression(expression: &BoundExpression, input_dtype: &DType) -> VortexResult { +pub(super) fn is_identity_expression( + expression: &BoundExpression, + input_dtype: &DType, +) -> VortexResult { if expression.is_root() { return Ok(expression.dtype() == input_dtype); } diff --git a/vortex-layout/src/plan/physical/plans/zoned.rs b/vortex-layout/src/plan/physical/plans/zoned.rs index 15bf8cd7dfd..01b4be577de 100644 --- a/vortex-layout/src/plan/physical/plans/zoned.rs +++ b/vortex-layout/src/plan/physical/plans/zoned.rs @@ -468,6 +468,9 @@ impl PlanParentReduceRule for ExpressionZonedRule { /// Whether `eval` is a boolean expression over abstract statistics alone, which a zone table can /// prune without reading the data it summarises. pub(crate) fn is_zone_pruning(eval: &Plan) -> VortexResult { + if !eval.dtype().is_boolean() { + return Ok(false); + } let mut contains_stat = false; let mut contains_root = false; eval.expression().clone().transform_down(|expression| { @@ -485,7 +488,7 @@ pub(crate) fn is_zone_pruning(eval: &Plan) -> VortexResult { contains_root |= expression.is_root(); Ok(Transformed::no(expression)) })?; - Ok(eval.dtype().is_boolean() && contains_stat && !contains_root) + Ok(contains_stat && !contains_root) } /// Replaces a zoned data plan by its data: a scan reads the data through it and never its zone diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index c86edcf4d30..0c439b058f1 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -648,12 +648,9 @@ fn empty_projection_prunes_row_idx_child_fields() -> VortexResult<()> { ); let plan = make_row_idx_plan(projection, make_plan(layout)?)?; + // The empty projection over the empty row domain is the identity, so only the domain remains. let optimized = optimize(plan, &SESSION)?; - let projection = optimized - .as_opt::() - .ok_or_else(|| vortex_err!("optimized plan has no projection expression"))?; - let child = projection.child_plan()?; - let empty_struct = child + let empty_struct = optimized .as_opt::() .ok_or_else(|| vortex_err!("empty projection did not prune the RowIdx child"))?; diff --git a/vortex-layout/src/scan/execution/adaptor.rs b/vortex-layout/src/scan/execution/adaptor.rs index fc6053d55ca..85b834bb74d 100644 --- a/vortex-layout/src/scan/execution/adaptor.rs +++ b/vortex-layout/src/scan/execution/adaptor.rs @@ -3,7 +3,6 @@ //! Adapts pipeline execution to the planner I/O and compute protocol. -use std::ops::Range; use std::sync::Arc; use rustc_hash::FxHashMap; @@ -25,11 +24,11 @@ use crate::plan::physical::pipeline::Turn; use crate::scan::execution::SegmentLocation; use crate::segments::SegmentId; -/// Rows a running plan produced, in row order. +/// Rows a running plan produced, in row order within its split. #[derive(Clone)] pub(super) struct Piece { - /// Orders the piece among the outputs of one scan. - pub(super) rows: Range, + /// The index of the split the rows belong to. + pub(super) split: usize, /// The selected rows the piece holds. pub(super) array: ArrayRef, } @@ -54,8 +53,6 @@ pub(super) struct ScanAdaptor { /// A turn already taken while batching reads, retained for the next compute. pending: Option, done: bool, - /// Orders the pieces handed out. - next_piece: u64, locations: Arc<[SegmentLocation]>, /// Published and undelivered requests, with the scan read and segment each one answers. outstanding: FxHashMap, @@ -72,7 +69,6 @@ impl ScanAdaptor { scan, pending: None, done: false, - next_piece: 0, locations, outstanding: FxHashMap::default(), failed: None, @@ -96,7 +92,7 @@ impl ScanAdaptor { } Ok(match self.step()? { Turn::Waiting | Turn::Done => AdaptorStep::Yield, - Turn::Output(_, array) => AdaptorStep::Piece(self.piece(array)), + Turn::Output(split, array) => AdaptorStep::Piece(Piece { split, array }), Turn::Read(read) => { // Publish every read the scan has ready together, so the owner can coalesce them. let mut batch = vec![self.request(read.id, read.segment_id)?]; @@ -130,15 +126,6 @@ impl ScanAdaptor { Ok(turn) } - fn piece(&mut self, array: ArrayRef) -> Piece { - let start = self.next_piece; - self.next_piece += 1; - Piece { - rows: start..start + 1, - array, - } - } - fn request(&mut self, read: ReadId, segment_id: SegmentId) -> VortexResult { let location = self.location(segment_id)?; let id = self.next_id; diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs index 382425d9bb0..3128b2b93cc 100644 --- a/vortex-layout/src/scan/execution/chunk.rs +++ b/vortex-layout/src/scan/execution/chunk.rs @@ -44,6 +44,29 @@ pub(crate) fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> u blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE } +/// The rows of a filtered chunk for a scan of `rows` running `concurrency` chunks at once: as +/// large as [`FILTERED_ROWS_PER_CHUNK`], but small enough that every chunk the scan can run at +/// once has rows to filter, and never below [`MAX_ROWS_PER_CHUNK`]. +pub(crate) fn filtered_rows_per_chunk(rows: u64, concurrency: usize) -> usize { + rows_per_concurrent_chunk(rows, concurrency).clamp(MAX_ROWS_PER_CHUNK, FILTERED_ROWS_PER_CHUNK) +} + +/// The rows of a projection chunk for a scan of `rows` running `concurrency` chunks at once: as +/// large as `max_rows`, but small enough that every chunk the scan can run at once has rows to +/// project, and never below [`ROW_BLOCK_SIZE`]. +pub(crate) fn projected_rows_per_chunk(rows: u64, concurrency: usize, max_rows: usize) -> usize { + rows_per_concurrent_chunk(rows, concurrency).clamp(ROW_BLOCK_SIZE, max_rows) +} + +/// `rows` split evenly over `concurrency` chunks, rounded up to the row grid. +fn rows_per_concurrent_chunk(rows: u64, concurrency: usize) -> usize { + let per_chunk = rows.div_ceil(concurrency.max(1) as u64); + let blocks = per_chunk.div_ceil(ROW_BLOCK_SIZE as u64); + usize::try_from(blocks) + .unwrap_or(usize::MAX) + .saturating_mul(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() { diff --git a/vortex-layout/src/scan/execution/filter.rs b/vortex-layout/src/scan/execution/filter.rs index d3d7b640892..69b0999482a 100644 --- a/vortex-layout/src/scan/execution/filter.rs +++ b/vortex-layout/src/scan/execution/filter.rs @@ -130,8 +130,7 @@ pub(super) struct FilterPlanner { selected: Option, filter: Arc, remaining: BitVec, - active: Option<(usize, ScanAdaptor)>, - pieces: Vec, + active: Option<(usize, MaskScan)>, announced: bool, /// Ids for hints, counting down from `u32::MAX` so they never meet an adaptor's fetch ids. next_hint: u32, @@ -163,7 +162,6 @@ impl FilterPlanner { selected: Some(selected), filter, active: None, - pieces: Vec::new(), announced: false, next_hint: u32::MAX, } @@ -193,15 +191,72 @@ pub(super) fn hint( } /// Executes a scan's boolean piece to a mask, nulls as false. -pub(super) fn piece_mask(session: &VortexSession, piece: Piece) -> VortexResult { +fn piece_mask(session: &VortexSession, piece: Piece) -> VortexResult { let mut ctx = session.create_execution_ctx(); piece.array.fill_null(false)?.execute::(&mut ctx) } +/// A boolean plan running over one split, its pieces gathered into one mask. +pub(super) struct MaskScan { + adaptor: ScanAdaptor, + pieces: Vec, +} + +/// What one [`MaskScan::compute`] produced. +pub(super) enum MaskStep { + /// The scan is still running: its owner's compute returns this. + Running(PlannerOutput), + /// The mask over the split's selected rows, nulls as false. + Finished(Mask), +} + +impl MaskScan { + pub(super) fn new(scan: &PreparedScan, plan: PlanRef, split: Split) -> VortexResult { + let pipeline = Scan::try_new(scan.session.clone(), plan, vec![split])? + .with_row_offset(scan.row_offset) + .with_value_cache(scan.values.clone()); + Ok(Self { + adaptor: ScanAdaptor::new(pipeline, Arc::clone(&scan.locations)), + pieces: Vec::new(), + }) + } + + /// The owner's state while the scan runs: once it has finished, its mask is ready to take. + pub(super) fn state(&self) -> State { + match self.adaptor.state() { + State::Done => State::NeedsCompute, + state => state, + } + } + + pub(super) fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { + self.adaptor.set_io_result(request, result); + } + + pub(super) fn compute(&mut self, session: &VortexSession) -> VortexResult { + if self.adaptor.state() == State::Done { + // A single piece is already the mask; concatenating would copy its bits. + let mask = match self.pieces.as_slice() { + [mask] => mask.clone(), + pieces => Mask::concat(pieces.iter())?, + }; + return Ok(MaskStep::Finished(mask)); + } + Ok(MaskStep::Running(match self.adaptor.compute()? { + AdaptorStep::Yield => PlannerOutput::Continue, + AdaptorStep::NeedsIO(batch) => PlannerOutput::NeedsIO(batch), + AdaptorStep::Piece(piece) => { + self.pieces.push(piece_mask(session, piece)?); + PlannerOutput::Continue + } + })) + } +} + impl IoConsumer for FilterPlanner { fn set_io_result(&mut self, request: IoRequestId, result: IoResult) { - if let Some((_, adaptor)) = &mut self.active { - adaptor.set_io_result(request, result); + if let Some((_, conjunct)) = &mut self.active { + conjunct.set_io_result(request, result); } } } @@ -212,8 +267,8 @@ impl Planner for FilterPlanner { return State::Done; } match &self.active { - Some((_, adaptor)) if adaptor.state() != State::Done => adaptor.state(), - _ => State::NeedsCompute, + Some((_, conjunct)) => conjunct.state(), + None => State::NeedsCompute, } } @@ -260,19 +315,12 @@ impl Planner for FilterPlanner { return Ok(PlannerOutput::NeedsIO(batch)); } } - 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) => { - self.pieces.push(piece_mask(&self.scan.session, piece)?); - PlannerOutput::Continue - } - }); - } + if let Some((index, conjunct)) = &mut self.active { + let result = match conjunct.compute(&self.scan.session)? { + MaskStep::Running(output) => return Ok(output), + MaskStep::Finished(result) => result, + }; let input = selected.mask.true_count(); - let result = Mask::concat(self.pieces.iter())?; vortex_ensure_eq!(result.len(), input); selected.mask = selected.mask.intersect_by_rank(&result); tracing::debug!( @@ -287,7 +335,6 @@ impl Planner for FilterPlanner { .expression .report_selectivity(*index, selected.mask.true_count() as f64 / input as f64); self.active = None; - self.pieces.clear(); } if selected.mask.all_false() { self.selected = None; @@ -295,20 +342,15 @@ impl Planner for FilterPlanner { } if let Some(index) = self.filter.expression.next_conjunct(&self.remaining) { self.remaining.set(index, false); - let scan = Scan::try_new( - self.scan.session.clone(), + let conjunct = MaskScan::new( + &self.scan, self.filter.predicates[index].clone(), - vec![Split { + Split { rows: selected.scope.rows.clone(), mask: selected.mask.clone(), - }], - )? - .with_row_offset(self.scan.row_offset) - .with_value_cache(self.scan.values.clone()); - self.active = Some(( - index, - ScanAdaptor::new(scan, Arc::clone(&self.scan.locations)), - )); + }, + )?; + self.active = Some((index, conjunct)); return Ok(PlannerOutput::Continue); } let selected = self.selected.take().vortex_expect("selected rows exist"); diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 0185f61d6a9..0a9ff0b7679 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -31,9 +31,10 @@ use vortex_scan::selection::Selection; use vortex_session::VortexSession; pub 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; +pub(crate) use self::chunk::filtered_rows_per_chunk; +pub(crate) use self::chunk::projected_rows_per_chunk; #[cfg(test)] pub(crate) use self::chunk::row_chunks; pub(crate) use self::chunk::rows_per_chunk; diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index 141b7377c26..f3da4b4218c 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -2,11 +2,12 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::sync::Arc; -use std::vec::IntoIter; +use vortex_array::ArrayRef; 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; @@ -18,7 +19,6 @@ use vortex_scan::planning::morsel::MorselOutput; 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 crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Split; @@ -26,46 +26,23 @@ use crate::plan::physical::selection::selected; use crate::scan::execution::PreparedScan; use crate::scan::execution::SelectedRows; use crate::scan::execution::adaptor::AdaptorStep; -use crate::scan::execution::adaptor::Piece; use crate::scan::execution::adaptor::ScanAdaptor; +use crate::scan::execution::chunk::row_chunks; 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. +/// Emits one morsel projecting a chunk's selected rows. pub(super) struct ProjectionPlanner { scan: Arc, - selected: IntoIter, + selected: Option, } impl ProjectionPlanner { /// Creates a planner for `selected`. pub(super) fn new(scan: Arc, selected: SelectedRows) -> Self { - let selected = if selected.mask.true_count() <= MAX_OUTPUT_ROWS { - vec![selected] - } else { - (0..selected.mask.len()) - .step_by(MAX_OUTPUT_ROWS) - .filter_map(|start| { - let end = (start + MAX_OUTPUT_ROWS).min(selected.mask.len()); - let mask = selected.mask.slice(start..end); - if mask.all_false() { - return None; - } - Some(SelectedRows { - scope: WorkScope { - file_ordinal: selected.scope.file_ordinal, - rows: selected.scope.rows.start + start as u64 - ..selected.scope.rows.start + end as u64, - }, - mask, - }) - }) - .collect() - }; Self { scan, - selected: selected.into_iter(), + selected: Some(selected), } } } @@ -76,7 +53,7 @@ impl IoConsumer for ProjectionPlanner { impl Planner for ProjectionPlanner { fn state(&self) -> State { - if !self.selected.as_slice().is_empty() { + if self.selected.is_some() { State::NeedsCompute } else { State::Done @@ -84,7 +61,7 @@ impl Planner for ProjectionPlanner { } fn compute(&mut self) -> VortexResult { - let Some(selected) = self.selected.next() else { + let Some(selected) = self.selected.take() else { return Ok(PlannerOutput::Done); }; if selected.mask.all_false() { @@ -97,14 +74,25 @@ 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. +/// Runs the projection plan over a chunk's selected rows on one scan, so the chunk compiles one +/// plan and shares its reads and dictionary values, and yields a batch per window of the rows, +/// in row order. +/// +/// Sparse selections stay one window. Dense selections are cut into windows of +/// `MAX_OUTPUT_ROWS` on a fixed grid of the file's rows, each a split of the scan, so kernels run +/// over a window at a time and converting a batch to Arrow copies no more than a window. struct ProjectionMorsel { scan: Arc, dtype: DType, selected: SelectedRows, adaptor: Option, - pieces: Vec, + /// The rows each split projects to, in split order. + split_rows: Vec, + /// The split whose pieces `pieces` holds. + split: usize, + pieces: Vec, + /// The rows `pieces` holds. + held: usize, done: bool, } @@ -116,25 +104,54 @@ impl ProjectionMorsel { scan, selected, adaptor: None, + split_rows: Vec::new(), + split: 0, pieces: Vec::new(), + held: 0, done: false, } } - fn finish(&mut self) -> VortexResult { - self.done = true; - self.pieces.sort_by_key(|piece| piece.rows.start); + /// The windows of the selected rows that select anything, each a split of the scan, with + /// the rows each projects to. + fn splits(&self) -> (Vec, Vec) { + let rows = &self.selected.scope.rows; + let mask = &self.selected.mask; + if mask.true_count() <= MAX_OUTPUT_ROWS { + let split = Split { + rows: rows.clone(), + mask: Mask::new_true(mask.len()), + }; + return (vec![split], vec![mask.true_count()]); + } + let offset = + |row: u64| usize::try_from(row - rows.start).vortex_expect("a work chunk fits usize"); + row_chunks(rows.clone(), MAX_OUTPUT_ROWS) + .into_iter() + .filter_map(|window| { + let (start, end) = (offset(window.start), offset(window.end)); + let selected = mask.slice(start..end).true_count(); + let split = Split { + rows: window, + mask: Mask::new_true(end - start), + }; + (selected > 0).then_some((split, selected)) + }) + .unzip() + } + + /// The pieces of the split held so far as one batch, or `None` when they hold no rows. + fn take_batch(&mut self) -> VortexResult> { + self.held = 0; let mut arrays: Vec<_> = std::mem::take(&mut self.pieces) .into_iter() - .map(|piece| piece.array) .filter(|array| !array.is_empty()) .collect(); - let array = match arrays.len() { - 0 => return Ok(MorselOutput::Done), - 1 => arrays.remove(0), - _ => ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array(), - }; - Ok(MorselOutput::Batch(array)) + Ok(match arrays.len() { + 0 => None, + 1 => Some(arrays.remove(0)), + _ => Some(ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array()), + }) } } @@ -163,6 +180,8 @@ impl Morsel for ProjectionMorsel { vortex_bail!("ProjectionMorsel: compute called after Done"); } let Some(adaptor) = self.adaptor.as_mut() else { + let (splits, split_rows) = self.splits(); + self.split_rows = split_rows; let scan = Scan::try_new( self.scan.session.clone(), selected( @@ -170,10 +189,7 @@ impl Morsel for ProjectionMorsel { self.selected.scope.rows.clone(), self.selected.mask.clone(), ), - vec![Split { - rows: self.selected.scope.rows.clone(), - mask: Mask::new_true(self.selected.mask.len()), - }], + splits, )? .with_row_offset(self.scan.row_offset) .with_value_cache(self.scan.values.clone()); @@ -181,14 +197,31 @@ impl Morsel for ProjectionMorsel { return Ok(MorselOutput::Continue); }; if adaptor.state() == State::Done { - return self.finish(); + return Ok(match self.take_batch()? { + Some(batch) => MorselOutput::Batch(batch), + None => { + self.done = true; + MorselOutput::Done + } + }); } Ok(match adaptor.compute()? { AdaptorStep::Yield => MorselOutput::Continue, AdaptorStep::NeedsIO(batch) => MorselOutput::NeedsIO(batch), AdaptorStep::Piece(piece) => { - self.pieces.push(piece); - MorselOutput::Continue + // One split runs at a time, so a piece of a later split ends the one before it. + let mut batch = None; + if piece.split != self.split { + batch = self.take_batch()?; + self.split = piece.split; + } + self.held += piece.array.len(); + self.pieces.push(piece.array); + // A split that has projected all its rows yields without waiting for the next. + if batch.is_none() && self.split_rows.get(self.split) == Some(&self.held) { + batch = self.take_batch()?; + } + batch.map_or(MorselOutput::Continue, MorselOutput::Batch) } }) } diff --git a/vortex-layout/src/scan/execution/prune.rs b/vortex-layout/src/scan/execution/prune.rs index 325f5f897d7..e014c0d2bb8 100644 --- a/vortex-layout/src/scan/execution/prune.rs +++ b/vortex-layout/src/scan/execution/prune.rs @@ -22,13 +22,11 @@ use vortex_scan::planning::planner::State; use vortex_scan::planning::planner::WorkScope; use super::PreparedScan; -use super::adaptor::AdaptorStep; -use super::adaptor::ScanAdaptor; +use super::filter::MaskScan; +use super::filter::MaskStep; use super::filter::hint; -use super::filter::piece_mask; use super::filter::segments_read; use crate::plan::physical::PlanRef; -use crate::plan::physical::pipeline::Scan; use crate::plan::physical::pipeline::Split; /// The planner a prune planner continues into, given the rows pruning kept, or `None` when @@ -49,8 +47,7 @@ pub struct PrunePlanner { /// The rows no pruning plan has excluded yet, over `rows`. mask: Mask, pruning: IntoIter, - active: Option, - pieces: Vec, + active: Option, announced: bool, /// Ids for hints, counting down from `u32::MAX` so they never meet an adaptor's fetch ids. next_hint: u32, @@ -74,7 +71,6 @@ impl PrunePlanner { mask: Mask::new_true(len), pruning: pruning.into_iter(), active: None, - pieces: Vec::new(), announced: false, next_hint: u32::MAX, done: false, @@ -84,8 +80,8 @@ impl PrunePlanner { 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); + if let Some(pruning) = &mut self.active { + pruning.set_io_result(request, result); } } } @@ -96,8 +92,8 @@ impl Planner for PrunePlanner { return State::Done; } match &self.active { - Some(adaptor) if adaptor.state() != State::Done => adaptor.state(), - _ => State::NeedsCompute, + Some(pruning) => pruning.state(), + None => State::NeedsCompute, } } @@ -129,19 +125,12 @@ impl Planner for PrunePlanner { 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 - } - }); - } + if let Some(pruning) = &mut self.active { + let excluded = match pruning.compute(&self.scan.session)? { + MaskStep::Running(output) => return Ok(output), + MaskStep::Finished(excluded) => excluded, + }; 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!( @@ -152,7 +141,6 @@ impl Planner for PrunePlanner { "v2 pruning plan evaluated" ); self.active = None; - self.pieces.clear(); self.pruned = true; } if self.mask.all_false() { @@ -160,17 +148,14 @@ impl Planner for PrunePlanner { return Ok(PlannerOutput::Done); } if let Some(plan) = self.pruning.next() { - let scan = Scan::try_new( - self.scan.session.clone(), + self.active = Some(MaskScan::new( + &self.scan, plan, - vec![Split { + 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 pruning plan has run: the planner this one continues into gets the kept rows. diff --git a/vortex-layout/src/scan/v2/cache.rs b/vortex-layout/src/scan/v2/cache.rs index c621bf19289..4e946cbc8b9 100644 --- a/vortex-layout/src/scan/v2/cache.rs +++ b/vortex-layout/src/scan/v2/cache.rs @@ -6,10 +6,12 @@ use std::sync::Arc; use std::sync::Weak; +use once_cell::sync::OnceCell; use parking_lot::Mutex; use rustc_hash::FxHashMap; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; +use vortex_error::VortexResult; use super::io::SegmentRanges; use crate::plan::physical::PlanRef; @@ -35,81 +37,62 @@ pub struct ScanCache(Arc); #[derive(Default)] struct Inner { - root: Mutex>, - /// Optimized plans of non-identity projections. Values are cached by the plan that produced - /// them, so a projection must reuse its plan for its values to be found again. - projections: Mutex>, - ranges: Mutex>, + root: OnceCell, + /// Optimized plans of projections. Values are cached by the plan that produced them, so a + /// projection must reuse its plan for its values to be found again. + projections: Mutex>>>, + ranges: OnceCell, /// Plans of filters, keyed by the exact filter expression. A pruning plan remembers the zones it /// excluded, so sharing the plans prunes a file's zones once for every scan of that filter. - filters: Mutex>>, + filters: Mutex>>>>, values: ValueCache, } impl ScanCache { - /// The lowered root plan an earlier scan kept. - pub(super) fn root(&self) -> Option { - self.0.root.lock().clone() - } - - /// Keeps `root`, returning the plan already kept if another scan got there first. - pub(super) fn insert_root(&self, root: PlanRef) -> PlanRef { - self.0.root.lock().get_or_insert(root).clone() + /// The lowered root plan, lowered by `lower` unless an earlier scan kept it. Scans that miss + /// together wait for one of them to lower it. + pub(super) fn root_or_try_init( + &self, + lower: impl FnOnce() -> VortexResult, + ) -> VortexResult { + self.0.root.get_or_try_init(lower).cloned() } - /// The optimized plan of `projection` an earlier scan kept. - pub(super) fn projection(&self, projection: &BoundExpression) -> Option { - self.0 - .projections - .lock() - .get(&ExactBoundExpr(projection.clone())) + /// The optimized plan of `projection`, planned by `plan` unless an earlier scan kept it. + pub(super) fn projection_or_try_init( + &self, + projection: &BoundExpression, + plan: impl FnOnce() -> VortexResult, + ) -> VortexResult { + cell(&self.0.projections, projection) + .get_or_try_init(plan) .cloned() } - /// Keeps the optimized plan of `projection`, returning the plan already kept if another scan - /// got there first. - pub(super) fn insert_projection(&self, projection: BoundExpression, plan: PlanRef) -> PlanRef { - self.0 - .projections - .lock() - .entry(ExactBoundExpr(projection)) - .or_insert(plan) - .clone() - } - - /// The segment byte-range index an earlier scan kept. - pub(super) fn ranges(&self) -> Option { - self.0.ranges.lock().clone() - } - - /// Keeps `ranges`, returning the index already kept if another scan got there first. - pub(super) fn insert_ranges(&self, ranges: SegmentRanges) -> SegmentRanges { - Arc::clone(self.0.ranges.lock().get_or_insert(ranges)) + /// The segment byte-range index, built by `index` unless an earlier scan kept it. + pub(super) fn ranges_or_init(&self, index: impl FnOnce() -> SegmentRanges) -> SegmentRanges { + Arc::clone(self.0.ranges.get_or_init(index)) } /// The plans of `filter` an earlier scan kept. + #[cfg(test)] pub(super) fn filter_plans(&self, filter: &BoundExpression) -> Option> { self.0 .filters .lock() .get(&ExactBoundExpr(filter.clone())) - .cloned() + .and_then(|plans| plans.get().cloned()) } - /// Keeps the plans of `filter`, returning the plans already kept if another scan got there - /// first. - pub(super) fn insert_filter_plans( + /// The plans of `filter`, planned by `plan` unless an earlier scan kept them. + pub(super) fn filter_plans_or_try_init( &self, - filter: BoundExpression, - plans: Arc, - ) -> Arc { - Arc::clone( - self.0 - .filters - .lock() - .entry(ExactBoundExpr(filter)) - .or_insert(plans), - ) + filter: &BoundExpression, + plan: impl FnOnce() -> VortexResult, + ) -> VortexResult> { + cell(&self.0.filters, filter) + .get_or_try_init(|| plan().map(Arc::new)) + .cloned() } pub(super) fn values(&self) -> ValueCache { @@ -122,6 +105,14 @@ impl ScanCache { } } +/// The cell of `key` in `map`, so scans that miss together build its entry once. +fn cell( + map: &Mutex>>>, + key: &BoundExpression, +) -> Arc> { + Arc::clone(map.lock().entry(ExactBoundExpr(key.clone())).or_default()) +} + /// 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. diff --git a/vortex-layout/src/scan/v2/io.rs b/vortex-layout/src/scan/v2/io.rs index 0e690ebfde7..abb24bdebf1 100644 --- a/vortex-layout/src/scan/v2/io.rs +++ b/vortex-layout/src/scan/v2/io.rs @@ -87,7 +87,7 @@ pub(super) struct SegmentIoSource { announced: Mutex>>, } -/// A read started ahead of the fetch that needs it. +/// A segment read the session shares among every request of the segment. type Prefetched = Shared>>>; impl SegmentIoSource { @@ -101,71 +101,69 @@ impl SegmentIoSource { announced: Mutex::default(), } } + + /// Registers a read of `segment`, taking over an announcement of it if there is one. + fn read(&self, segment: SegmentId) -> Prefetched { + self.announced + .lock() + .get_mut(&segment) + .and_then(VecDeque::pop_front) + .unwrap_or_else(|| self.segments.request(segment)) + .map_err(Arc::new) + .boxed() + .shared() + } } impl IoSource for SegmentIoSource { fn submit(&self, owner: IoOwnerId, batch: IoBatch) -> VortexResult<()> { - let reads = self.reads.lock(); - let mut ready = self.ready.lock(); - let mut prefetched = self.prefetched.lock(); - for request in batch { - let IoTarget::Range { offset, len, .. } = request.target else { - vortex_bail!( - "SegmentIoSource only serves byte ranges, not {:?}", - request.target - ); - }; - let segment = *self - .ranges - .get(&(offset, len)) - .ok_or_else(|| vortex_err!("No segment at bytes {offset}+{len}"))?; - if request.intent == IoIntent::Announce { - self.announced - .lock() - .entry(segment) - .or_default() - .push_back(self.segments.request(segment)); - continue; - } - if request.intent == IoIntent::Prefetch { - prefetched.entry(segment).or_insert_with(|| { - let read = self - .announced - .lock() - .get_mut(&segment) - .and_then(VecDeque::pop_front) - .unwrap_or_else(|| self.segments.request(segment)) - .map_err(Arc::new) - .boxed() - .shared(); - // Polling registers the read as wanted, so the source starts it now. - drop( - read.clone() - .poll_unpin(&mut Context::from_waker(Waker::noop())), + let mut started = Vec::with_capacity(batch.len()); + { + let mut prefetched = self.prefetched.lock(); + for request in batch { + let IoTarget::Range { offset, len, .. } = request.target else { + vortex_bail!( + "SegmentIoSource only serves byte ranges, not {:?}", + request.target ); - read - }); - continue; - } - // 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(|| { + }; + let segment = *self + .ranges + .get(&(offset, len)) + .ok_or_else(|| vortex_err!("No segment at bytes {offset}+{len}"))?; + if request.intent == IoIntent::Announce { 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() + .entry(segment) + .or_default() + .push_back(self.segments.request(segment)); + continue; + } + // 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 read = prefetched + .entry(segment) + .or_insert_with(|| self.read(segment)) + .clone(); + let fetch = (request.intent == IoIntent::Fetch).then_some(request.request); + started.push((segment, read, fetch)); + } + } + + // Every read of the batch is registered before any is polled, so a coalescing source + // sees the whole batch, announcements included, before it picks a read to issue. + let reads = self.reads.lock(); + let mut ready = self.ready.lock(); + for (segment, mut read, fetch) in started { + let Some(id) = fetch else { + // Polling registers the read as wanted, so the source starts it now. + drop(read.poll_unpin(&mut Context::from_waker(Waker::noop()))); + continue; + }; + let mut bytes = read .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. if let Poll::Ready(result) = bytes.poll_unpin(&mut Context::from_waker(Waker::noop())) { @@ -235,6 +233,7 @@ mod tests { use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; + use std::task::Poll; use futures::FutureExt; use futures::channel::oneshot; @@ -298,6 +297,64 @@ mod tests { } } + /// Logs each segment's registration and first poll; its reads never complete. + #[derive(Default)] + struct LoggingSegments(Arc>>); + + impl SegmentSource for LoggingSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.0.lock().push(format!("register {id}")); + let log = Arc::clone(&self.0); + let mut polled = false; + future::poll_fn(move |_| { + if !std::mem::replace(&mut polled, true) { + log.lock().push(format!("poll {id}")); + } + Poll::Pending + }) + .boxed() + } + } + + /// A coalescing source sees every read of a batch, announcements included, before any of + /// them is polled and so made eligible to issue. + #[test] + fn batch_registers_every_read_before_polling() -> VortexResult<()> { + let segments = LoggingSegments::default(); + let log = Arc::clone(&segments.0); + let locations = [0, 4, 8].map(|offset| SegmentLocation { + offset, + length: 4, + alignment: Alignment::none(), + }); + let io = SegmentIoSource::new(Arc::new(segments), segment_ranges(&locations)); + let batch = [ + (IoIntent::Prefetch, 0), + (IoIntent::Fetch, 4), + (IoIntent::Announce, 8), + ] + .into_iter() + .enumerate() + .map(|(id, (intent, offset))| IoRequest { + intent, + request: IoRequestId(u32::try_from(id).unwrap_or(u32::MAX)), + target: IoTarget::range(offset, 4), + }) + .collect(); + io.submit(IoOwnerId(0), batch)?; + assert_eq!( + *log.lock(), + [ + "register SegmentId(0)", + "register SegmentId(1)", + "register SegmentId(2)", + "poll SegmentId(0)", + "poll SegmentId(1)", + ] + ); + Ok(()) + } + #[tokio::test] async fn pending_read_wakes_after_an_immediate_completion() -> VortexResult<()> { let (sender, receiver) = oneshot::channel(); diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index f0ad78034ae..6eb2a6381db 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -36,6 +36,7 @@ use vortex_error::VortexResult; use vortex_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use vortex_session::VortexSession; +use vortex_utils::parallelism::get_available_parallelism; pub use self::cache::ScanCache; pub use self::cache::WeakScanCache; @@ -44,7 +45,6 @@ use crate::plan::physical::lower; use crate::plan::physical::optimize; use crate::plan::physical::plan_row_idx_expression; pub use crate::scan::execution::ChunkPlanner; -use crate::scan::execution::FILTERED_ROWS_PER_CHUNK; pub use crate::scan::execution::FilterPlans; use crate::scan::execution::MAX_ROWS_PER_CHUNK; pub use crate::scan::execution::PreparedScan; @@ -52,6 +52,8 @@ pub use crate::scan::execution::PruneNext; pub use crate::scan::execution::PrunePlanner; use crate::scan::execution::ROW_BLOCK_SIZE; pub use crate::scan::execution::SegmentLocation; +use crate::scan::execution::filtered_rows_per_chunk; +use crate::scan::execution::projected_rows_per_chunk; use crate::scan::execution::rows_per_chunk; use crate::scan::v2::io::segment_ranges; use crate::segments::SegmentSource; @@ -96,6 +98,7 @@ pub struct ScanBuilder { filter: Option, row_range: Option>, selection: Selection, + concurrency: usize, row_offset: u64, } @@ -109,6 +112,7 @@ impl ScanBuilder { file, row_range: None, selection: Selection::default(), + concurrency: get_available_parallelism().unwrap_or(1), row_offset: 0, } } @@ -152,6 +156,15 @@ impl ScanBuilder { self.with_selection(Selection::IncludeByIndex(row_indices)) } + /// Sets the number of the scan's chunks its driver runs at once, one per worker thread by + /// default, which sizes a filtered scan's chunks so each of them has rows to filter. Must be + /// greater than zero. + pub fn with_concurrency(mut self, concurrency: usize) -> Self { + assert!(concurrency > 0); + self.concurrency = concurrency; + self + } + /// Sets the root row offset used by row-index expressions. pub fn with_row_offset(mut self, row_offset: u64) -> Self { self.row_offset = row_offset; @@ -164,52 +177,42 @@ impl ScanBuilder { pub fn prepare(self) -> VortexResult { let file = self.file; let cache = file.cache.unwrap_or_default(); - let root = match cache.root() { - Some(root) => root, - None => cache.insert_root(lower(&file.layout)?), - }; + let root = cache.root_or_try_init(|| lower(&file.layout))?; let row_count = root.row_count(); + let session = &self.session; let filter = self .filter - .map(|filter| -> VortexResult<_> { - match cache.filter_plans(&filter) { - Some(plans) => Ok(plans), - None => { - let plans = - Arc::new(FilterPlans::new(filter.clone(), &root, &self.session)?); - Ok(cache.insert_filter_plans(filter, plans)) - } - } + .map(|filter| { + cache.filter_plans_or_try_init(&filter, || { + FilterPlans::new(filter.clone(), &root, session) + }) }) .transpose()?; // 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 = self.projection.is_root() && self.projection.dtype() == root.dtype(); - let projection = match cache.projection(&self.projection) { - Some(plan) => plan, - None => { - let plan = optimize( - plan_row_idx_expression(self.projection.clone(), root)?, - &self.session, - )?; - cache.insert_projection(self.projection, plan) - } - }; + let projection = cache.projection_or_try_init(&self.projection, || { + optimize( + plan_row_idx_expression(self.projection.clone(), root.clone())?, + session, + ) + })?; + let rows = self.row_range.unwrap_or(0..row_count); let chunk_rows = if filter.is_some() { - FILTERED_ROWS_PER_CHUNK - } else if identity { - rows_per_chunk(row_count, &file.locations) + filtered_rows_per_chunk(rows.end - rows.start, self.concurrency) } else { - MAX_ROWS_PER_CHUNK - }; - let ranges = match cache.ranges() { - Some(ranges) => ranges, - None => cache.insert_ranges(segment_ranges(&file.locations)), + let max_rows = if identity { + rows_per_chunk(row_count, &file.locations) + } else { + MAX_ROWS_PER_CHUNK + }; + projected_rows_per_chunk(rows.end - rows.start, self.concurrency, max_rows) }; + let ranges = cache.ranges_or_init(|| segment_ranges(&file.locations)); let values = cache.values(); Ok(PreparedScan { session: self.session, - rows: self.row_range.unwrap_or(0..row_count), + rows, selection: self.selection, segments: file.segments, ranges, diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index eff68dfe996..9015e37b9a9 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -72,6 +72,8 @@ use crate::layouts::zoned::writer::ZonedLayoutOptions; use crate::layouts::zoned::writer::ZonedStrategy; use crate::scan::arrow::to_record_batch; use crate::scan::execution::SegmentLocation; +use crate::scan::execution::filtered_rows_per_chunk; +use crate::scan::execution::projected_rows_per_chunk; use crate::scan::execution::row_chunks; use crate::scan::scan_builder::ScanBuilder as V1ScanBuilder; use crate::segments::SegmentFuture; @@ -495,6 +497,31 @@ fn chunk_byte_budget(#[case] rows: u64, #[case] bytes: u32, #[case] expected: us assert_eq!(rows_per_chunk(rows, &locations), expected); } +/// A filtered chunk shrinks so every chunk a scan runs at once has rows, within its bounds. +#[rstest] +#[case::one_at_a_time(6_000_000, 1, 1_048_576)] +#[case::split_across_concurrency(262_144, 4, 65_536)] +#[case::rounded_to_blocks(300_000, 4, 81_920)] +#[case::at_least_a_window(100_000, 96, 65_536)] +fn filtered_chunk_rows(#[case] rows: u64, #[case] concurrency: usize, #[case] expected: usize) { + assert_eq!(filtered_rows_per_chunk(rows, concurrency), expected); +} + +/// A projection chunk shrinks so every chunk a scan runs at once has rows, within its bounds: +/// the 1.5M rows of TPC-H SF10's `customer` cut into 6 partitions of 16 chunks each must not +/// run 4 chunks of 64Ki rows each. +#[rstest] +#[case::one_at_a_time(6_000_000, 1, 65_536)] +#[case::split_across_concurrency(250_000, 16, 16_384)] +#[case::rounded_to_blocks(100_000, 4, 32_768)] +#[case::at_least_a_block(10_000, 96, 8_192)] +fn projected_chunk_rows(#[case] rows: u64, #[case] concurrency: usize, #[case] expected: usize) { + assert_eq!( + projected_rows_per_chunk(rows, concurrency, MAX_ROWS_PER_CHUNK), + expected + ); +} + #[test] fn smaller_chunks_keep_file_alignment() { assert_eq!( @@ -550,7 +577,12 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { reads: AtomicUsize::new(0), }); let segments = Arc::clone(&counted) as Arc; - let scan = Arc::new(builder(&session, &segments, &layout)?.prepare()?); + // One chunk at a time, so chunks are as large as a projection chunk on any machine. + let scan = Arc::new( + builder(&session, &segments, &layout)? + .with_concurrency(1) + .prepare()?, + ); let chunks = row_chunks(scan.rows().clone(), scan.chunk_rows); assert_eq!( chunks.len(), @@ -618,6 +650,52 @@ async fn every_row_is_returned() -> VortexResult<()> { Ok(()) } +/// A dense chunk over one stored segment is projected by one morsel, a batch per output window. +#[tokio::test] +async fn dense_chunk_yields_a_batch_per_window() -> VortexResult<()> { + const ROWS: i32 = 150_000; + let session = new_session().with_tokio(); + let segments = Arc::new(TestSegments::default()); + let (mut sequence, eof) = SequenceId::root().split(); + let array = PrimitiveArray::from_iter(0..ROWS).into_array(); + let dtype = array.dtype().clone(); + let layout = chunked_flat() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(dtype, stream::iter([Ok((sequence.advance(), array))])) + .sendable(), + eof, + &session, + ) + .await?; + let segments: Arc = segments; + let mut scan = builder(&session, &segments, &layout)?.prepare()?; + scan.chunk_rows = usize::try_from(ROWS)?; + let scan = Arc::new(scan); + let batches = run_batches(&scan, 0..layout.row_count()).await?; + assert_eq!( + batches + .iter() + .map(|batch| batch.array.len()) + .collect::>(), + [65_536, 65_536, 18_928] + ); + let actual = ChunkedArray::try_new( + batches + .into_iter() + .map(|batch| batch.array) + .collect::>(), + layout.dtype().clone(), + )?; + assert_arrays_eq!( + actual, + PrimitiveArray::from_iter(0..ROWS), + &mut session.create_execution_ctx() + ); + Ok(()) +} + #[tokio::test] async fn primitive_column() -> VortexResult<()> { let session = new_session().with_tokio();