feat(file): a coalescing IO service for planning-protocol scans

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::request::IoService;
 use vortex_layout::LayoutReader;
 use vortex_layout::scan::layout::LayoutReaderDataSource;
 use vortex_layout::scan::scan_builder::ScanBuilder;
@@ -48,6 +49,8 @@ pub struct VortexFile {
     metadata: Arc<HashMap<String, ByteBuffer>>,
     /// None id LayoutReader caching is turned off
     layout_reader_cache: Option<OnceLock<Arc<dyn LayoutReader>>>,
+    /// Serves the planning protocol's reads, when the file was opened over a reader.
+    scan_io: Option<Arc<dyn IoService>>,
 }
 
 fn layout_reader(
@@ -84,9 +87,21 @@ impl VortexFile {
             session,
             metadata: Arc::new(HashMap::new()),
             layout_reader_cache: None,
+            scan_io: None,
         }
     }
 
+    pub(crate) fn with_scan_io(mut self, scan_io: Arc<dyn IoService>) -> Self {
+        self.scan_io = Some(scan_io);
+        self
+    }
+
+    /// The service that serves the planning protocol's reads of this file, when the file was
+    /// opened over a reader.
+    pub fn scan_io(&self) -> Option<&Arc<dyn IoService>> {
+        self.scan_io.as_ref()
+    }
+
     pub(crate) fn with_metadata(mut self, metadata: Arc<HashMap<String, ByteBuffer>>) -> Self {
         self.metadata = metadata;
         self
@@ -103,6 +118,7 @@ impl VortexFile {
             session: self.session,
             metadata: self.metadata,
             layout_reader_cache: Some(OnceLock::new()),
+            scan_io: self.scan_io,
         }
     }
 
@@ -158,6 +174,8 @@ impl VortexFile {
     /// replacement source.
     pub fn with_segment_source(mut self, segment_source: Arc<dyn SegmentSource>) -> Self {
         self.segment_source = segment_source;
+        // Reads now go to the replacement, which the old read driver knows nothing of.
+        self.scan_io = None;
         if self.layout_reader_cache.is_some() {
             self.layout_reader_cache = Some(OnceLock::new());
         }

vortex-file/src/open.rs

@@ -14,6 +14,7 @@ use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_err;
 use vortex_io::VortexReadAt;
+use vortex_io::request::IoService;
 use vortex_io::session::RuntimeSessionExt;
 use vortex_layout::segments::InstrumentedSegmentCache;
 use vortex_layout::segments::NoOpSegmentCache;
@@ -35,6 +36,7 @@ use crate::MAX_POSTSCRIPT_SIZE;
 use crate::VortexFile;
 use crate::footer::Footer;
 use crate::segments::BufferSegmentSource;
+use crate::segments::FileIoService;
 use crate::segments::FileSegmentSource;
 use crate::segments::InitialReadSegmentCache;
 use crate::segments::RequestMetrics;
@@ -281,6 +283,7 @@ impl VortexOpenOptions {
     ///
     /// This is the common path for files, object stores, and custom random-access sources.
     pub async fn open_read<R: VortexReadAt + Clone>(self, reader: R) -> VortexResult<VortexFile> {
+        let has_segment_cache = self.segment_cache.is_some();
         let segment_cache = self
             .segment_cache
             .clone()
@@ -327,13 +330,26 @@ impl VortexOpenOptions {
 
         let metrics = RequestMetrics::new(metrics_registry.as_ref(), self.labels);
 
+        // Scans that drive the planning protocol read through their own service when it is
+        // selected, and through the segment source's read driver otherwise.
+        let batch_io = FileIoService::enabled().then(|| {
+            Arc::new(FileIoService::open(
+                footer.segment_specs_with_metadata(),
+                reader.clone(),
+                self.session.handle(),
+                metrics.clone(),
+            )) as Arc<dyn IoService>
+        });
+
         // Create a segment source backed by the VortexRead implementation.
-        let segment_source = Arc::new(SharedSegmentSource::new(FileSegmentSource::open(
+        let file_source = FileSegmentSource::open(
             footer.segment_specs_with_metadata(),
             reader,
             self.session.handle(),
             metrics,
-        )));
+        );
+        let scan_io = batch_io.unwrap_or_else(|| Arc::new(file_source.scan_io()));
+        let segment_source = Arc::new(SharedSegmentSource::new(file_source));
 
         // Wrap up the segment source to first resolve segments from the initial read cache.
         let segment_source: Arc<dyn SegmentSource> = Arc::new(SegmentCacheSourceAdapter::new(
@@ -348,6 +364,13 @@ impl VortexOpenOptions {
         };
         let file =
             VortexFile::new(footer, segment_source, self.session.clone()).with_metadata(metadata);
+        // The direct range service bypasses SegmentSource and its cache. Let V2 use the
+        // segment adapter when a caller explicitly configured a cache.
+        let file = if has_segment_cache {
+            file
+        } else {
+            file.with_scan_io(scan_io)
+        };
         Ok(if self.cache_layout_reader {
             file.with_caching()
         } else {