perf!: run scheduler initialize eagerly in async read_tasks (#6710)

BREAKING CHANGE: the rust file reader's read methods are now async. This
is to allow the caller control over when the scheduler initialization
and inline scheduling occurs so that they can parallelize this work
across fragments, if appropriate.

## Summary

Fix the v6→v7 inline-scheduling regression by running the decode
scheduler's `initialize` I/O eagerly on the awaiting task instead of
smuggling it into the returned stream's first poll.

This is offered as an alternative to #6709 (which reverted #6637
entirely): we keep the inline-scheduling optimization for small reads,
but make the work explicit and properly parallelized across fragments.

## Background

#6637 introduced an "inline scheduling" path that, for small reads,
attached the scheduler future to the front of the returned stream via
`flatten_stream` and only ran it on first poll. Combined with the
per-fragment `try_flatten` in `FilteredReadExec`
(`rust/lance/src/io/exec/filtered_read.rs:455`) and `LanceScanExec` —
both of which poll one inner stream at a time — this serialized the
scheduler's `initialize` I/O across fragments.

`StructuralPrimitiveFieldScheduler::initialize`
(`rust/lance-encoding/src/encodings/logical/primitive.rs:3422`) does a
real `io.submit_request(...).await` for chunk metadata. The cache is
per-file (the `FieldDataCacheKey` is column-scoped within a file's
metadata cache), so every fragment open misses. With 800 small fragments
× tens of ms of S3 latency, the inline path was catastrophic.

## Repro

