feat(scan): run the V2 executor on the segment IO protocol

vortex-datafusion/src/persistent/opener.rs

@@ -434,12 +434,13 @@ impl FileOpener for VortexOpener {
                 Ok(RecordBatch::from(arrow.as_struct().clone()))
             };
             let stream = if scan_v2 {
-                let mut scan_file = vxf.scan_file();
-                scan_file.segments = shared_segment_source(
+                let shared = shared_segment_source(
                     &segment_sources,
                     &file.object_meta.location,
-                    scan_file.segments,
+                    vxf.segment_source(),
                 );
+                let vxf = vxf.clone().with_segment_source(shared);
+                let mut scan_file = vxf.scan_file();
                 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)

vortex-datafusion/src/persistent/scan_v2.rs

@@ -25,8 +25,8 @@ 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::runtime::Handle;
+use vortex::io::segments::SegmentIoSession;
 use vortex::layout::scan::v2::PreparedScan;
 use vortex::scan::planning::driver::Progress;
 use vortex::scan::planning::driver::Run;
@@ -128,7 +128,7 @@ struct RunState {
     handle: Handle,
     /// Taken while a CPU step runs on its task and put back with its result.
     run: Option<Run>,
-    io: Arc<dyn IoSource>,
+    io: Arc<dyn SegmentIoSession>,
     map: MapBatch,
     /// Whether a spawned chunk runs on this run, as another root, instead of being yielded.
     inline: bool,
@@ -143,7 +143,7 @@ fn run_stream(
     handle: Handle,
     root: Box<dyn Planner>,
     scope: WorkScope,
-    io: Arc<dyn IoSource>,
+    io: Arc<dyn SegmentIoSession>,
     map: MapBatch,
     inline: bool,
 ) -> BoxStream<'static, VortexResult<Yielded>> {
@@ -211,7 +211,7 @@ fn chunk_stream(
     handle: &Handle,
     root: Box<dyn Planner>,
     scope: WorkScope,
-    io: Arc<dyn IoSource>,
+    io: Arc<dyn SegmentIoSession>,
     map: MapBatch,
 ) -> BoxStream<'static, VortexResult<(u64, RecordBatch)>> {
     let (mut send, recv) = mpsc::channel(1);

vortex-duckdb/src/scan_v2.rs

@@ -27,11 +27,11 @@ use parking_lot::Mutex;
 use vortex::array::ArrayRef;
 use vortex::error::VortexExpect;
 use vortex::error::VortexResult;
-use vortex::io::request::IoConsumer;
-use vortex::io::request::IoRequestId;
-use vortex::io::request::IoResult;
-use vortex::io::request::IoService;
+use vortex::error::vortex_panic;
 use vortex::io::runtime::BlockingRuntime as _;
+use vortex::io::segments::IoKey;
+use vortex::io::segments::SegmentIoService;
+use vortex::io::segments::SourceResult;
 use vortex::layout::scan::v2::ChunkPlanner;
 use vortex::layout::scan::v2::PreparedScan;
 use vortex::layout::scan::v2::PrunePlanner;
@@ -59,7 +59,7 @@ const PART_ROW_BLOCK: u64 = 65_536;
 /// One file's V2 scan, shared by the threads scanning it.
 pub(crate) struct FileScan {
     scan: Arc<PreparedScan>,
-    io: Arc<dyn IoService>,
+    io: Arc<dyn SegmentIoService>,
     /// The run of the root planner, locked while a thread takes the chunk it spawns next.
     root: Mutex<Run>,
     /// Set once the root has spawned its last chunk.
@@ -223,11 +223,11 @@ impl PartPlanner {
     }
 }
 
-impl IoConsumer for PartPlanner {
-    fn set_io_result(&mut self, _request: IoRequestId, _result: IoResult) {}
-}
-
 impl Planner for PartPlanner {
+    fn deliver(&mut self, key: IoKey, _result: SourceResult) {
+        vortex_panic!("PartPlanner: delivery of {key:?}, though it never fetches");
+    }
+
     fn state(&self) -> State {
         if self.parts.as_slice().is_empty() {
             State::Done

vortex-file/src/file.rs

@@ -17,6 +17,7 @@ use vortex_array::dtype::FieldMask;
 use vortex_array::expr::Expression;
 use vortex_buffer::ByteBuffer;
 use vortex_error::VortexResult;
+use vortex_io::segments::RegisteredSource;
 use vortex_layout::LayoutReader;
 use vortex_layout::scan::layout::LayoutReaderDataSource;
 use vortex_layout::scan::scan_builder::ScanBuilder;
@@ -32,6 +33,7 @@ use vortex_utils::aliases::hash_map::HashMap;
 use crate::FileStatistics;
 use crate::footer::Footer;
 use crate::pruning::can_prune_file_stats;
+use crate::segment_io::FileSegmentIo;
 use crate::source::FileDataSource;
 use crate::v2::FileStatsLayoutReader;
 
@@ -56,14 +58,24 @@ pub struct VortexFile {
     scan_cache: Option<FileScanCache>,
 }
 
-/// What repeated V2 scans of a file share: its segment locations, read from the footer once, and
-/// the plan state derived from them.
+/// What repeated V2 scans of a file share: its segment locations, read from the footer once, its
+/// registration with the session's [`FileSegmentIo`], and the plan state derived from them.
 #[derive(Clone)]
 struct FileScanCache {
     locations: Arc<[SegmentLocation]>,
+    segment_io: RegisteredSource,
     plans: ScanCache,
 }
 
+/// Registers `segment_source` with the session's [`FileSegmentIo`]. The file's size is unknown
+/// when it was opened with a cached footer, and the V2 scan only reads segments.
+fn register_segments(
+    session: &VortexSession,
+    segment_source: &Arc<dyn SegmentSource>,
+) -> RegisteredSource {
+    FileSegmentIo::of(session).register(Arc::clone(segment_source), None)
+}
+
 fn segment_locations(footer: &Footer) -> Arc<[SegmentLocation]> {
     footer
         .segment_specs_with_metadata()
@@ -127,6 +139,7 @@ impl VortexFile {
     pub fn with_caching(self) -> Self {
         let scan_cache = FileScanCache {
             locations: segment_locations(&self.footer),
+            segment_io: register_segments(&self.session, &self.segment_source),
             plans: ScanCache::default(),
         };
         Self {
@@ -195,6 +208,7 @@ impl VortexFile {
             self.layout_reader_cache = Some(OnceLock::new());
         }
         if let Some(cache) = &mut self.scan_cache {
+            cache.segment_io = register_segments(&self.session, &self.segment_source);
             cache.plans = ScanCache::default();
         }
         self
@@ -259,19 +273,29 @@ impl VortexFile {
     }
 
     /// The parts of this file the planner-driven scan executor reads through: its root layout,
-    /// where each of its segments lives, and its segment source.
+    /// where each of its segments lives, and its registration with the session's
+    /// [`FileSegmentIo`], which reads through the file's segment source.
     ///
     /// Pass this to [`vortex_layout::scan::v2::ScanBuilder::new`] to build a V2 scan explicitly.
-    /// With [caching](Self::with_caching), every returned [`ScanFile`] shares one [`ScanCache`].
+    /// With [caching](Self::with_caching), every returned [`ScanFile`] shares one registration
+    /// and one [`ScanCache`].
     pub fn scan_file(&self) -> ScanFile {
-        let (locations, cache) = match &self.scan_cache {
-            Some(cache) => (Arc::clone(&cache.locations), Some(cache.plans.clone())),
-            None => (segment_locations(&self.footer), None),
+        let (locations, segment_io, cache) = match &self.scan_cache {
+            Some(cache) => (
+                Arc::clone(&cache.locations),
+                cache.segment_io.clone(),
+                Some(cache.plans.clone()),
+            ),
+            None => (
+                segment_locations(&self.footer),
+                register_segments(&self.session, &self.segment_source),
+                None,
+            ),
         };
         ScanFile {
             layout: Arc::clone(self.footer.layout()),
             locations,
-            segments: Arc::clone(&self.segment_source),
+            segment_io,
             cache,
         }
     }

vortex-file/src/lib.rs

@@ -112,6 +112,7 @@ mod open;
 mod pruning;
 mod read;
 /// Segment sources, caches, and sinks used by file readers and writers.
+pub mod segment_io;
 pub mod segments;
 mod source;
 mod strategy;