From 690c861e9330eeaa7dedba6a973f286b7fe56d4c Mon Sep 17 00:00:00 2001 From: backurs Date: Mon, 16 Feb 2026 08:33:04 -0800 Subject: [PATCH 1/4] copy bftrees to prefix_path + bftreepaths location --- .../async_/bf_tree/neighbor_provider.rs | 4 +- .../graph/provider/async_/bf_tree/provider.rs | 59 +++++++++++++++++-- .../async_/bf_tree/quant_vector_provider.rs | 4 +- .../async_/bf_tree/vector_provider.rs | 4 +- 4 files changed, 60 insertions(+), 11 deletions(-) diff --git a/diskann-providers/src/model/graph/provider/async_/bf_tree/neighbor_provider.rs b/diskann-providers/src/model/graph/provider/async_/bf_tree/neighbor_provider.rs index fa19850d53..b853465df6 100644 --- a/diskann-providers/src/model/graph/provider/async_/bf_tree/neighbor_provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/bf_tree/neighbor_provider.rs @@ -54,8 +54,8 @@ impl NeighborProvider { /// Create a snapshot of the adjacency list index /// - pub fn snapshot(&self) { - self.adjacency_list_index.snapshot(); + pub fn snapshot(&self) -> std::path::PathBuf { + self.adjacency_list_index.snapshot() } /// Return the maximum degree (number of neighbors per vector) diff --git a/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs b/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs index 8650b969fb..a31a1d78cc 100644 --- a/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs @@ -1897,8 +1897,28 @@ where } // Save vectors and neighbors - self.full_vectors.snapshot(); - self.neighbor_provider.snapshot(); + let vectors_snapshot_path = self.full_vectors.snapshot(); + let neighbors_snapshot_path = self.neighbor_provider.snapshot(); + + // Copy snapshot files to the target prefix location if they differ + let target_vectors_path = BfTreePaths::vectors_bftree(&saved_params.prefix); + if vectors_snapshot_path != target_vectors_path { + std::fs::copy(&vectors_snapshot_path, &target_vectors_path).map_err(|e| { + ANNError::log_index_error(format!( + "Failed to copy vectors from {:?} to {:?}: {}", + vectors_snapshot_path, target_vectors_path, e + )) + })?; + } + let target_neighbors_path = BfTreePaths::neighbors_bftree(&saved_params.prefix); + if neighbors_snapshot_path != target_neighbors_path { + std::fs::copy(&neighbors_snapshot_path, &target_neighbors_path).map_err(|e| { + ANNError::log_index_error(format!( + "Failed to copy neighbors from {:?} to {:?}: {}", + neighbors_snapshot_path, target_neighbors_path, e + )) + })?; + } // Save delete bitmap { @@ -2035,9 +2055,38 @@ where } // Save vectors, neighbors, and quant vectors - self.full_vectors.snapshot(); - self.neighbor_provider.snapshot(); - self.quant_vectors.snapshot(); + let vectors_snapshot_path = self.full_vectors.snapshot(); + let neighbors_snapshot_path = self.neighbor_provider.snapshot(); + let quant_snapshot_path = self.quant_vectors.snapshot(); + + // Copy snapshot files to the target prefix location if they differ + let target_vectors_path = BfTreePaths::vectors_bftree(&saved_params.prefix); + if vectors_snapshot_path != target_vectors_path { + std::fs::copy(&vectors_snapshot_path, &target_vectors_path).map_err(|e| { + ANNError::log_index_error(format!( + "Failed to copy vectors from {:?} to {:?}: {}", + vectors_snapshot_path, target_vectors_path, e + )) + })?; + } + let target_neighbors_path = BfTreePaths::neighbors_bftree(&saved_params.prefix); + if neighbors_snapshot_path != target_neighbors_path { + std::fs::copy(&neighbors_snapshot_path, &target_neighbors_path).map_err(|e| { + ANNError::log_index_error(format!( + "Failed to copy neighbors from {:?} to {:?}: {}", + neighbors_snapshot_path, target_neighbors_path, e + )) + })?; + } + let target_quant_path = BfTreePaths::quant_bftree(&saved_params.prefix); + if quant_snapshot_path != target_quant_path { + std::fs::copy(&quant_snapshot_path, &target_quant_path).map_err(|e| { + ANNError::log_index_error(format!( + "Failed to copy quant from {:?} to {:?}: {}", + quant_snapshot_path, target_quant_path, e + )) + })?; + } // Save PQ table metadata and data using PQStorage format let filename = BfTreePaths::pq_pivots_bin(&saved_params.prefix); diff --git a/diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs b/diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs index a0582748b7..e64e7e7a94 100644 --- a/diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs @@ -76,8 +76,8 @@ impl QuantVectorProvider { /// Create a snapshot of the quant vector index /// - pub fn snapshot(&self) { - self.quant_vector_index.snapshot(); + pub fn snapshot(&self) -> std::path::PathBuf { + self.quant_vector_index.snapshot() } /// Create a new instance from an existing BfTree (for loading from snapshot) diff --git a/diskann-providers/src/model/graph/provider/async_/bf_tree/vector_provider.rs b/diskann-providers/src/model/graph/provider/async_/bf_tree/vector_provider.rs index 1a67034b5e..e7231076bd 100644 --- a/diskann-providers/src/model/graph/provider/async_/bf_tree/vector_provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/bf_tree/vector_provider.rs @@ -104,8 +104,8 @@ impl VectorProvider { /// Create a snapshot of the vector index /// #[inline(always)] - pub fn snapshot(&self) { - self.vector_index.snapshot(); + pub fn snapshot(&self) -> std::path::PathBuf { + self.vector_index.snapshot() } /// Set vector with Id, `i``, to `v` From 19cf17395665336e2ca3443b78c727476dee746a Mon Sep 17 00:00:00 2001 From: backurs Date: Tue, 17 Feb 2026 07:44:11 -0800 Subject: [PATCH 2/4] test when bf-tree snapshot path is different from saving path --- .../graph/provider/async_/bf_tree/provider.rs | 22 +++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs b/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs index a31a1d78cc..8c966c87fc 100644 --- a/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs @@ -2504,12 +2504,19 @@ mod tests { let storage = FileStorageProvider; - provider.save_with(&storage, &prefix).await.unwrap(); + // Save to a different prefix to exercise the snapshot copy logic + let save_dir = tempdir().unwrap(); + let save_prefix = save_dir + .path() + .join("saved_bf_tree_provider") + .to_string_lossy() + .to_string(); + provider.save_with(&storage, &save_prefix).await.unwrap(); // Load using trait method (includes delete bitmap) let loaded_provider = BfTreeProvider::::load_with( &storage, - &prefix.clone(), + &save_prefix, ) .await .unwrap(); @@ -2672,13 +2679,20 @@ mod tests { let storage = FileStorageProvider; - provider.save_with(&storage, &prefix).await.unwrap(); + // Save to a different prefix to exercise the snapshot copy logic + let save_dir = tempdir().unwrap(); + let save_prefix = save_dir + .path() + .join("saved_bf_tree_provider_quant") + .to_string_lossy() + .to_string(); + provider.save_with(&storage, &save_prefix).await.unwrap(); // Load using trait method (includes delete bitmap and quantization) let loaded_provider = BfTreeProvider::::load_with( &storage, - &prefix.clone(), + &save_prefix, ) .await .unwrap(); From 0851515717608752e17019baaea39026fd1a804a Mon Sep 17 00:00:00 2001 From: backurs Date: Tue, 17 Feb 2026 09:46:05 -0800 Subject: [PATCH 3/4] extract function --- .../graph/provider/async_/bf_tree/provider.rs | 83 +++++++++---------- 1 file changed, 38 insertions(+), 45 deletions(-) diff --git a/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs b/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs index 8c966c87fc..8544549b03 100644 --- a/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs @@ -1852,6 +1852,24 @@ impl BfTreePaths { } } +/// Copy a snapshot file to the target path if they differ. +/// This handles the case where the index was built with a different prefix +/// than the one being saved to. +fn copy_snapshot_if_needed( + snapshot_path: &std::path::Path, + target_path: &std::path::Path, +) -> ANNResult<()> { + if snapshot_path != target_path { + std::fs::copy(snapshot_path, target_path).map_err(|e| { + ANNError::log_index_error(format!( + "Failed to copy snapshot from {:?} to {:?}: {}", + snapshot_path, target_path, e + )) + })?; + } + Ok(()) +} + // SaveWith/LoadWith for BfTreeProvider with TableDeleteProviderAsync impl SaveWith for BfTreeProvider @@ -1901,24 +1919,14 @@ where let neighbors_snapshot_path = self.neighbor_provider.snapshot(); // Copy snapshot files to the target prefix location if they differ - let target_vectors_path = BfTreePaths::vectors_bftree(&saved_params.prefix); - if vectors_snapshot_path != target_vectors_path { - std::fs::copy(&vectors_snapshot_path, &target_vectors_path).map_err(|e| { - ANNError::log_index_error(format!( - "Failed to copy vectors from {:?} to {:?}: {}", - vectors_snapshot_path, target_vectors_path, e - )) - })?; - } - let target_neighbors_path = BfTreePaths::neighbors_bftree(&saved_params.prefix); - if neighbors_snapshot_path != target_neighbors_path { - std::fs::copy(&neighbors_snapshot_path, &target_neighbors_path).map_err(|e| { - ANNError::log_index_error(format!( - "Failed to copy neighbors from {:?} to {:?}: {}", - neighbors_snapshot_path, target_neighbors_path, e - )) - })?; - } + copy_snapshot_if_needed( + &vectors_snapshot_path, + &BfTreePaths::vectors_bftree(&saved_params.prefix), + )?; + copy_snapshot_if_needed( + &neighbors_snapshot_path, + &BfTreePaths::neighbors_bftree(&saved_params.prefix), + )?; // Save delete bitmap { @@ -2060,33 +2068,18 @@ where let quant_snapshot_path = self.quant_vectors.snapshot(); // Copy snapshot files to the target prefix location if they differ - let target_vectors_path = BfTreePaths::vectors_bftree(&saved_params.prefix); - if vectors_snapshot_path != target_vectors_path { - std::fs::copy(&vectors_snapshot_path, &target_vectors_path).map_err(|e| { - ANNError::log_index_error(format!( - "Failed to copy vectors from {:?} to {:?}: {}", - vectors_snapshot_path, target_vectors_path, e - )) - })?; - } - let target_neighbors_path = BfTreePaths::neighbors_bftree(&saved_params.prefix); - if neighbors_snapshot_path != target_neighbors_path { - std::fs::copy(&neighbors_snapshot_path, &target_neighbors_path).map_err(|e| { - ANNError::log_index_error(format!( - "Failed to copy neighbors from {:?} to {:?}: {}", - neighbors_snapshot_path, target_neighbors_path, e - )) - })?; - } - let target_quant_path = BfTreePaths::quant_bftree(&saved_params.prefix); - if quant_snapshot_path != target_quant_path { - std::fs::copy(&quant_snapshot_path, &target_quant_path).map_err(|e| { - ANNError::log_index_error(format!( - "Failed to copy quant from {:?} to {:?}: {}", - quant_snapshot_path, target_quant_path, e - )) - })?; - } + copy_snapshot_if_needed( + &vectors_snapshot_path, + &BfTreePaths::vectors_bftree(&saved_params.prefix), + )?; + copy_snapshot_if_needed( + &neighbors_snapshot_path, + &BfTreePaths::neighbors_bftree(&saved_params.prefix), + )?; + copy_snapshot_if_needed( + &quant_snapshot_path, + &BfTreePaths::quant_bftree(&saved_params.prefix), + )?; // Save PQ table metadata and data using PQStorage format let filename = BfTreePaths::pq_pivots_bin(&saved_params.prefix); From 665937c14ff36089db2fef3e44a6ab9a3034ba08 Mon Sep 17 00:00:00 2001 From: backurs Date: Tue, 17 Feb 2026 11:09:56 -0800 Subject: [PATCH 4/4] tokio --- .../graph/provider/async_/bf_tree/provider.rs | 57 +++++++++++-------- 1 file changed, 33 insertions(+), 24 deletions(-) diff --git a/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs b/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs index 8544549b03..76b0b0b01e 100644 --- a/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs @@ -1855,17 +1855,21 @@ impl BfTreePaths { /// Copy a snapshot file to the target path if they differ. /// This handles the case where the index was built with a different prefix /// than the one being saved to. -fn copy_snapshot_if_needed( - snapshot_path: &std::path::Path, - target_path: &std::path::Path, +async fn copy_snapshot_if_needed( + snapshot_path: std::path::PathBuf, + target_path: std::path::PathBuf, ) -> ANNResult<()> { if snapshot_path != target_path { - std::fs::copy(snapshot_path, target_path).map_err(|e| { - ANNError::log_index_error(format!( - "Failed to copy snapshot from {:?} to {:?}: {}", - snapshot_path, target_path, e - )) - })?; + tokio::task::spawn_blocking(move || { + std::fs::copy(&snapshot_path, &target_path).map_err(|e| { + ANNError::log_index_error(format!( + "Failed to copy snapshot from {:?} to {:?}: {}", + snapshot_path, target_path, e + )) + }) + }) + .await + .map_err(|e| ANNError::log_index_error(format!("Blocking copy task failed: {}", e)))??; } Ok(()) } @@ -1920,13 +1924,15 @@ where // Copy snapshot files to the target prefix location if they differ copy_snapshot_if_needed( - &vectors_snapshot_path, - &BfTreePaths::vectors_bftree(&saved_params.prefix), - )?; + vectors_snapshot_path, + BfTreePaths::vectors_bftree(&saved_params.prefix), + ) + .await?; copy_snapshot_if_needed( - &neighbors_snapshot_path, - &BfTreePaths::neighbors_bftree(&saved_params.prefix), - )?; + neighbors_snapshot_path, + BfTreePaths::neighbors_bftree(&saved_params.prefix), + ) + .await?; // Save delete bitmap { @@ -2069,17 +2075,20 @@ where // Copy snapshot files to the target prefix location if they differ copy_snapshot_if_needed( - &vectors_snapshot_path, - &BfTreePaths::vectors_bftree(&saved_params.prefix), - )?; + vectors_snapshot_path, + BfTreePaths::vectors_bftree(&saved_params.prefix), + ) + .await?; copy_snapshot_if_needed( - &neighbors_snapshot_path, - &BfTreePaths::neighbors_bftree(&saved_params.prefix), - )?; + neighbors_snapshot_path, + BfTreePaths::neighbors_bftree(&saved_params.prefix), + ) + .await?; copy_snapshot_if_needed( - &quant_snapshot_path, - &BfTreePaths::quant_bftree(&saved_params.prefix), - )?; + quant_snapshot_path, + BfTreePaths::quant_bftree(&saved_params.prefix), + ) + .await?; // Save PQ table metadata and data using PQStorage format let filename = BfTreePaths::pq_pivots_bin(&saved_params.prefix);