Recover worker merges and route new-table catchup through masters

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

@@ -35,7 +35,9 @@ use lance_index::mem_wal::ShardManifest;
 use crate::generic_codec::{batch_to_rows, rows_to_batch, Row};
 use crate::merge_budget::MergeMemoryBudget;
 use crate::store::{CompactionConfig, CompactionStats};
-use crate::store_base::{ListSource, PreparedMerge, StorageBase, StorageBaseOptions};
+use crate::store_base::{
+    is_not_found_error, ListSource, PreparedMerge, StorageBase, StorageBaseOptions,
+};
 use lance_context_api::schema_spec::{SchemaSpec, ID_COLUMN};
 use lance_context_api::MAX_BATCH_GET_IDS;
 
@@ -444,26 +446,38 @@ impl GenericStore {
         if let Some(filter) = filter {
             // Keep point/filtered reads on this handle's view. The server's
             // refresh-on-miss path refreshes the cached handle itself; reading
-            // a newer clone here would bypass that cache update.
-            let refs: Vec<&str> = columns.iter().map(String::as_str).collect();
-            let mut scanner = self
-                .base
-                .lsm_scanner()
-                .await?
-                .project(&refs)?
-                .filter(filter)?;
-            if limit.is_some() || offset.is_some() {
-                scanner = scanner.limit(
-                    map_i64_bound("limit", limit)?,
-                    map_i64_bound("offset", offset)?,
-                )?;
+            // a newer clone on every lookup would bypass that cache update.
+            let result = async {
+                let snapshots = self.base.wal_shard_snapshots().await?;
+                self.scan_filtered_snapshot(
+                    self.base.dataset.clone(),
+                    snapshots,
+                    filter,
+                    limit,
+                    offset,
+                    columns,
+                )
+                .await
             }
-            let mut stream = scanner.try_into_stream().await?;
-            let mut rows = Vec::new();
-            while let Some(batch) = stream.try_next().await? {
-                rows.extend(batch_to_rows(&self.spec, &batch)?);
+            .await;
+            if !self.base.is_version_pinned()
+                && matches!(&result, Err(error) if is_not_found_error(error)
+                    && !matches!(error, LanceError::InvalidInput { .. } | LanceError::SchemaMismatch { .. }))
+            {
+                // GC can remove an opened generation during the scan. An
+                // error is not an absent row: retry a fresh consistent view.
+                // Successful reads and ordinary misses retain cache semantics.
+                return self
+                    .base
+                    .read_consistent(ListSource::All, |dataset, snapshots| {
+                        self.scan_filtered_snapshot(
+                            dataset, snapshots, filter, limit, offset, columns,
+                        )
+                    })
+                    .boxed()
+                    .await;
             }
-            return Ok(rows);
+            return result;
         }
 
         // Only unfiltered lists use this compatibility recovery. Keep retries
@@ -509,6 +523,34 @@ impl GenericStore {
             .await
     }
 
+    async fn scan_filtered_snapshot(
+        &self,
+        dataset: lance::Dataset,
+        snapshots: Vec<ShardSnapshot>,
+        filter: &str,
+        limit: Option<usize>,
+        offset: Option<usize>,
+        columns: &[String],
+    ) -> LanceResult<Vec<Row>> {
+        let refs: Vec<&str> = columns.iter().map(String::as_str).collect();
+        let mut scanner =
+            StorageBase::lsm_scanner_for_dataset(&dataset, ID_COLUMN, ListSource::All, snapshots)
+                .project(&refs)?
+                .filter(filter)?;
+        if limit.is_some() || offset.is_some() {
+            scanner = scanner.limit(
+                map_i64_bound("limit", limit)?,
+                map_i64_bound("offset", offset)?,
+            )?;
+        }
+        let mut stream = scanner.try_into_stream().await?;
+        let mut rows = Vec::new();
+        while let Some(batch) = stream.try_next().await? {
+            rows.extend(batch_to_rows(&self.spec, &batch)?);
+        }
+        Ok(rows)
+    }
+
     /// Narrow compatibility path for Lance 9's broken batched PK sidecars.
     /// Physical scans supply only candidate keys; the native point planner
     /// still decides visibility/newest-wins, including replay and tombstones.
@@ -850,6 +892,68 @@ mod tests {
     use serde_json::json;
     use tempfile::TempDir;
 
+    #[tokio::test]
+    async fn filtered_point_snapshot_recovers_peer_gc_and_preserves_filter() {
+        use std::sync::atomic::{AtomicUsize, Ordering};
+        let dir = TempDir::new().unwrap();
+        let uri = dir.path().to_str().unwrap();
+        let writer = GenericStore::open(uri, spec(), sealing()).await.unwrap();
+        writer
+            .add(&[
+                row(json!({"id":"wanted", "user_id":"value"})),
+                row(json!({"id":"other", "user_id":"hidden"})),
+            ])
+            .await
+            .unwrap();
+        let reader = GenericStore::open_existing(uri, Default::default())
+            .await
+            .unwrap();
+        let writer = tokio::sync::Mutex::new(writer);
+        let attempts = AtomicUsize::new(0);
+        let columns = vec!["id".to_string(), "user_id".to_string()];
+        let rows = reader
+            .base
+            .read_consistent(ListSource::All, |dataset, snapshots| {
+                let reader = &reader;
+                let writer = &writer;
+                let attempts = &attempts;
+                let columns = &columns;
+                async move {
+                    if attempts.fetch_add(1, Ordering::SeqCst) == 0 {
+                        writer.lock().await.cleanup_wal().await.unwrap();
+                        for snapshot in &snapshots {
+                            for generation in &snapshot.sstables {
+                                let path = reader
+                                    .base
+                                    .flushed_generation_uri(snapshot.shard_id, &generation.path);
+                                tokio::time::timeout(std::time::Duration::from_secs(2), async {
+                                    while std::path::Path::new(&path).exists() {
+                                        tokio::task::yield_now().await;
+                                    }
+                                })
+                                .await
+                                .unwrap();
+                            }
+                        }
+                    }
+                    reader
+                        .scan_filtered_snapshot(
+                            dataset,
+                            snapshots,
+                            "id = 'wanted'",
+                            Some(1),
+                            None,
+                            columns,
+                        )
+                        .await
+                }
+            })
+            .await
+            .unwrap();
+        assert_eq!(attempts.load(Ordering::SeqCst), 2);
+        assert_eq!(rows, vec![row(json!({"id":"wanted", "user_id":"value"}))]);
+    }
+
     fn spec() -> SchemaSpec {
         SchemaSpec::new(vec![
             (

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

@@ -66,7 +66,7 @@ use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::SessionContext;
 use datafusion::sql::parser::{DFParser, Statement as DFStatement};
 use datafusion::sql::sqlparser::ast::Statement as SqlStatement;
-use futures::{stream, StreamExt, TryStreamExt};
+use futures::{stream, FutureExt, StreamExt, TryStreamExt};
 use lance::dataset::mem_wal::{LsmScanner, ShardManifestStore, ShardSnapshot};
 use lance::dataset::optimize::CompactionMetrics;
 use lance::dataset::scanner::MaterializationStyle;
@@ -1555,21 +1555,37 @@ impl RolloutStore {
         id: &str,
         source: ListSource,
     ) -> LanceResult<Option<RolloutRecord>> {
-        let shard_snapshots = match source {
-            ListSource::Fragments => Vec::new(),
-            ListSource::Wal | ListSource::All => self.wal_shard_snapshots().await?,
-        };
+        if source == ListSource::Fragments {
+            // Keep the cached base hit cheap. A WAL-only miss still reaches
+            // the server's existing refresh-on-miss path, updating its cache.
+            return self
+                .scan_one_by_id_snapshot(id, source, self.base.dataset.clone(), Vec::new())
+                .await;
+        }
+        self.base
+            .read_consistent(source, |dataset, snapshots| {
+                self.scan_one_by_id_snapshot(id, source, dataset, snapshots)
+            })
+            .boxed()
+            .await
+    }
+
+    async fn scan_one_by_id_snapshot(
+        &self,
+        id: &str,
+        source: ListSource,
+        dataset: Dataset,
+        shard_snapshots: Vec<ShardSnapshot>,
+    ) -> LanceResult<Option<RolloutRecord>> {
         let escaped_id = id.replace('\'', "''");
-        let columns = self
-            .non_blob_columns()
+        let columns = Self::non_blob_columns_for(&dataset)
             .into_iter()
             .filter(|column| {
                 source == ListSource::Fragments || !CLAIM_CHECK_COLUMNS.contains(&column.as_str())
             })
             .collect::<Vec<_>>();
         let refs: Vec<&str> = columns.iter().map(String::as_str).collect();
-        let scanner = self
-            .lsm_scanner_for_source(source, shard_snapshots)
+        let scanner = StorageBase::lsm_scanner_for_dataset(&dataset, "id", source, shard_snapshots)
             .project(&refs)?
             .filter(&format!("id = '{}'", escaped_id))?;
         let mut stream = scanner.try_into_stream().await?;
@@ -2525,6 +2541,65 @@ mod tests {
         assert_eq!(attempts.load(Ordering::SeqCst), 2);
     }
 
+    #[tokio::test]
+    async fn point_snapshot_recovers_peer_gc_without_changing_wal_only_semantics() {
+        use std::sync::atomic::{AtomicUsize, Ordering};
+        for source in [ListSource::All, ListSource::Wal] {
+            let dir = tempfile::tempdir().unwrap();
+            let uri = dir.path().to_str().unwrap();
+            let writer = RolloutStore::open(uri).await.unwrap();
+            writer.add(&[assistant_record("point")]).await.unwrap();
+            writer.flush().await.unwrap();
+            let reader = RolloutStore::open(uri).await.unwrap();
+            let writer = tokio::sync::Mutex::new(writer);
+            let attempts = AtomicUsize::new(0);
+            let record = reader
+                .base
+                .read_consistent(source, |dataset, snapshots| {
+                    let reader = &reader;
+                    let writer = &writer;
+                    let attempts = &attempts;
+                    async move {
+                        if attempts.fetch_add(1, Ordering::SeqCst) == 0 {
+                            writer.lock().await.cleanup_own_shard().await.unwrap();
+                            for snapshot in &snapshots {
+                                for generation in &snapshot.sstables {
+                                    let path = reader.flushed_generation_uri(
+                                        snapshot.shard_id,
+                                        &generation.path,
+                                    );
+                                    tokio::time::timeout(
+                                        std::time::Duration::from_secs(2),
+                                        async {
+                                            while std::path::Path::new(&path).exists() {
+                                                tokio::task::yield_now().await;
+                                            }
+                                        },
+                                    )
+                                    .await
+                                    .unwrap();
+                                }
+                            }
+                        }
+                        reader
+                            .scan_one_by_id_snapshot("point", source, dataset, snapshots)
+                            .await
+                    }
+                })
+                .await
+                .unwrap();
+            assert_eq!(attempts.load(Ordering::SeqCst), 2);
+            if source == ListSource::All {
+                assert_eq!(record.unwrap().id, "point");
+            } else {
+                assert!(
+                    record.is_none(),
+                    "WAL-only must not return the merged base row"
+                );
+            }
+        }
+    }
+
     #[tokio::test]
     async fn completed_list_snapshot_survives_later_merge_commits() {
         use std::sync::atomic::{AtomicUsize, Ordering};
@@ -2692,6 +2767,71 @@ mod tests {
         assert_eq!(attempts.load(Ordering::SeqCst), 2);
     }
 
+    #[tokio::test]
+    async fn list_survives_several_progressing_peer_merges() {
+        use std::sync::atomic::{AtomicUsize, Ordering};
+
+        let dir = tempfile::tempdir().unwrap();
+        let uri = dir.path().to_str().unwrap();
+        let writer = RolloutStore::open_with_options(
+            uri,
+            RolloutStoreOptions {
+                shard_id: Some("continuous-merge-writer".into()),
+                ..Default::default()
+            },
+        )
+        .await
+        .unwrap();
+        writer.add(&[assistant_record("row-0")]).await.unwrap();
+        writer.flush().await.unwrap();
+        let reader = RolloutStore::open(uri).await.unwrap();
+        let writer = tokio::sync::Mutex::new(writer);
+        let attempts = AtomicUsize::new(0);
+
+        let rows = reader
+            .base
+            .read_consistent(ListSource::All, |dataset, snapshots| {
+                let writer = &writer;
+                let reader = &reader;
+                let attempts = &attempts;
+                async move {
+                    let attempt = attempts.fetch_add(1, Ordering::SeqCst);
+                    if attempt < 4 {
+                        // Each failed snapshot loses a DIFFERENT generation to
+                        // a real successful merge. This is useful progress, not
+                        // a permanently missing file to retry indefinitely.
+                        let mut writer = writer.lock().await;
+                        assert_eq!(writer.cleanup_own_shard().await.unwrap(), 1);
+                        for snapshot in &snapshots {
+                            for generation in &snapshot.sstables {
+                                let path = reader
+                                    .flushed_generation_uri(snapshot.shard_id, &generation.path);
+                                tokio::time::timeout(std::time::Duration::from_secs(2), async {
+                                    while std::path::Path::new(&path).exists() {
+                                        tokio::time::sleep(std::time::Duration::from_millis(1))
+                                            .await;
+                                    }
+                                })
+                                .await
+                                .expect("the actual merge GC must invalidate this snapshot");
+                            }
+                        }
+                        writer
+                            .add(&[assistant_record(&format!("row-{}", attempt + 1))])
+                            .await
+                            .unwrap();
+                        writer.flush().await.unwrap();
+                    }
+                    reader.list_non_blob_snapshot(dataset, snapshots).await
+                }
+            })
+            .await
+            .expect("healthy merge progress must not exhaust a three-attempt read limit");
+        assert_eq!(attempts.load(Ordering::SeqCst), 5);
+        let ids: HashSet<_> = rows.into_iter().map(|row| row.id).collect();
+        assert_eq!(ids, (0..5).map(|i| format!("row-{i}")).collect());
+    }
+
     #[tokio::test]
     async fn consistent_read_bounds_missing_generation_retries() {
         use std::sync::atomic::{AtomicUsize, Ordering};