diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs index 640a2d25ccb..5d7b017dfb1 100644 --- a/vortex-layout/src/scan/execution/chunk.rs +++ b/vortex-layout/src/scan/execution/chunk.rs @@ -5,6 +5,7 @@ use std::ops::Range; use std::sync::Arc; +use std::sync::LazyLock; use std::vec::IntoIter; use vortex_error::VortexResult; @@ -20,7 +21,9 @@ use vortex_scan::selection::Selection; use super::PreparedScan; use super::SegmentLocation; +use super::filter::segments_read; use super::selection::plan_selection; +use super::selection::selected_rows; /// The row grid every chunk boundary lies on. pub(crate) const ROW_BLOCK_SIZE: usize = 8192; @@ -68,6 +71,7 @@ pub struct ChunkPlanner { selection: Selection, /// The rows zone pruning kept, over `rows`, or `None` when nothing was pruned. kept: Option, + announced: bool, } impl ChunkPlanner { @@ -86,8 +90,50 @@ impl ChunkPlanner { chunks, selection, kept, + announced: false, } } + + /// Announces the filter's and the projection's segments over every chunk still to plan. + fn announce(&self) -> VortexResult<()> { + let kept = self.kept.as_ref().map(|kept| (kept, self.rows.start)); + let predicates = self + .scan + .filter + .iter() + .flat_map(|filter| filter.predicates()); + let plans: Vec<_> = predicates + .chain(std::iter::once(&self.scan.projection)) + .collect(); + let mut segments = Vec::new(); + for chunk in self.chunks.as_slice() { + let scope = WorkScope { + file_ordinal: 0, + rows: chunk.clone(), + }; + let Some(selected) = selected_rows(scope, &self.selection, kept) else { + continue; + }; + segments.extend(segments_read( + &self.scan.session, + plans.iter().copied(), + selected.scope.rows, + &selected.mask, + )?); + } + segments.sort_unstable(); + segments.dedup(); + self.scan.announce(&segments) + } +} + +/// Whether a scan announces every chunk's reads before planning the first. Setting +/// `VORTEX_SCAN_EARLY_ANNOUNCE=0` leaves each chunk to announce its own, for comparisons. +fn early_announce() -> bool { + static ENABLED: LazyLock = LazyLock::new(|| { + !std::env::var("VORTEX_SCAN_EARLY_ANNOUNCE").is_ok_and(|value| value == "0") + }); + *ENABLED } impl Planner for ChunkPlanner { @@ -104,6 +150,12 @@ impl Planner for ChunkPlanner { } fn compute(&mut self) -> VortexResult { + if !self.announced { + self.announced = true; + if early_announce() { + self.announce()?; + } + } let kept = self.kept.as_ref().map(|kept| (kept, self.rows.start)); for chunk in self.chunks.by_ref() { let scope = WorkScope { diff --git a/vortex-layout/src/scan/execution/filter.rs b/vortex-layout/src/scan/execution/filter.rs index 9de4686bfc7..0eb3c4be0ee 100644 --- a/vortex-layout/src/scan/execution/filter.rs +++ b/vortex-layout/src/scan/execution/filter.rs @@ -97,6 +97,11 @@ impl FilterPlans { pub(crate) fn pruning(&self) -> impl Iterator { self.pruning.iter().flatten() } + + /// The plan of every conjunct. + pub(super) fn predicates(&self) -> impl Iterator { + self.predicates.iter() + } } /// A pruning plan must not fetch data when its requested statistics are unavailable. diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index d73fc6e4289..007c6ac869c 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -21,7 +21,14 @@ mod selection; use std::ops::Range; use std::sync::Arc; +use parking_lot::Mutex; use vortex_buffer::Alignment; +use vortex_error::VortexResult; +use vortex_io::segments::IoHandle; +use vortex_io::segments::IoIntent; +use vortex_io::segments::IoKey; +use vortex_io::segments::IoOwnerId; +use vortex_io::segments::IoRequest; use vortex_io::segments::RegisteredSource; use vortex_io::segments::SegmentIoService; use vortex_io::segments::SegmentSourceId; @@ -46,6 +53,10 @@ pub use self::prune::PrunePlanner; use crate::plan::physical::PlanRef; use crate::plan::physical::pipeline::ValueCache; use crate::scan::v2::ScanCache; +use crate::segments::SegmentId; + +/// The owner a scan's own announcements are submitted as. Nothing waits on them. +const SCAN_OWNER: IoOwnerId = IoOwnerId(u64::MAX); /// Where a segment's bytes live in the file. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -86,6 +97,8 @@ pub struct PreparedScan { pub(super) values: ValueCache, /// The rows of one chunk, the unit of work a chunk planner emits. pub(super) chunk_rows: usize, + /// Announcements of the scan's reads, held while any planner of the scan lives. + pub(super) announcements: Mutex>, } impl PreparedScan { @@ -115,6 +128,25 @@ impl PreparedScan { self.segment_io.id() } + /// Announces `segments` for the rest of the scan, so the IO service can widen the reads of + /// every chunk with them. The announcements read nothing, and are held while the scan lives: + /// each chunk runs on a driver of its own, and announcing per chunk would let one chunk's + /// read coalesce only with its own segments. + pub(super) fn announce(&self, segments: &[SegmentId]) -> VortexResult<()> { + if segments.is_empty() { + return Ok(()); + } + let batch: Vec<_> = segments + .iter() + .map(|segment| { + IoRequest::new(IoKey::segment(self.source(), **segment), IoIntent::Announce) + }) + .collect(); + let handles = self.io().session().submit(SCAN_OWNER, &batch)?; + self.announcements.lock().extend(handles); + Ok(()) + } + /// 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 { diff --git a/vortex-layout/src/scan/execution/selection.rs b/vortex-layout/src/scan/execution/selection.rs index d4bd8f0f41e..1f1f4088bf3 100644 --- a/vortex-layout/src/scan/execution/selection.rs +++ b/vortex-layout/src/scan/execution/selection.rs @@ -24,6 +24,18 @@ pub(super) fn plan_selection( selection: &Selection, kept: Option<(&Mask, u64)>, ) -> Option<(WorkScope, Box)> { + let selected = selected_rows(scope, selection, kept)?; + let scope = selected.scope.clone(); + Some((scope, scan.split_planner(selected))) +} + +/// The rows of `scope` that `selection` selects and `kept`, the rows zone pruning left from +/// `first_row` on, keeps, trimmed to the first and last of them, or `None` when there are none. +pub(super) fn selected_rows( + scope: WorkScope, + selection: &Selection, + kept: Option<(&Mask, u64)>, +) -> Option { 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) @@ -36,8 +48,7 @@ pub(super) fn plan_selection( 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(); - Some((scope, scan.split_planner(selected))) + Some(selected) } fn select_rows(mut scope: WorkScope, selection: &Selection) -> Option { diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index d6d8a625e7e..015fda424db 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -28,6 +28,7 @@ mod tests; use std::ops::Range; use std::sync::Arc; +use parking_lot::Mutex; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_error::VortexExpect; @@ -209,6 +210,7 @@ impl ScanBuilder { row_offset: self.row_offset, values, chunk_rows, + announcements: Mutex::default(), }) } } diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index 88283741693..75ff6992ee7 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -971,3 +971,23 @@ async fn shared_cache_prunes_zones_once() -> VortexResult<()> { assert_eq!(reads, vec![2, 1]); Ok(()) } + +/// Before planning its first chunk, a scan announces every segment its chunks will read, and holds +/// the announcements while it lives. +#[tokio::test] +async fn a_scan_announces_every_chunk_before_planning_one() -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, table().as_ref()).await?; + let mut ids = Vec::new(); + for layout in layout.depth_first_traversal() { + ids.extend(layout?.segment_ids()); + } + ids.sort_unstable(); + ids.dedup(); + let scan = Arc::new(builder(&session, &segments, &layout)?.prepare()?); + let rows = scan.rows().clone(); + let mut planner = compose(&scan, rows); + assert!(matches!(planner.compute()?, PlannerOutput::Spawn(..))); + assert_eq!(scan.announcements.lock().len(), ids.len()); + Ok(()) +}