feat: add priority to I/O scheduler (#2315)

This also renames store scheduler to scan scheduler. I'm thinking I
don't want to get into the thorny issue of how to prioritize I/O
requests across different scans. So, with this approach, if multiple
scans are running at the same time, then they will overschedule (and let
the OS deal with it). This can be revisited in the future.

Closes #1958

Cargo.toml

@@ -103,6 +103,7 @@ http = "0.2.9"
 itertools = "0.12"
 lazy_static = "1"
 log = "0.4"
+mockall = { version = "0.12.1" }
 mock_instant = { version = "0.3.1", features = ["sync"] }
 moka = "0.11"
 num-traits = "0.2"

python/src/file.rs

@@ -23,7 +23,7 @@ use lance_file::v2::{
     reader::{BufferDescriptor, CachedFileMetadata, FileReader},
     writer::{FileWriter, FileWriterOptions},
 };
-use lance_io::{scheduler::StoreScheduler, ReadBatchParams};
+use lance_io::{scheduler::ScanScheduler, ReadBatchParams};
 use object_store::path::Path;
 use pyo3::{
     exceptions::{PyIOError, PyRuntimeError, PyValueError},
@@ -267,7 +267,7 @@ impl LanceFileReader {
         let io_parallelism = std::env::var("IO_THREADS")
             .map(|val| val.parse::<u32>().unwrap_or(8))
             .unwrap_or(8);
-        let scheduler = StoreScheduler::new(Arc::new(object_store), io_parallelism);
+        let scheduler = ScanScheduler::new(Arc::new(object_store), io_parallelism);
         let file = scheduler.open_file(&path).await.infer_error()?;
         let inner = FileReader::try_open(file, None).await.infer_error()?;
         Ok(Self {

rust/lance-encoding/src/decoder.rs

@@ -537,8 +537,12 @@ impl DecodeBatchScheduler {
 
         let range = range.start as u32..range.end as u32;
 
-        self.root_scheduler
-            .schedule_ranges(&[range.clone()], scheduler, &sink)?;
+        self.root_scheduler.schedule_ranges(
+            &[range.clone()],
+            scheduler,
+            &sink,
+            range.start as u64,
+        )?;
 
         trace!("Finished scheduling of range {:?}", range);
         Ok(())
@@ -567,8 +571,11 @@ impl DecodeBatchScheduler {
                 format!("{}, ..., {}", indices[0], indices[indices.len() - 1])
             }
         );
+        if indices.is_empty() {
+            return Ok(());
+        }
         self.root_scheduler
-            .schedule_take(indices, scheduler, &sink)?;
+            .schedule_take(indices, scheduler, &sink, indices[0] as u64)?;
         trace!("Finished scheduling take of {} rows", indices.len());
         Ok(())
     }
@@ -740,10 +747,13 @@ pub trait PhysicalPageScheduler: Send + Sync + std::fmt::Debug {
     /// * `range` - the range of row offsets (relative to start of page) requested
     ///             these must be ordered and must not overlap
     /// * `scheduler` - a scheduler to submit the I/O request to
+    /// * `top_level_row` - the row offset of the top level field currently being
+    ///   scheduled.  This can be used to assign priority to I/O requests
     fn schedule_ranges(
         &self,
         ranges: &[Range<u32>],
         scheduler: &dyn EncodingsIo,
+        top_level_row: u64,
     ) -> BoxFuture<'static, Result<Box<dyn PhysicalPageDecoder>>>;
 }
 
@@ -780,6 +790,7 @@ pub trait LogicalPageScheduler: Send + Sync + std::fmt::Debug {
         ranges: &[Range<u32>],
         scheduler: &Arc<dyn EncodingsIo>,
         sink: &mpsc::UnboundedSender<Box<dyn LogicalPageDecoder>>,
+        top_level_row: u64,
     ) -> Result<()>;
     /// Schedules I/O for the requested rows (identified by row offsets from start of page)
     /// TODO: implement this using schedule_ranges
@@ -788,6 +799,7 @@ pub trait LogicalPageScheduler: Send + Sync + std::fmt::Debug {
         indices: &[u32],
         scheduler: &Arc<dyn EncodingsIo>,
         sink: &mpsc::UnboundedSender<Box<dyn LogicalPageDecoder>>,
+        top_level_row: u64,
     ) -> Result<()>;
     /// The number of rows covered by this page
     fn num_rows(&self) -> u32;

rust/lance-encoding/src/encodings/logical/binary.rs

@@ -47,11 +47,12 @@ impl LogicalPageScheduler for BinaryPageScheduler {
         ranges: &[std::ops::Range<u32>],
         scheduler: &Arc<dyn crate::EncodingsIo>,
         sink: &tokio::sync::mpsc::UnboundedSender<Box<dyn crate::decoder::LogicalPageDecoder>>,
+        top_level_row: u64,
     ) -> Result<()> {
         trace!("Scheduling binary for {} ranges", ranges.len());
         let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
         self.varbin_scheduler
-            .schedule_ranges(ranges, scheduler, &tx)?;
+            .schedule_ranges(ranges, scheduler, &tx, top_level_row)?;
 
         while let Some(decoder) = rx.recv().now_or_never() {
             let wrapped = BinaryPageDecoder {
@@ -69,6 +70,7 @@ impl LogicalPageScheduler for BinaryPageScheduler {
         indices: &[u32],
         scheduler: &Arc<dyn crate::EncodingsIo>,
         sink: &tokio::sync::mpsc::UnboundedSender<Box<dyn crate::decoder::LogicalPageDecoder>>,
+        top_level_row: u64,
     ) -> Result<()> {
         trace!("Scheduling binary for {} indices", indices.len());
         self.schedule_ranges(
@@ -78,6 +80,7 @@ impl LogicalPageScheduler for BinaryPageScheduler {
                 .collect::<Vec<_>>(),
             scheduler,
             sink,
+            top_level_row,
         )
     }
 

rust/lance-encoding/src/encodings/logical/fixed_size_list.rs

@@ -53,6 +53,7 @@ impl LogicalPageScheduler for FslPageScheduler {
         ranges: &[Range<u32>],
         scheduler: &Arc<dyn EncodingsIo>,
         sink: &mpsc::UnboundedSender<Box<dyn LogicalPageDecoder>>,
+        top_level_row: u64,
     ) -> Result<()> {
         let expanded_ranges = ranges
             .iter()
@@ -64,7 +65,7 @@ impl LogicalPageScheduler for FslPageScheduler {
         );
         let (tx, mut rx) = mpsc::unbounded_channel();
         self.items_scheduler
-            .schedule_ranges(&expanded_ranges, scheduler, &tx)?;
+            .schedule_ranges(&expanded_ranges, scheduler, &tx, top_level_row)?;
         let inner_page_decoder = rx.blocking_recv().unwrap();
         sink.send(Box::new(FslPageDecoder {
             inner: inner_page_decoder,
@@ -79,6 +80,7 @@ impl LogicalPageScheduler for FslPageScheduler {
         indices: &[u32],
         scheduler: &Arc<dyn EncodingsIo>,
         sink: &mpsc::UnboundedSender<Box<dyn LogicalPageDecoder>>,
+        top_level_row: u64,
     ) -> Result<()> {
         self.schedule_ranges(
             &indices
@@ -87,6 +89,7 @@ impl LogicalPageScheduler for FslPageScheduler {
                 .collect::<Vec<_>>(),
             scheduler,
             sink,
+            top_level_row,
         )
     }
 

rust/lance-encoding/src/encodings/logical/list.rs

@@ -227,6 +227,7 @@ impl LogicalPageScheduler for ListPageScheduler {
         ranges: &[std::ops::Range<u32>],
         scheduler: &Arc<dyn EncodingsIo>,
         sink: &mpsc::UnboundedSender<Box<dyn LogicalPageDecoder>>,
+        top_level_row: u64,
     ) -> Result<()> {
         // TODO: Shortcut here if the request covers the entire range (can be determined by
         // the first_invalid_offset).  If this is the case we don't need any indirect I/O.  We
@@ -258,7 +259,7 @@ impl LogicalPageScheduler for ListPageScheduler {
         // to this page.
         let (tx, mut rx) = mpsc::unbounded_channel();
         self.offsets_scheduler
-            .schedule_ranges(&offsets_ranges, scheduler, &tx)?;
+            .schedule_ranges(&offsets_ranges, scheduler, &tx, top_level_row)?;
         let mut scheduled_offsets = rx.try_recv().unwrap();
         let items_schedulers = self.items_schedulers.clone();
         let ranges = ranges.to_vec();
@@ -319,7 +320,17 @@ impl LogicalPageScheduler for ListPageScheduler {
                     // All requested items are past this page, continue
                     row_offset += next_scheduler.num_rows() as u64;
                     if !next_item_ranges.is_empty() {
-                        next_scheduler.schedule_ranges(&next_item_ranges, &scheduler, &tx)?;
+                        // Note: we are providing the same top_level_row to ALL items pages referenced by
+                        // this offsets page.  This gives them higher priority.
+                        // TODO: Ideally we would ALSO have a guarantee from the scheduler that items with
+                        // the same top_level_row are scheduled in FCFS order but I don't think it works
+                        // that way.  Still, this is probably good enough for a while
+                        next_scheduler.schedule_ranges(
+                            &next_item_ranges,
+                            &scheduler,
+                            &tx,
+                            top_level_row,
+                        )?;
                         next_item_ranges.clear();
                     }
                     next_scheduler = item_schedulers.pop_front().unwrap();
@@ -342,14 +353,24 @@ impl LogicalPageScheduler for ListPageScheduler {
                     next_item_ranges.push(page_range);
                     row_offset += next_scheduler.num_rows() as u64;
                     if !next_item_ranges.is_empty() {
-                        next_scheduler.schedule_ranges(&next_item_ranges, &scheduler, &tx)?;
+                        next_scheduler.schedule_ranges(
+                            &next_item_ranges,
+                            &scheduler,
+                            &tx,
+                            top_level_row,
+                        )?;
                         next_item_ranges.clear();
                     }
                     next_scheduler = item_schedulers.pop_front().unwrap();
                 }
             }
             if !next_item_ranges.is_empty() {
-                next_scheduler.schedule_ranges(&next_item_ranges, &scheduler, &tx)?;
+                next_scheduler.schedule_ranges(
+                    &next_item_ranges,
+                    &scheduler,
+                    &tx,
+                    top_level_row,
+                )?;
             }
             let mut item_decoders = Vec::new();
             drop(tx);
@@ -388,6 +409,7 @@ impl LogicalPageScheduler for ListPageScheduler {
         indices: &[u32],
         scheduler: &Arc<dyn EncodingsIo>,
         sink: &mpsc::UnboundedSender<Box<dyn LogicalPageDecoder>>,
+        top_level_row: u64,
     ) -> Result<()> {
         trace!("Scheduling list offsets for {} indices", indices.len());
         self.schedule_ranges(
@@ -397,6 +419,7 @@ impl LogicalPageScheduler for ListPageScheduler {
                 .collect::<Vec<_>>(),
             scheduler,
             sink,
+            top_level_row,
         )
     }
 }

rust/lance-encoding/src/encodings/logical/primitive.rs

@@ -80,12 +80,13 @@ impl LogicalPageScheduler for PrimitivePageScheduler {
         ranges: &[std::ops::Range<u32>],
         scheduler: &Arc<dyn EncodingsIo>,
         sink: &mpsc::UnboundedSender<Box<dyn LogicalPageDecoder>>,
+        top_level_row: u64,
     ) -> Result<()> {
         let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
         trace!("Scheduling ranges {:?} from physical page", ranges);
-        let physical_decoder = self
-            .physical_decoder
-            .schedule_ranges(ranges, scheduler.as_ref());
+        let physical_decoder =
+            self.physical_decoder
+                .schedule_ranges(ranges, scheduler.as_ref(), top_level_row);
 
         let logical_decoder = PrimitiveFieldDecoder {
             data_type: self.data_type.clone(),
@@ -104,6 +105,7 @@ impl LogicalPageScheduler for PrimitivePageScheduler {
         indices: &[u32],
         scheduler: &Arc<dyn EncodingsIo>,
         sink: &mpsc::UnboundedSender<Box<dyn LogicalPageDecoder>>,
+        top_level_row: u64,
     ) -> Result<()> {
         trace!(
             "Scheduling take of {} indices from physical page",
@@ -116,6 +118,7 @@ impl LogicalPageScheduler for PrimitivePageScheduler {
                 .collect::<Vec<_>>(),
             scheduler,
             sink,
+            top_level_row,
         )
     }
 }

rust/lance-encoding/src/encodings/logical/struct.rs

@@ -110,6 +110,7 @@ impl LogicalPageScheduler for SimpleStructScheduler {
         ranges: &[Range<u32>],
         scheduler: &Arc<dyn EncodingsIo>,
         sink: &mpsc::UnboundedSender<Box<dyn LogicalPageDecoder>>,
+        top_level_row: u64,
     ) -> Result<()> {
         for range in ranges.iter().cloned() {
             let mut rows_to_read = range.end - range.start;
@@ -156,6 +157,8 @@ impl LogicalPageScheduler for SimpleStructScheduler {
             // The downside of the current algorithm is that many tiny I/O batches means less opportunity for in-batch coalescing.
             // Then again, if our outer batch coalescing is super good then maybe we don't bother
 
+            let mut current_top_level_row = top_level_row;
+
             while rows_to_read > 0 {
                 let mut min_rows_added = u32::MAX;
                 for (col_idx, field_scheduler) in self.children.iter().enumerate() {
@@ -183,7 +186,12 @@ impl LogicalPageScheduler for SimpleStructScheduler {
                             page_range_start,
                             next_page
                         );
-                        next_page.schedule_ranges(&[page_range], scheduler, sink)?;
+                        next_page.schedule_ranges(
+                            &[page_range],
+                            scheduler,
+                            sink,
+                            current_top_level_row,
+                        )?;
 
                         status.rows_queued += rows_to_take;
                         status.rows_to_take -= rows_to_take;
@@ -199,6 +207,7 @@ impl LogicalPageScheduler for SimpleStructScheduler {
                     panic!("Error in scheduling logic, panic to avoid infinite loop");
                 }
                 rows_to_read -= min_rows_added;
+                current_top_level_row += min_rows_added as u64;
                 for field_status in &mut field_status {
                     field_status.rows_queued -= min_rows_added;
                 }
@@ -216,6 +225,7 @@ impl LogicalPageScheduler for SimpleStructScheduler {
         indices: &[u32],
         scheduler: &Arc<dyn EncodingsIo>,
         sink: &mpsc::UnboundedSender<Box<dyn LogicalPageDecoder>>,
+        top_level_row: u64,
     ) -> Result<()> {
         trace!("Scheduling struct decode of {} indices", indices.len());
 
@@ -236,7 +246,7 @@ impl LogicalPageScheduler for SimpleStructScheduler {
         let mut rows_to_read = indices.len() as u32;
 
         // NOTE: See schedule_range for a description of the scheduling algorithm
-
+        let mut current_top_level_row = top_level_row;
         while rows_to_read > 0 {
             let mut min_rows_added = u32::MAX;
             for (col_idx, field_scheduler) in self.children.iter().enumerate() {
@@ -269,7 +279,12 @@ impl LogicalPageScheduler for SimpleStructScheduler {
                     // We should be guaranteed to get at least one page
                     let next_page = next_page.unwrap();
 
-                    next_page.schedule_take(&indices_in_page, scheduler, sink)?;
+                    next_page.schedule_take(
+                        &indices_in_page,
+                        scheduler,
+                        sink,
+                        current_top_level_row,
+                    )?;
 
                     let rows_scheduled = indices_in_page.len() as u32;
                     status.rows_queued += rows_scheduled;
@@ -281,6 +296,7 @@ impl LogicalPageScheduler for SimpleStructScheduler {
                 panic!("Error in scheduling logic, panic to avoid infinite loop");
             }
             rows_to_read -= min_rows_added;
+            current_top_level_row += min_rows_added as u64;
             for field_status in &mut field_status {
                 field_status.rows_queued -= min_rows_added;
             }

rust/lance-encoding/src/encodings/physical/basic.rs

@@ -123,18 +123,27 @@ impl PhysicalPageScheduler for BasicPageScheduler {
         &self,
         ranges: &[std::ops::Range<u32>],
         scheduler: &dyn EncodingsIo,
+        top_level_row: u64,
     ) -> BoxFuture<'static, Result<Box<dyn PhysicalPageDecoder>>> {
         let validity_future = match &self.mode {
             SchedulerNullStatus::None(_) | SchedulerNullStatus::All => None,
             SchedulerNullStatus::Some(schedulers) => {
                 trace!("Scheduling ranges {:?} from validity", ranges);
-                Some(schedulers.validity.schedule_ranges(ranges, scheduler))
+                Some(
+                    schedulers
+                        .validity
+                        .schedule_ranges(ranges, scheduler, top_level_row),
+                )
             }
         };
 
         let values_future = if let Some(values_scheduler) = self.mode.values_scheduler() {
             trace!("Scheduling range {:?} from values", ranges);
-            Some(values_scheduler.schedule_ranges(ranges, scheduler).boxed())
+            Some(
+                values_scheduler
+                    .schedule_ranges(ranges, scheduler, top_level_row)
+                    .boxed(),
+            )
         } else {
             trace!("No values fetch needed since values all null");
             None

rust/lance-encoding/src/encodings/physical/bitmap.rs

@@ -35,6 +35,7 @@ impl PhysicalPageScheduler for DenseBitmapScheduler {
         &self,
         ranges: &[Range<u32>],
         scheduler: &dyn EncodingsIo,
+        top_level_row: u64,
     ) -> BoxFuture<'static, Result<Box<dyn PhysicalPageDecoder>>> {
         let mut min = u64::MAX;
         let mut max = 0;
@@ -62,7 +63,7 @@ impl PhysicalPageScheduler for DenseBitmapScheduler {
             min,
             max
         );
-        let bytes = scheduler.submit_request(byte_ranges);
+        let bytes = scheduler.submit_request(byte_ranges, top_level_row);
 
         async move {
             let bytes = bytes.await?;