fix: fence concurrent MV refreshes

rust/lancedb/src/materialized_view/refresh.rs

@@ -21,6 +21,9 @@ use futures::StreamExt;
 use lance::Dataset;
 use lance::dataset::mem_wal::DatasetMemWalExt;
 use lance::dataset::transaction::{Operation, Transaction};
+use lance::dataset::write::merge_insert::inserted_rows::{
+    KeyExistenceFilter, KeyExistenceFilterBuilder, KeyValue,
+};
 use lance::dataset::{CommitBuilder, InsertBuilder, WriteDestination, WriteMode, WriteParams};
 use lance_datafusion::planner::Planner;
 use serde::{Deserialize, Serialize};
@@ -233,6 +236,7 @@ async fn replace_fragments(dataset: Dataset, stream: SendableRecordBatchStream)
             message: "expected an append while staging replacement rows".into(),
         });
     };
+    let inserted_rows_filter = refresh_conflict_filter(&dataset)?;
     let committed = CommitBuilder::new(WriteDestination::Dataset(dataset))
         .execute(Transaction::new(
             read_version,
@@ -244,7 +248,7 @@ async fn replace_fragments(dataset: Dataset, stream: SendableRecordBatchStream)
                 compacted_sstables: Vec::new(),
                 fields_for_preserving_frag_bitmap: Vec::new(),
                 update_mode: None,
-                inserted_rows_filter: None,
+                inserted_rows_filter: Some(inserted_rows_filter),
                 updated_fragment_offsets: None,
             },
             None,
@@ -262,6 +266,28 @@ async fn replace_fragments(dataset: Dataset, stream: SendableRecordBatchStream)
     Ok(committed)
 }
 
+const REFRESH_TOKEN_ID: u64 = u64::MAX;
+
+fn refresh_conflict_filter(dataset: &Dataset) -> Result<KeyExistenceFilter> {
+    let field_id = dataset
+        .schema()
+        .fields
+        .first()
+        .map(|field| field.id)
+        .ok_or_else(|| Error::Runtime {
+            message: "a materialized view must have at least one field".into(),
+        })?;
+    let mut filter = KeyExistenceFilterBuilder::new(vec![field_id]);
+    filter
+        .insert(KeyValue::UInt64(REFRESH_TOKEN_ID))
+        .map_err(|error| Error::Runtime {
+            message: format!(
+                "failed to build the materialized-view refresh conflict filter: {error}"
+            ),
+        })?;
+    Ok(filter.build())
+}
+
 /// Reject MemWAL/LSM state because an ordinary dataset scan cannot see its
 /// un-compacted tiers.
 pub(crate) async fn ensure_no_mem_wal(dataset: &Dataset, role: &str, name: &str) -> Result<()> {
@@ -540,3 +566,104 @@ fn unnest_batch(batch: &RecordBatch, list_column: &str) -> Result<RecordBatch> {
     ));
     Ok(RecordBatch::try_new(schema, columns)?)
 }
+
+#[cfg(test)]
+mod tests {
+    use arrow_array::record_batch;
+    use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+
+    use super::*;
+    use crate::connect;
+
+    #[tokio::test]
+    async fn raced_first_full_refresh_lands_one_result() {
+        let directory = tempfile::tempdir().unwrap();
+        let connection = connect(directory.path().to_str().unwrap())
+            .execute()
+            .await
+            .unwrap();
+        let source_batch = record_batch!(("x", Int32, [1, 2, 3])).unwrap();
+        connection
+            .create_table("src", source_batch.clone())
+            .execute()
+            .await
+            .unwrap();
+        let view = connection
+            .create_materialized_view("copy", "src")
+            .with_no_data(true)
+            .execute()
+            .await
+            .unwrap();
+        let schema = view.table().schema().await.unwrap();
+        let result_batch =
+            RecordBatch::try_new(schema.clone(), source_batch.columns().to_vec()).unwrap();
+        let planned = view
+            .table()
+            .as_native()
+            .unwrap()
+            .dataset
+            .get()
+            .await
+            .unwrap()
+            .as_ref()
+            .clone();
+        assert!(planned.get_fragments().is_empty());
+
+        let (entered_tx, entered_rx) = tokio::sync::oneshot::channel();
+        let (release_tx, release_rx) = tokio::sync::oneshot::channel();
+        let loser_batch = result_batch.clone();
+        let blocked = futures::stream::once(async move {
+            entered_tx.send(()).unwrap();
+            release_rx.await.unwrap();
+            Ok::<_, DataFusionError>(loser_batch)
+        });
+        let stream = Box::pin(RecordBatchStreamAdapter::new(schema.clone(), blocked));
+        let loser_table = view.table().clone();
+        let loser =
+            tokio::spawn(
+                async move { replace_materialized_view_fragments(&loser_table, stream).await },
+            );
+        entered_rx.await.unwrap();
+
+        let winner_stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new(
+            schema,
+            futures::stream::iter([Ok(result_batch)]),
+        ));
+        let planned = Arc::new(planned);
+        let staged = InsertBuilder::new(WriteDestination::Dataset(planned.clone()))
+            .with_params(&WriteParams {
+                mode: WriteMode::Append,
+                ..Default::default()
+            })
+            .execute_uncommitted_stream(winner_stream)
+            .await
+            .unwrap();
+        let Operation::Append { fragments } = staged.operation else {
+            panic!("expected append staging");
+        };
+        let filter = refresh_conflict_filter(&planned).unwrap();
+        CommitBuilder::new(WriteDestination::Dataset(planned.clone()))
+            .execute(Transaction::new(
+                planned.version().version,
+                Operation::Update {
+                    removed_fragment_ids: Vec::new(),
+                    updated_fragments: Vec::new(),
+                    new_fragments: fragments,
+                    fields_modified: Vec::new(),
+                    compacted_sstables: Vec::new(),
+                    fields_for_preserving_frag_bitmap: Vec::new(),
+                    update_mode: None,
+                    inserted_rows_filter: Some(filter),
+                    updated_fragment_offsets: None,
+                },
+                None,
+            ))
+            .await
+            .unwrap();
+        release_tx.send(()).unwrap();
+
+        loser.await.unwrap().unwrap_err();
+        let current = connection.open_table("copy").execute().await.unwrap();
+        assert_eq!(current.count_rows(None).await.unwrap(), 3);
+    }
+}

