From dc7bbcf2105593dc0d2c2ae2a64aaf2227eb22bd Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Fri, 2 Oct 2026 23:34:11 -0700 Subject: [PATCH] fix: fence concurrent MV refreshes --- rust/lancedb/src/materialized_view/refresh.rs | 129 +++++++++++++++++- rust/lancedb/src/remote/sql_test.rs | 85 ++++++++++++ rust/lancedb/src/remote/table.rs | 5 +- 3 files changed, 215 insertions(+), 4 deletions(-) diff --git a/rust/lancedb/src/materialized_view/refresh.rs b/rust/lancedb/src/materialized_view/refresh.rs index ca90f3214..f9634837a 100644 --- a/rust/lancedb/src/materialized_view/refresh.rs +++ b/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 { + 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 { )); 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); + } +} diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index d20ba495f..fbc43234a 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/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::(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(); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 36545a593..79511d55d 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -2156,7 +2156,7 @@ impl BaseTable for RemoteTable { 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 BaseTable for RemoteTable { 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,