Support u64 vertex ids in the bftree provider

diskann-bftree/src/id.rs

@@ -0,0 +1,74 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+//! Vertex-id abstraction for the bf-tree provider.
+
+use diskann::utils::{IntoUsize, VectorId};
+use diskann::{ANNError, ANNResult};
+
+/// Identifier type usable as a `BfTreeProvider` vertex id.
+///
+/// This bundles the bounds the core algorithm requires of an id ([`VectorId`])
+/// with the ability to convert *to* `usize` ([`IntoUsize`], used to key the
+/// per-vector stores) and *from* a zero-based index. The provider mints ids
+/// densely from `0..total`, so it needs a way to build an `I` from a counter.
+///
+/// Implemented for `u32` (the default, capping at ~4.29B vertices) and `u64`
+/// (for billion-scale-and-beyond, larger-than-memory datasets). On a 64-bit
+/// target `u64` covers every representable `usize`, so its conversions never
+/// fail.
+pub trait BfTreeId: VectorId + IntoUsize {
+    /// Build an id from a zero-based index, truncating on overflow.
+    ///
+    /// Only call this for indices already known to fit (e.g. ids drawn from
+    /// `0..total`, which the provider guarantees fit by construction).
+    fn from_index(index: usize) -> Self;
+
+    /// Build an id from a zero-based index, returning `None` if it does not fit.
+    fn try_from_index(index: usize) -> Option<Self>;
+}
+
+impl BfTreeId for u32 {
+    #[inline(always)]
+    fn from_index(index: usize) -> Self {
+        index as u32
+    }
+
+    #[inline(always)]
+    fn try_from_index(index: usize) -> Option<Self> {
+        u32::try_from(index).ok()
+    }
+}
+
+impl BfTreeId for u64 {
+    #[inline(always)]
+    fn from_index(index: usize) -> Self {
+        index as u64
+    }
+
+    #[inline(always)]
+    fn try_from_index(index: usize) -> Option<Self> {
+        u64::try_from(index).ok()
+    }
+}
+
+/// Validate that a provider holding `total` ids can represent every id in `0..total`.
+///
+/// `BfTreeProvider::iter` mints ids densely via the infallible (truncating)
+/// [`BfTreeId::from_index`]; callers must guarantee the range fits in `I`. This check
+/// enforces that guarantee up front (at construction and load) so the truncating
+/// conversion can never silently wrap a real id.
+pub(crate) fn validate_id_capacity<I: BfTreeId>(total: usize) -> ANNResult<()> {
+    if let Some(last) = total.checked_sub(1) {
+        if I::try_from_index(last).is_none() {
+            return Err(ANNError::log_index_error(format!(
+                "provider capacity of {total} ids exceeds the maximum representable by the \
+                 {}-byte vertex id type",
+                std::mem::size_of::<I>()
+            )));
+        }
+    }
+    Ok(())
+}

diskann-bftree/src/lib.rs

@@ -9,13 +9,16 @@
 //! [`DataProvider`](diskann::provider::DataProvider) trait, enabling indexes that can
 //! transparently spill to disk for datasets larger than available memory.
 
+pub mod id;
 pub mod neighbors;
 pub mod provider;
 pub mod quant;
 pub mod vectors;
 
 mod locks;
 
+pub use id::BfTreeId;
+
 // Accessors
 pub use provider::{
     AsVectorDtype, BfTreePaths, BfTreeProvider, BfTreeProviderParameters, CreateQuantProvider,

diskann-bftree/src/neighbors.rs

@@ -35,8 +35,11 @@ impl<I: VectorId + IntoUsize> HasId for NeighborProvider<I> {
 impl<I: VectorId + IntoUsize> NeighborProvider<I> {
     /// Create a new instance based on bf-tree Config directly.
     pub fn new_with_config(max_degree: u32, config: Config) -> ANNResult<Self> {
-        let key_size = std::mem::size_of::<u32>();
-        let value_size = (max_degree as usize + 1) * std::mem::size_of::<u32>();
+        // Records are keyed by an `I`-width id and store a `dim`-cell `I`-width value
+        // (`dim == max_degree + 1`), so size the validation by the actual id width
+        // rather than assuming `u32`.
+        let key_size = std::mem::size_of::<I>();
+        let value_size = (max_degree as usize + 1) * std::mem::size_of::<I>();
         crate::validate_record_size("neighbor_provider", &config, key_size, value_size)?;
 
         let adj_list_index = BfTree::with_config(config, None).map_err(ConfigError)?;
@@ -407,6 +410,42 @@ mod tests {
         }
     }
 
+    /// Exercise the `u64` id path beyond the `u32` range: vertex ids and neighbor
+    /// values above `u32::MAX` must round-trip, and ids that share their low 32 bits
+    /// must not collide (proving the bf-tree key uses the full 8-byte width).
+    #[tokio::test]
+    async fn test_u64_high_bit_ids() {
+        let locks = Arc::new(StripedLocks::new());
+        let neighbor_provider =
+            NeighborProvider::<u64>::new_with_config(6, Config::default()).unwrap();
+        let mut scratch = neighbor_provider.scratch(&locks);
+
+        let high: u64 = (u32::MAX as u64) + 1;
+        let big_id: u64 = high + 7;
+        let big_neighbors: Vec<u64> = vec![high, high + 1, u64::MAX, 3];
+        scratch.write_neighbors(big_id, &big_neighbors).unwrap();
+
+        let mut result = AdjacencyList::with_capacity(10);
+        neighbor_provider
+            .get_neighbors(big_id, &mut result)
+            .unwrap();
+        assert_eq!(&*big_neighbors, &*result);
+
+        // `low` and `low | (1 << 32)` share their low 32 bits; with an 8-byte key they
+        // are distinct entries. A truncating 4-byte key would alias them.
+        let low: u64 = 1;
+        let aliased: u64 = low | (1u64 << 32);
+        scratch.write_neighbors(low, &[10, 11]).unwrap();
+        scratch.write_neighbors(aliased, &[20, 21]).unwrap();
+
+        neighbor_provider.get_neighbors(low, &mut result).unwrap();
+        assert_eq!(&[10u64, 11], &*result);
+        neighbor_provider
+            .get_neighbors(aliased, &mut result)
+            .unwrap();
+        assert_eq!(&[20u64, 21], &*result);
+    }
+
     /// Test corner cases of appending to neighbor list
     #[tokio::test]
     async fn test_neighbor_accessors() {