fix(disk): honor explicit PiPNN build selection

Author: SeliMeli

Cargo.lock

@@ -76,26 +76,9 @@ where
         index_writer: DiskIndexWriter,
     ) -> ANNResult<Self> {
         #[cfg(feature = "pipnn")]
-        let disk_build_param = {
-            let mut disk_build_param = disk_build_param;
-            if let Some(config) = disk_build_param.pipnn_config() {
-                config.validate()?;
-                let estimate = disk_build_param.use_vamana_if_pipnn_exceeds(
-                    index_configuration.max_points,
-                    index_configuration.dim,
-                    std::mem::size_of::<Data::VectorDataType>(),
-                    index_configuration.num_threads,
-                )?;
-                let selected = disk_build_param.build_algorithm();
-                info!(
-                    estimated_peak_bytes = estimate,
-                    memory_limit_bytes = disk_build_param.build_memory_limit().in_bytes(),
-                    algorithm = %selected,
-                    "Selected graph build algorithm"
-                );
-            }
-            disk_build_param
-        };
+        if let Some(config) = disk_build_param.pipnn_config() {
+            config.validate()?;
+        }
 
         let pq_storage = PQStorage::new(
             &(index_writer.get_index_path_prefix() + "_pq_pivots.bin"),

diskann-disk/src/build/builder/build.rs

@@ -17,10 +17,7 @@ use vfs::MemoryFS;
 
 use crate::{
     build::{
-        builder::{
-            build::DiskIndexBuilder,
-            core::{determine_build_strategy, IndexBuildStrategy},
-        },
+        builder::build::DiskIndexBuilder,
         configuration::{MemoryBudget, NumPQChunks, PiPNNParameters},
     },
     data_model::AdHoc,
@@ -147,15 +144,15 @@ fn pipnn_graph_adapter_writes_real_point_header() {
 }
 
 #[test]
-fn pipnn_disk_build_falls_back_to_complete_vamana_pipeline() {
+fn explicit_pipnn_selection_is_not_replaced_by_memory_budget() {
     let storage = VirtualStorageProvider::new_memory();
     let (points, dimensions) = (256, 8);
     write_data(&storage, points, dimensions);
     let mut builder = builder(&storage, points, dimensions, 0.000001, 1.3, pipnn());
 
     assert!(matches!(
         builder.disk_build_param.build_algorithm(),
-        crate::BuildAlgorithm::Vamana
+        crate::BuildAlgorithm::PiPNN(_)
     ));
     assert_eq!(
         builder.disk_build_param.build_quantization(),
@@ -164,21 +161,12 @@ fn pipnn_disk_build_falls_back_to_complete_vamana_pipeline() {
     assert_eq!(builder.index_configuration.config.pruned_degree().get(), 32);
     assert_eq!(builder.index_configuration.config.l_build().get(), 50);
     assert_eq!(builder.index_configuration.config.alpha(), 1.3);
-    assert!(matches!(
-        determine_build_strategy::<AdHoc<f32>>(
-            &builder.index_configuration,
-            builder.disk_build_param.build_memory_limit().in_bytes() as f64,
-            builder.disk_build_param.build_quantization(),
-        ),
-        IndexBuildStrategy::Merged
-    ));
-
     builder.build().unwrap();
     assert!(storage.exists(&get_disk_index_file("/index")));
 }
 
 #[test]
-fn pipnn_disk_build_rejects_invalid_config_before_fallback() {
+fn pipnn_disk_build_rejects_invalid_config() {
     let storage = VirtualStorageProvider::new_memory();
     let invalid = PiPNNParameters {
         c_max: 0,
@@ -193,7 +181,7 @@ fn pipnn_disk_build_rejects_invalid_config_before_fallback() {
     let writer = DiskIndexWriter::new("/data.fbin".into(), "/index".into(), None, 4096).unwrap();
 
     let error = match DiskIndexBuilder::<AdHoc<f32>, _>::new(&storage, params, config, writer) {
-        Ok(_) => panic!("invalid PiPNN config must not silently fall back to Vamana"),
+        Ok(_) => panic!("invalid PiPNN config must be rejected"),
         Err(error) => error,
     };
 

diskann-disk/src/build/builder/build/pipnn/tests.rs

@@ -8,8 +8,6 @@
 use std::num::NonZeroUsize;
 
 use diskann::ANNError;
-#[cfg(feature = "pipnn")]
-use diskann::ANNResult;
 use thiserror::Error;
 
 #[cfg(feature = "pipnn")]
@@ -111,8 +109,8 @@ impl NumPQChunks {
 /// Parameters specific for disk index construction.
 #[derive(Clone, PartialEq, Debug)]
 pub struct DiskIndexBuildParameters {
-    /// Limit on graph-construction memory. PiPNN falls back to Vamana when its
-    /// estimated one-shot peak exceeds this value.
+    /// Memory budget for disk-index pipeline stages that support bounded work.
+    /// Explicit one-shot PiPNN selection is never silently replaced.
     build_memory_limit: MemoryBudget,
 
     /// Number of PQ chunks stored in-memory for search and to be generated during build.
@@ -146,8 +144,8 @@ impl DiskIndexBuildParameters {
 
     /// Create parameters for one-shot PiPNN graph construction.
     ///
-    /// PiPNN uses the common search-PQ and disk-layout pipeline. The memory
-    /// budget selects Vamana when the estimated PiPNN peak does not fit.
+    /// PiPNN uses the common search-PQ and disk-layout pipeline. Its one-shot
+    /// graph build is not governed by the pipeline memory budget.
     #[cfg(feature = "pipnn")]
     pub fn new_pipnn(
         build_memory_limit: MemoryBudget,
@@ -204,116 +202,6 @@ impl DiskIndexBuildParameters {
             BuildAlgorithm::Vamana => None,
         }
     }
-
-    #[cfg(feature = "pipnn")]
-    pub(crate) fn use_vamana_if_pipnn_exceeds(
-        &mut self,
-        npoints: usize,
-        dimensions: usize,
-        element_size: usize,
-        num_threads: usize,
-    ) -> ANNResult<usize> {
-        let parameters = match &self.build_algorithm {
-            BuildAlgorithm::PiPNN(parameters) => parameters,
-            BuildAlgorithm::Vamana => {
-                return Err(ANNError::log_index_error(
-                    "memory selection requires PiPNN parameters",
-                ));
-            }
-        };
-        let estimate =
-            estimate_pipnn_peak_memory(parameters, npoints, dimensions, element_size, num_threads);
-        if estimate.is_none_or(|bytes| bytes > self.build_memory_limit.in_bytes()) {
-            self.build_algorithm = BuildAlgorithm::Vamana;
-        }
-        Ok(estimate.unwrap_or(usize::MAX))
-    }
-}
-
-#[cfg(feature = "pipnn")]
-fn estimate_pipnn_peak_memory(
-    config: &PiPNNParameters,
-    npoints: usize,
-    dimensions: usize,
-    element_size: usize,
-    num_threads: usize,
-) -> Option<usize> {
-    const PROPORTIONAL_HEADROOM_PERCENT: u128 = 108;
-    const PROCESS_HEADROOM: u128 = 16 * 1024 * 1024;
-    const PER_WORKER_HEADROOM: u128 = 28 * 1024 * 1024;
-
-    let workers = if num_threads == 0 {
-        std::thread::available_parallelism().map_or(1, NonZeroUsize::get)
-    } else {
-        num_threads
-    } as u128;
-    let copies = config
-        .fanout
-        .iter()
-        .try_fold(config.replicas as u128, |copies, &fanout| {
-            copies.checked_mul(fanout as u128)
-        })?;
-    let dataset = (dimensions as u128).checked_mul(element_size as u128)?;
-    let leaf_ids = copies.checked_mul(size_of::<u32>() as u128)?;
-    let offered = copies.checked_mul(config.k as u128)?.checked_mul(2)?;
-    let candidate_capacity = offered.checked_next_power_of_two()?;
-    let candidate_storage = candidate_capacity.checked_mul(size_of::<u32>() as u128)?;
-    let candidate_metadata =
-        size_of::<std::sync::Mutex<diskann::graph::AdjacencyList<u32>>>() as u128;
-
-    let partition_per_point = dataset
-        .checked_add(16)?
-        .checked_add(leaf_ids.checked_mul(2)?)?;
-    let leaf_per_point = dataset
-        .checked_add(candidate_metadata)?
-        .checked_add(candidate_storage)?
-        .checked_add(leaf_ids)?;
-    let finalization_per_point = dataset
-        .checked_add(candidate_metadata)?
-        .checked_add(size_of::<Vec<u32>>() as u128)?
-        .checked_add(candidate_storage)?;
-
-    let leaf_size = config.c_max.min(npoints).max(1) as u128;
-    let leaf_scratch = leaf_size
-        .checked_mul(dimensions as u128)?
-        .checked_mul(size_of::<f32>() as u128)?
-        .checked_add(leaf_size.checked_mul(leaf_size)?.checked_mul(5)?)?
-        .checked_add(leaf_size.checked_mul(config.k as u128)?.checked_mul(24)?)?
-        .checked_add(leaf_size.checked_mul(20)?)?
-        .checked_add(68)?;
-    let partition_rows = (npoints as u128).min(1_024);
-    let partition_leaders = (npoints as u128).min(1_000);
-    let partition_scratch = partition_leaders
-        .checked_mul(dimensions as u128)?
-        .checked_mul(size_of::<f32>() as u128)?
-        .checked_add(
-            partition_rows
-                .checked_mul(dimensions as u128)?
-                .checked_mul(size_of::<f32>() as u128)?,
-        )?
-        .checked_add(
-            partition_rows
-                .checked_mul(partition_leaders)?
-                .checked_mul(size_of::<f32>() as u128)?,
-        )?;
-
-    let points = npoints as u128;
-    let structural = points
-        .checked_mul(partition_per_point)?
-        .checked_add(workers.checked_mul(partition_scratch)?)?
-        .max(
-            points
-                .checked_mul(leaf_per_point)?
-                .checked_add(workers.checked_mul(leaf_scratch)?)?,
-        )
-        .max(points.checked_mul(finalization_per_point)?);
-    let proportional = structural
-        .checked_mul(PROPORTIONAL_HEADROOM_PERCENT)?
-        .div_ceil(100);
-    let allocator_floor = structural
-        .checked_add(PROCESS_HEADROOM)?
-        .checked_add(workers.checked_mul(PER_WORKER_HEADROOM)?)?;
-    usize::try_from(proportional.max(allocator_floor)).ok()
 }
 
 #[cfg(test)]
@@ -407,47 +295,12 @@ mod dataset_test {
         let params = DiskIndexBuildParameters::new_pipnn(budget, pq, parameters);
 
         assert_eq!(params.pipnn_config(), Some(config));
+        assert_eq!(params.build_memory_limit(), budget);
         assert_eq!(params.search_pq_chunks(), pq);
         assert_eq!(
             params.data_compression_chunk_vector_count(),
             DEFAULT_DATA_COMPRESSION_CHUNK_VECTOR_COUNT
         );
         assert!(matches!(params.build_algorithm(), BuildAlgorithm::PiPNN(_)));
     }
-
-    #[cfg(feature = "pipnn")]
-    #[test]
-    fn pipnn_memory_estimate_covers_measured_bigann10m_peak() {
-        const MEASURED_PEAK: usize = 8_656_004 * 1024;
-        let parameters = PiPNNParameters {
-            c_max: 512,
-            c_min: 64,
-            p_samp: 0.01,
-            fanout: vec![10, 3],
-            k: 2,
-            replicas: 1,
-        };
-
-        let estimate = estimate_pipnn_peak_memory(&parameters, 10_000_000, 128, 2, 16).unwrap();
-
-        assert!(estimate >= MEASURED_PEAK);
-        assert!(estimate <= MEASURED_PEAK * 120 / 100);
-    }
-
-    #[cfg(feature = "pipnn")]
-    #[test]
-    fn pipnn_falls_back_to_vamana_above_budget() {
-        let budget = MemoryBudget::try_from_gb(1.0).unwrap();
-        let pq = NumPQChunks::new_with(1, 128).unwrap();
-        let mut params =
-            DiskIndexBuildParameters::new_pipnn(budget, pq, PiPNNParameters::default());
-
-        let estimate = params
-            .use_vamana_if_pipnn_exceeds(10_000_000, 128, 2, 16)
-            .unwrap();
-
-        assert!(estimate > budget.in_bytes());
-        assert!(matches!(params.build_algorithm(), BuildAlgorithm::Vamana));
-        assert_eq!(params.build_quantization(), &QuantizationType::FP);
-    }
 }