diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 8eabae064d..85ae57a6a9 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -1571,6 +1571,10 @@ mod disk_provider_tests { "/disk_index_search/disk_index_sift_learn_R4_L50_A1.2_truth_search_pq_compressed.bin"; const TEST_TRUTH_RESULT_10PTS_128DIM: &str = "/disk_index_search/disk_index_10pts_idx_uint32_truth_search_res.bin"; + /// Search results for an index built by the current builder, unlike + /// [`TEST_TRUTH_RESULT_10PTS_128DIM`], which matches the prebuilt [`TEST_INDEX_128DIM`]. + const TEST_TRUTH_RESULT_ASSOCIATED_DATA_10PTS_128DIM: &str = + "/disk_index_search/disk_index_sift_learn_R4_L50_A1.2_associated_data_10pts_truth_search_res.bin"; const TEST_QUERY_10PTS_128DIM: &str = "/disk_index_search/disk_index_sample_query_10pts.fbin"; const TEST_INDEX_PREFIX_100DIM: &str = "/disk_index_search/256pts_100dim_f32_truth_Index"; @@ -1727,7 +1731,7 @@ mod disk_provider_tests { index_search_engine: &search_engine, thread_num: 1, query_file_path: TEST_QUERY_10PTS_128DIM, - truth_result_file_path: TEST_TRUTH_RESULT_10PTS_128DIM, + truth_result_file_path: TEST_TRUTH_RESULT_ASSOCIATED_DATA_10PTS_128DIM, k: 10, l: 20, }, @@ -1741,7 +1745,7 @@ mod disk_provider_tests { index_search_engine: &search_engine, thread_num: 5, query_file_path: TEST_QUERY_10PTS_128DIM, - truth_result_file_path: TEST_TRUTH_RESULT_10PTS_128DIM, + truth_result_file_path: TEST_TRUTH_RESULT_ASSOCIATED_DATA_10PTS_128DIM, k: 10, l: 20, }, diff --git a/diskann-disk/src/storage/disk_index_writer.rs b/diskann-disk/src/storage/disk_index_writer.rs index 6908bb8682..f699b109d4 100644 --- a/diskann-disk/src/storage/disk_index_writer.rs +++ b/diskann-disk/src/storage/disk_index_writer.rs @@ -284,7 +284,7 @@ impl DiskIndexWriter { let num_nbrs: u32 = Self::get_neighbors_number(state)?; // sanity checks on num_nbrs - debug_assert!(num_nbrs > 0); + debug_assert!(num_nbrs > 0 || state.num_pts == 1); debug_assert!(num_nbrs <= state.max_degree); let num_nbrs_start = state.dims as usize * mem::size_of::(); @@ -559,8 +559,9 @@ impl DiskIndexWriter { #[cfg(test)] mod disk_index_storage_test { use crate::test_utils::GraphDataF32VectorU32Data; + use diskann_providers::model::graph::provider::async_::SimpleNeighborProviderAsync; use diskann_providers::storage::VirtualStorageProvider; - use diskann_utils::test_data_root; + use diskann_utils::{io::write_bin, test_data_root, views::rowmajor}; use vfs::OverlayFS; use super::*; @@ -735,6 +736,50 @@ mod disk_index_storage_test { .expect("Failed to delete file"); } + // A 1-point graph has no neighbors; creating its disk layout must succeed. + #[test] + fn create_disk_layout_one_point_graph() { + let storage_provider = VirtualStorageProvider::new_memory(); + + const DATA_FILE: &str = "/one_point.fbin"; + const INDEX_PREFIX: &str = "/one_point_index"; + const DIMS: usize = 4; + const MAX_DEGREE: u32 = 4; + + write_bin( + rowmajor::Ref::try_from_data([1.0f32; DIMS].as_slice(), 1, DIMS).unwrap(), + &mut storage_provider.create_for_write(DATA_FILE).unwrap(), + ) + .unwrap(); + + let index_writer = DiskIndexWriter::new( + DATA_FILE.to_string(), + INDEX_PREFIX.to_string(), + Option::None, + DEFAULT_DISK_SECTOR_LEN, + ) + .unwrap(); + + // A single node with no neighbors, which is also the start point. + SimpleNeighborProviderAsync::new(1, 0, MAX_DEGREE, 1.0) + .save_direct(&storage_provider, 0, &index_writer.get_mem_index_file()) + .unwrap(); + + index_writer + .create_disk_layout::(&storage_provider) + .unwrap(); + + let mut disk_layout: Vec = Vec::new(); + storage_provider + .open_reader(index_writer.disk_index_file().as_str()) + .unwrap() + .read_to_end(&mut disk_layout) + .unwrap(); + + let header = GraphHeader::try_from(&disk_layout[8..]).unwrap(); + assert_eq!(header.metadata().num_pts, 1); + } + #[test] fn disk_index_writer_rejects_small_block_size() { let small_block_size = GraphMetadata::get_size() - 1; diff --git a/diskann-providers/src/model/graph/provider/async_/simple_neighbor_provider.rs b/diskann-providers/src/model/graph/provider/async_/simple_neighbor_provider.rs index 724ec684c7..3ad46973a2 100644 --- a/diskann-providers/src/model/graph/provider/async_/simple_neighbor_provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/simple_neighbor_provider.rs @@ -305,6 +305,7 @@ impl storage::bin::GetAdjacencyList for SimpleNeighborProviderAsync { /// /// This adaptor handles these differences by: /// - Substituting the virtual start point ID with an actual dataset ID when found in adjacency lists +/// - Merging the virtual start point's out-edges into the actual start point's adjacency list /// - Excluding the virtual point from the total count (subtracting 1 from length) /// - Setting additional_points to 0 because the virtual point is omitted from the serialized graph /// @@ -315,6 +316,36 @@ struct DiskAdaptor<'a> { actual_start_point: u32, } +impl DiskAdaptor<'_> { + /// The disk graph folds the virtual start point into `actual_start_point`, so the latter + /// must inherit the virtual start point's out-edges. Otherwise, disk search would begin + /// from a different neighborhood than in-memory search, and if `actual_start_point` only + /// linked to the virtual start point, it would have no out-edges at all. + /// + /// The virtual start point's edges come first because they are the edges in-memory search + /// expands first. The merged list is capped at the longer of the two in-memory lists so + /// the serialized max degree does not grow. + fn merged_start_point_adjacency_list(&self, own: &[u32]) -> ANNResult> { + let mut merged = AdjacencyList::new(); + self.provider + .get_neighbors_sync(self.inmem_start_point as usize, &mut merged)?; + + let cap = own.len().max(merged.len()); + let is_start_point = + |n: &u32| *n == self.inmem_start_point || *n == self.actual_start_point; + + merged.retain(|n| !is_start_point(n)); + for neighbor in own.iter().copied().filter(|n| !is_start_point(n)) { + if merged.len() == cap { + break; + } + merged.push(neighbor); + } + + Ok(merged.into()) + } +} + impl storage::bin::GetAdjacencyList for DiskAdaptor<'_> { type Element = u32; type Item<'item> @@ -328,6 +359,11 @@ impl storage::bin::GetAdjacencyList for DiskAdaptor<'_> { let mut list: Vec<_> = list.into(); let node_id = u32::try_from(i)?; + + if node_id == self.actual_start_point { + return self.merged_start_point_adjacency_list(&list); + } + let mut seen_actual_start_point = false; list.retain_mut(|neighbor| { if *neighbor == self.inmem_start_point { @@ -417,6 +453,41 @@ mod tests { assert_eq!(adaptor.get_adjacency_list(1).unwrap(), vec![3]); } + #[test] + fn disk_adaptor_merges_virtual_start_point_edges_into_actual_start_point() { + let provider = SimpleNeighborProviderAsync::new(4, 1, 4, 1.0); + // The actual start point only links to the virtual start point, which links to + // everything else. + provider.set_neighbors_sync(0, &[4]).unwrap(); + provider.set_neighbors_sync(1, &[4, 2]).unwrap(); + provider.set_neighbors_sync(4, &[0, 1, 2, 3]).unwrap(); + + let adaptor = DiskAdaptor { + provider: &provider, + inmem_start_point: 4, + actual_start_point: 0, + }; + + assert_eq!(adaptor.get_adjacency_list(0).unwrap(), vec![1, 2, 3]); + assert_eq!(adaptor.get_adjacency_list(1).unwrap(), vec![0, 2]); + } + + #[test] + fn disk_adaptor_caps_merged_start_point_edges() { + let provider = SimpleNeighborProviderAsync::new(6, 1, 4, 1.0); + provider.set_neighbors_sync(0, &[6, 4, 5]).unwrap(); + provider.set_neighbors_sync(6, &[0, 1, 2, 3]).unwrap(); + + let adaptor = DiskAdaptor { + provider: &provider, + inmem_start_point: 6, + actual_start_point: 0, + }; + + // Inherited edges come first, and the list is no longer than the longest input list. + assert_eq!(adaptor.get_adjacency_list(0).unwrap(), vec![1, 2, 3, 4]); + } + #[tokio::test] async fn test_save_load() { let max_degree = 5; diff --git a/test_data/disk_index_build/truth_sift_learn_R4_L50_disk.index b/test_data/disk_index_build/truth_sift_learn_R4_L50_disk.index index fe18e3353c..1062a55d6a 100644 --- a/test_data/disk_index_build/truth_sift_learn_R4_L50_disk.index +++ b/test_data/disk_index_build/truth_sift_learn_R4_L50_disk.index @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:53fb1b3e2e468099411b9c201cf6fdcc35f31e00113a4c3a2eb0354087188851 +oid sha256:9487c3e804a119321680d54cfac96a35daa426c32dddf709a00ee2ef50798c59 size 155648 diff --git a/test_data/disk_index_search/disk_index_sift_learn_R4_L50_A1.2_associated_data_10pts_truth_search_res.bin b/test_data/disk_index_search/disk_index_sift_learn_R4_L50_A1.2_associated_data_10pts_truth_search_res.bin new file mode 100644 index 0000000000..9b09176fbf --- /dev/null +++ b/test_data/disk_index_search/disk_index_sift_learn_R4_L50_A1.2_associated_data_10pts_truth_search_res.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7cc8d33124a8e0e353f68f4051fa3d974383446fb1d13b1ee8ca2aa0a262251e +size 408