fix: demand pipeline at production scale

crates/lance-context-core/src/store_base.rs

@@ -224,6 +224,10 @@ pub struct ShardWatermark {
     pub sealed_through: u64,
     /// Flushed generations not yet merged into the base table.
     pub pending_generations: usize,
+    /// Highest generation the base table's MemWAL index records as merged
+    /// for this shard, when the caller read the base index. `None` means
+    /// "not read", not "nothing merged".
+    pub merged_through: Option<u64>,
 }
 
 pub fn is_pending_generations_exceeded(err: &LanceError) -> bool {
@@ -948,6 +952,7 @@ impl StorageBase {
                     .max()
                     .unwrap_or(0),
                 pending_generations: manifest.flushed_generations.len(),
+                merged_through: None,
             }))
     }
 
@@ -2091,13 +2096,23 @@ impl StorageBase {
         if self.is_version_pinned() {
             return Ok(Vec::new());
         }
+        // The base table's MemWAL index is the authority on what has been
+        // merged. A scan must report it, or a lost completion event leaves
+        // a fully merged shard looking pending forever, and an old table's
+        // first contact reads generation 1001 as 1001 pending.
+        let merged = Arc::new(
+            crate::rollout_append::watermarks(&self.dataset)
+                .await
+                .unwrap_or_default(),
+        );
         let object_store = self.dataset.object_store(None).await?;
         let branch_path = self.dataset.branch_location().path.clone();
         let shard_ids = self.dataset.list_mem_wal_latest_shard_ids().await?;
         let marks: Vec<Option<ShardWatermark>> = stream::iter(shard_ids)
             .map(|shard_id| {
                 let object_store = object_store.clone();
                 let branch_path = branch_path.clone();
+                let merged = merged.clone();
                 async move {
                     let manifest_store = ShardManifestStore::new(
                         object_store,
@@ -2116,6 +2131,7 @@ impl StorageBase {
                                 .max()
                                 .unwrap_or(0),
                             pending_generations: manifest.flushed_generations.len(),
+                            merged_through: merged.get(&manifest.shard_id).copied(),
                         }
                     }))
                 }

crates/lance-context-master/src/demand_publish.rs

@@ -65,7 +65,7 @@ pub(crate) fn spawn_scan_demand(
                     shard: mark.shard_id.to_string(),
                     sealed_through: mark.sealed_through,
                     sealed_bytes_through: 0,
-                    merged_through: None,
+                    merged_through: mark.merged_through,
                     flushed_at_ms: now_ms,
                     writer_epoch: mark.writer_epoch,
                     source: EventSource::Scan,

crates/lance-context-master/src/executors.rs

@@ -16,7 +16,7 @@ use std::collections::BTreeMap;
 use std::sync::Arc;
 use std::time::Duration;
 
-use etcd_client::{GetOptions, PutOptions};
+use etcd_client::PutOptions;
 use lance_context_api::TaskKind;
 use serde::{Deserialize, Serialize};
 
@@ -322,38 +322,23 @@ pub(crate) async fn load_headroom(
     client: &etcd_client::Client,
     keys: &Keys,
 ) -> Result<BTreeMap<String, Headroom>, String> {
-    let mut client = client.clone();
-    let hbs = client
-        .get(
-            keys.executors_prefix(),
-            Some(GetOptions::new().with_prefix()),
-        )
-        .await
-        .map_err(|e| e.to_string())?;
+    let (hbs, _) = crate::planner::read_prefix_paged(client, &keys.executors_prefix()).await?;
     let prefix = keys.executors_prefix();
     let heartbeats: Vec<(String, Heartbeat)> = hbs
-        .kvs()
         .iter()
         .filter_map(|kv| {
-            let id = std::str::from_utf8(kv.key())
+            let id = std::str::from_utf8(&kv.key)
                 .ok()?
                 .strip_prefix(prefix.as_str())?
                 .to_string();
-            let hb: Heartbeat = serde_json::from_slice(kv.value()).ok()?;
+            let hb: Heartbeat = serde_json::from_slice(&kv.value).ok()?;
             (hb.v == SCHEMA_VERSION).then_some((id, hb))
         })
         .collect();
-    let asg = client
-        .get(
-            keys.assignments_prefix(),
-            Some(GetOptions::new().with_prefix()),
-        )
-        .await
-        .map_err(|e| e.to_string())?;
+    let (asg, _) = crate::planner::read_prefix_paged(client, &keys.assignments_prefix()).await?;
     let assignments: Vec<Assignment> = asg
-        .kvs()
         .iter()
-        .filter_map(|kv| serde_json::from_slice::<Assignment>(kv.value()).ok())
+        .filter_map(|kv| serde_json::from_slice::<Assignment>(&kv.value).ok())
         .filter(|a| a.v == SCHEMA_VERSION)
         .collect();
     Ok(headroom(