rust/lancedb/src/remote/sql_test.rs

@@ -446,6 +446,91 @@ async fn catalog_connections_use_explicit_sql_endpoint_and_database_scope() {
     server.await.unwrap().unwrap();
 }
 
+#[tokio::test]
+async fn remote_refresh_twice_keeps_latest_selector() {
+    use crate::remote::client::test_utils::client_with_handler;
+    use crate::remote::table::RemoteTable;
+    use crate::table::BaseTable;
+
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let address = listener.local_addr().unwrap();
+    let incoming = futures::stream::try_unfold(listener, |listener| async {
+        let (socket, _) = listener.accept().await?;
+        Ok::<_, std::io::Error>(Some((socket, listener)))
+    });
+    let service = TestSqlService::default();
+    let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
+    let server = tokio::spawn(
+        tonic::transport::Server::builder()
+            .add_service(FlightServiceServer::new(service))
+            .serve_with_incoming_shutdown(incoming, async {
+                let _ = shutdown_rx.await;
+            }),
+    );
+    let sql_client = SqlClient::new(
+        "analytics".into(),
+        None,
+        "test-key".into(),
+        None,
+        Some(format!("grpc://{address}")),
+        ClientConfig::default(),
+    );
+    let selectors = Arc::new(std::sync::Mutex::new(Vec::new()));
+    let captured = selectors.clone();
+    let client = client_with_handler(move |request| {
+        let body = request
+            .body()
+            .and_then(reqwest::Body::as_bytes)
+            .and_then(|bytes| serde_json::from_slice::<serde_json::Value>(bytes).ok())
+            .unwrap();
+        captured.lock().unwrap().push(body["version"].clone());
+        let text = match request.url().path() {
+            "/v1/table/ns$view/describe/" => r#"{"version": 42, "schema": {"fields": []}}"#,
+            "/v1/table/ns$view/count_rows/" => "2",
+            other => panic!("unexpected REST request: {other}"),
+        };
+        http::Response::builder()
+            .status(200)
+            .body(text.to_string())
+            .unwrap()
+    });
+    let table = RemoteTable::new_with_sql_client(
+        client,
+        "view".into(),
+        vec!["ns".into()],
+        "ns$view".into(),
+        Default::default(),
+        Some(sql_client),
+    );
+
+    let first = table
+        .refresh_materialized_view_async(None)
+        .await
+        .unwrap()
+        .wait()
+        .await
+        .unwrap();
+    let second = table
+        .refresh_materialized_view_async(None)
+        .await
+        .unwrap()
+        .wait()
+        .await
+        .unwrap();
+
+    assert_eq!(first.version, 42);
+    assert_eq!(second.version, 42);
+    assert!(
+        selectors
+            .lock()
+            .unwrap()
+            .iter()
+            .all(serde_json::Value::is_null)
+    );
+    shutdown_tx.send(()).unwrap();
+    server.await.unwrap().unwrap();
+}
+
 #[tokio::test]
 async fn submits_polls_fetches_cancels_and_reuses_client() {
     let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();

rust/lancedb/src/remote/table.rs

@@ -2156,7 +2156,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
         let identifier = self.identifier.clone();
         let server_version = self.server_version.clone();
         let freshness = self.freshness.clone();
-        let version_cache = self.version.clone();
+        let freshness_request = self.snapshot_freshness_headers();
         Ok(Job::spawned(tokio::spawn(async move {
             let query = sql_client.submit(&statement, &namespace).await?;
             let mut reader = query.reader().await?;
@@ -2176,8 +2176,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
                 u64::try_from(table.count_rows(None).await?).map_err(|_| Error::Runtime {
                     message: "materialized-view row count exceeds u64".to_string(),
                 })?;
-            *version_cache.write().await = Some(version);
-            track_read_version(&freshness, version);
+            freshness_request.observe_version(&freshness, version);
             Ok(RefreshMaterializedViewResult {
                 mode: crate::materialized_view::RefreshMode::Rebuild,
                 rows_written,