[gist](https://gist.github.com/wkalt/e080fc9ddff6edd8eaee5ab50a069fbe) —
400k rows / 800 fragments × 500 rows, KNN brute force, no index:

|                       | before fix    | after fix       |
|-----------------------|---------------|-----------------|
| default               | ~60–66 ms     | **~48–52 ms**   |
| `LANCE_INLINE_SCHEDULING_THRESHOLD=0` (spawn) | ~46 ms | ~50–52 ms |

Default and spawn are now matched. Cross-fragment-count ablation shows
no regression at any scale (default tracks spawn ±2 ms across
rpf=500/2000/8000/50000).

## Approach

The goal was to make the scheduling work explicit, not "smuggled into
the poll of the first batch."

1. **`schedule_and_decode` is now `async`**
(`rust/lance-encoding/src/decoder.rs`). It awaits
`DecodeBatchScheduler::try_new` (which runs `initialize`) before
returning. For the inline branch, it then runs the synchronous
`schedule_ranges` / `schedule_take` work in line, leaving a fully primed
decode stream. The non-inline branch still spawns the scheduling task so
it can overlap with decoding.

2. **Cascade async through the file-reader surface.** All of
`FileReader::read_tasks`, `read_range`, `read_ranges`, `read_stream`,
and `read_stream_projected` are now `async`. Each got a "Why is this
async?" doc paragraph explaining that the decode scheduler's metadata
I/O happens on the awaiting task rather than on the consumer that polls
the stream.

3. **`GenericFileReader` trait methods return `BoxFuture<'_,
Result<ReadBatchTaskStream>>`.** V1Reader, the v2 adapter `Reader`, and
`NullReader` updated. `FragmentReader::{read_range, read_all,
read_ranges, take_range}` and `new_read_impl` are now async;
`new_read_impl` uses `try_join_all` so per-data-file `initialize` I/Os
run concurrently within a fragment.

4. **Callers updated** in `scan.rs` (v1 + v2 paths), `filtered_read.rs`,
`dataset/updater.rs` (`Updater::try_new` made async), `lance-index`
(shufflers, distributed index merger, scalar lance_format, vector
storage), benches, and `python/src/file.rs`.

The fix relies on the existing `SpawnedTask::spawn` of `read_fragment`
in `FilteredReadExec` and the `tokio::spawn` of the open task in
`LanceScanExec`: the per-fragment task now also drives `initialize`, so
all fragments' scheduling I/Os run in parallel up to
`fragment_readahead`.

## Behavior change

Errors from `initialize` (e.g. corrupted metadata, transient I/O) now
surface from the `read_*` await instead of from the first stream item.
Existing callers that match on the result of `read_*` keep working;
callers that previously assumed the construction was infallible and only
the stream could error will now see the error one step earlier.

## Test plan

- [x] `cargo check --workspace --tests --benches` — clean
- [x] Repro gist — regression resolved (numbers above)
- [x] Cross-fragment-count ablation — no regression at any scale
- [x] Python tests: 526 passed across `test_dataset`,
`test_scalar_index`, `test_blob`, `test_filter`, `test_file`,
`test_fragment`, `test_vector_index` (minus the timing-fragile test
below)
- [ ] Cloud / S3 verification — would expect a much larger improvement
than local

## One known test failure to flag


`test_create_index_progress_callback_error_before_completion_propagates`
fails after this change. It is a **pre-existing timing race exposed by
the speedup**, not a correctness break:

- The test registers `fail_on_tag="start:train_ivf"` and expects
`create_index` to raise.
- Mechanism: Rust calls `progress.stage_start("train_ivf").await`, which
only does a sync channel send — the callback's error surfaces later,
when Python's `block_on_pumping` (`python/src/executor.rs:200-247`)
calls `pump()` between 100 ms `tokio::time::sleep`s.
- After this change the default-mode operation completes inside the
first 100 ms slice often enough that pump doesn't get a chance to
surface the error mid-flight. It hits the post-completion branch at
`executor.rs:238-244`, which logs and ignores errors from the final pump
(`"Ignoring progress callback error after operation completed
successfully"`). Running with `LANCE_INLINE_SCHEDULING_THRESHOLD=0`
(spawn) makes the test pass.
- Earlier perf commits on main may already have exposed variants of this
on some platforms; commit 87ef5e2 fixed a related case.

The clean fix is in `block_on_pumping` (propagate the final pump's error
rather than ignoring it), but that's outside this refactor's scope and
changes the contract of "operation succeeded but a callback later
errored". Happy to land that as a separate PR if reviewers want.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: westonpace

python/src/file.rs

@@ -734,8 +734,6 @@ impl LanceFileReader {
         batch_size: u32,
         batch_readahead: u32,
     ) -> PyResult<PyArrowType<Box<dyn RecordBatchReader + Send>>> {
-        // read_stream is a synchronous method but it launches tasks and needs to be
-        // run in the context of a tokio runtime
         let inner = self.inner.clone();
         let stream = rt().block_on(None, async move {
             inner
@@ -745,6 +743,7 @@ impl LanceFileReader {
                     batch_readahead,
                     FilterExpression::no_filter(),
                 )
+                .await
                 .infer_error()
         })??;
         Ok(PyArrowType(Box::new(LanceReaderAdapter(stream))))

rust/lance-encoding/src/decoder.rs

@@ -263,13 +263,28 @@ const BATCH_SIZE_BYTES_WARNING: u64 = 10 * 1024 * 1024;
 const ENV_LANCE_STRUCTURAL_BATCH_DECODE_SPAWN_MODE: &str =
     "LANCE_STRUCTURAL_BATCH_DECODE_SPAWN_MODE";
 const ENV_LANCE_READ_CACHE_REPETITION_INDEX: &str = "LANCE_READ_CACHE_REPETITION_INDEX";
+const ENV_LANCE_INLINE_SCHEDULING_THRESHOLD: &str = "LANCE_INLINE_SCHEDULING_THRESHOLD";
+
+// If a request is for at most this many rows we skip the scheduler-task spawn
+// and run scheduling inline as part of the `schedule_and_decode` await.
+const DEFAULT_INLINE_SCHEDULING_THRESHOLD: u64 = 16 * 1024;
 
 fn default_cache_repetition_index() -> bool {
     static DEFAULT_CACHE_REPETITION_INDEX: OnceLock<bool> = OnceLock::new();
     *DEFAULT_CACHE_REPETITION_INDEX
         .get_or_init(|| parse_env_as_bool(ENV_LANCE_READ_CACHE_REPETITION_INDEX, true))
 }
 
+fn inline_scheduling_threshold() -> u64 {
+    static THRESHOLD: OnceLock<u64> = OnceLock::new();
+    *THRESHOLD.get_or_init(|| {
+        std::env::var(ENV_LANCE_INLINE_SCHEDULING_THRESHOLD)
+            .ok()
+            .and_then(|v| v.trim().parse::<u64>().ok())
+            .unwrap_or(DEFAULT_INLINE_SCHEDULING_THRESHOLD)
+    })
+}
+
 /// Top-level encoding message for a page.  Wraps both the
 /// legacy pb::ArrayEncoding and the newer pb::PageLayout
 ///
@@ -1956,13 +1971,32 @@ pub struct DecoderConfig {
     pub cache_repetition_index: bool,
     /// Whether to validate decoded data
     pub validate_on_decode: bool,
+    /// Override the strategy used to dispatch the scheduling work in
+    /// [`schedule_and_decode`].
+    ///
+    /// `schedule_and_decode` always awaits the scheduler's `initialize` (which
+    /// performs metadata I/O) before returning.  This flag controls what
+    /// happens with the subsequent (synchronous) work of pushing decoder
+    /// messages into the channel that feeds the decode stream.
+    ///
+    /// * `None` - default behavior: the scheduling work runs inline (as part
+    ///   of the `schedule_and_decode` await) when the request is small
+    ///   (controlled by the `LANCE_INLINE_SCHEDULING_THRESHOLD` env var) and
+    ///   is dispatched onto a spawned task otherwise.
+    /// * `Some(true)` - always run scheduling inline.  The await of
+    ///   `schedule_and_decode` does not return until every decoder message
+    ///   has been queued.
+    /// * `Some(false)` - always spawn a task for scheduling so that it can
+    ///   overlap with consumption of the decode stream.
+    pub inline_scheduling: Option<bool>,
 }
 
 impl Default for DecoderConfig {
     fn default() -> Self {
         Self {
             cache_repetition_index: default_cache_repetition_index(),
             validate_on_decode: false,
+            inline_scheduling: None,
         }
     }
 }
@@ -2089,7 +2123,7 @@ pub fn create_decode_iterator(
     }
 }
 
-fn create_scheduler_decoder(
+async fn create_scheduler_decoder(
     column_infos: Vec<Arc<ColumnInfo>>,
     requested_rows: RequestedRows,
     filter: FilterExpression,
@@ -2120,28 +2154,35 @@ fn create_scheduler_decoder(
         config.batch_size_bytes,
     )?;
 
-    let scheduler_handle = tokio::task::spawn(async move {
-        let mut decode_scheduler = match DecodeBatchScheduler::try_new(
-            target_schema.as_ref(),
-            &column_indices,
-            &column_infos,
-            &vec![],
-            num_rows,
-            config.decoder_plugins,
-            config.io.clone(),
-            config.cache,
-            &filter,
-            &config.decoder_config,
-        )
-        .await
-        {
-            Ok(scheduler) => scheduler,
-            Err(e) => {
-                let _ = tx.send(Err(e));
-                return;
-            }
-        };
+    // The scheduler's `initialize` may perform I/O to load column metadata
+    // unless that metadata is already in the cache.  This metadata loading
+    // happens as part of this call and should be parallelized if reading
+    // multiple files.
+    let mut decode_scheduler = DecodeBatchScheduler::try_new(
+        target_schema.as_ref(),
+        &column_indices,
+        &column_infos,
+        &vec![],
+        num_rows,
+        config.decoder_plugins,
+        config.io.clone(),
+        config.cache,
+        &filter,
+        &config.decoder_config,
+    )
+    .await?;
 
+    // For small requests the scheduling cost is dwarfed by the overhead of
+    // spawning a task, so we run scheduling inline (still as part of this
+    // await) before returning.  The threshold is configurable via
+    // `LANCE_INLINE_SCHEDULING_THRESHOLD`, and callers can force either
+    // strategy via `DecoderConfig::inline_scheduling`.
+    let inline_scheduling = config
+        .decoder_config
+        .inline_scheduling
+        .unwrap_or_else(|| num_rows <= inline_scheduling_threshold());
+
+    if inline_scheduling {
         match requested_rows {
             RequestedRows::Ranges(ranges) => {
                 decode_scheduler.schedule_ranges(&ranges, &filter, tx, config.io)
@@ -2150,26 +2191,50 @@ fn create_scheduler_decoder(
                 decode_scheduler.schedule_take(&indices, &filter, tx, config.io)
             }
         }
-    });
-
-    Ok(check_scheduler_on_drop(decode_stream, scheduler_handle))
+        Ok(decode_stream)
+    } else {
+        // Spawn the (still synchronous) scheduling work so that decoder
+        // messages can stream into the channel while the consumer is
+        // already pulling from the decode stream.
+        let scheduling = async move {
+            match requested_rows {
+                RequestedRows::Ranges(ranges) => {
+                    decode_scheduler.schedule_ranges(&ranges, &filter, tx, config.io)
+                }
+                RequestedRows::Indices(indices) => {
+                    decode_scheduler.schedule_take(&indices, &filter, tx, config.io)
+                }
+            }
+        };
+        let scheduler_handle = tokio::task::spawn(scheduling);
+        Ok(check_scheduler_on_drop(decode_stream, scheduler_handle))
+    }
 }
 
-/// Launches a scheduler on a dedicated (spawned) task and creates a decoder to
-/// decode the scheduled data and returns the decoder as a stream of record batches.
+/// Initializes the scheduler, schedules the requested rows, and returns a
+/// stream of decode tasks for the resulting batches.
+///
+/// This is a convenience function that creates both the scheduler and the
+/// decoder, which can be a little tricky to get right.
+///
+/// # Why is this async?
 ///
-/// This is a convenience function that creates both the scheduler and the decoder
-/// which can be a little tricky to get right.
-pub fn schedule_and_decode(
+/// Constructing the scheduler runs `initialize` which will perform I/O
+/// unless the data required is already in the file metadata cache.
+///
+/// When `DecoderConfig::inline_scheduling` resolves to `true`, the
+/// subsequent (synchronous) scheduling work also runs before this function
+/// returns, leaving a fully primed decode stream.
+pub async fn schedule_and_decode(
     column_infos: Vec<Arc<ColumnInfo>>,
     requested_rows: RequestedRows,
     filter: FilterExpression,
     column_indices: Vec<u32>,
     target_schema: Arc<Schema>,
     config: SchedulerDecoderConfig,
-) -> BoxStream<'static, ReadBatchTask> {
+) -> Result<BoxStream<'static, ReadBatchTask>> {
     if requested_rows.num_rows() == 0 {
-        return stream::empty().boxed();
+        return Ok(stream::empty().boxed());
     }
 
     // If the user requested any ranges that are empty, ignore them.  They are pointless and
@@ -2178,27 +2243,19 @@ pub fn schedule_and_decode(
 
     let io = config.io.clone();
 
-    // For convenience we really want this method to be a snchronous method where all
-    // errors happen on the stream.  There is some async initialization that must happen
-    // when creating a scheduler.  We wrap that all up in the very first task.
-    match create_scheduler_decoder(
+    let stream = create_scheduler_decoder(
         column_infos,
         requested_rows,
         filter,
         column_indices,
         target_schema,
         config,
-    ) {
-        // Keep the io alive until the stream is dropped or finishes.  Otherwise the
-        // I/O drops as soon as the scheduling is finished and the I/O loop terminates.
-        Ok(stream) => stream.finally(move || drop(io)).boxed(),
-        // If the initialization failed make it look like a failed task
-        Err(e) => stream::once(std::future::ready(ReadBatchTask {
-            num_rows: 0,
-            task: std::future::ready(Err(e)).boxed(),
-        }))
-        .boxed(),
-    }
+    )
+    .await?;
+
+    // Keep the io alive until the stream is dropped or finishes.  Otherwise the
+    // I/O drops as soon as the scheduling is finished and the I/O loop terminates.
+    Ok(stream.finally(move || drop(io)).boxed())
 }
 
 pub static WAITER_RT: LazyLock<tokio::runtime::Runtime> = LazyLock::new(|| {

rust/lance-file/benches/reader.rs

@@ -89,6 +89,7 @@ fn bench_reader(c: &mut Criterion) {
                             None,
                             FilterExpression::no_filter(),
                         )
+                        .await
                         .unwrap();
                     let stats = Arc::new(Mutex::new((0, 0)));
                     let mut stream = stream
@@ -305,6 +306,7 @@ fn read_task(
                 None,
                 FilterExpression::no_filter(),
             )
+            .await
             .unwrap();
         let stats = Arc::new(Mutex::new((0, 0)));
         let mut stream = stream.then(|batch_task| {