perf(scan): announce every chunk's reads when a V2 scan starts

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<Mask>,
+    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<bool> = 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<PlannerOutput> {
+        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 {

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

@@ -97,6 +97,11 @@ impl FilterPlans {
     pub(crate) fn pruning(&self) -> impl Iterator<Item = &PlanRef> {
         self.pruning.iter().flatten()
     }
+
+    /// The plan of every conjunct.
+    pub(super) fn predicates(&self) -> impl Iterator<Item = &PlanRef> {
+        self.predicates.iter()
+    }
 }
 
 /// A pruning plan must not fetch data when its requested statistics are unavailable.

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<Vec<IoHandle>>,
 }
 
 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<Self>, selected: SelectedRows) -> Box<dyn Planner> {

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<dyn Planner>)> {
+    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<SelectedRows> {
     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<SelectedRows> {

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(),
         })
     }
 }

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(())
+}