Add PiPNN: high-performance alternative index builder (7.6-10x faster graph build)

Cargo.toml

@@ -24,6 +24,8 @@ members = [
     "diskann-benchmark",
     "diskann-tools",
     "vectorset",
+    # PiPNN
+    "diskann-pipnn",
 ]
 
 default-members = [
@@ -60,6 +62,8 @@ diskann = { path = "diskann", version = "0.49.1" }
 diskann-providers = { path = "diskann-providers", default-features = false, version = "0.49.1" }
 diskann-disk = { path = "diskann-disk", version = "0.49.1" }
 diskann-label-filter = { path = "diskann-label-filter", version = "0.49.1" }
+# PiPNN
+diskann-pipnn = { path = "diskann-pipnn", version = "0.49.1" }
 # Infra
 diskann-benchmark-runner = { path = "diskann-benchmark-runner", version = "0.49.1" }
 diskann-benchmark-core = { path = "diskann-benchmark-core", version = "0.49.1" }

diskann-benchmark/Cargo.toml

@@ -30,7 +30,8 @@ diskann-vector.workspace = true
 diskann-wide.workspace = true
 diskann-label-filter.workspace = true
 diskann-tools = { workspace = true }
-diskann-disk = { workspace = true, optional = true }
+diskann-disk = { workspace = true, optional = true, features = ["pipnn"] }
+diskann-pipnn = { workspace = true }
 cfg-if.workspace = true
 diskann-benchmark-runner = { workspace = true }
 opentelemetry = { workspace = true, optional = true }

diskann-benchmark/src/backend/disk_index/build.rs

@@ -91,10 +91,11 @@ where
 
     let metadata = load_metadata_from_file(storage_provider, &data_path)?;
 
-    let build_parameters = DiskIndexBuildParameters::new(
+    let build_parameters = DiskIndexBuildParameters::new_with_algorithm(
         MemoryBudget::try_from_gb(params.build_ram_limit_gb)?,
         params.quantization_type,
         NumPQChunks::new_with(params.num_pq_chunks.get(), metadata.ndims())?,
+        params.build_algorithm.clone(),
     );
 
     let index_configuration = IndexConfiguration::new(

diskann-benchmark/src/inputs/disk.rs

@@ -10,6 +10,8 @@ use diskann_benchmark_runner::{
     files::InputFile, utils::datatype::DataType, CheckDeserialization, Checker,
 };
 #[cfg(feature = "disk-index")]
+use diskann_disk::BuildAlgorithm;
+#[cfg(feature = "disk-index")]
 use diskann_disk::QuantizationType;
 use diskann_providers::storage::{get_compressed_pq_file, get_disk_index_file, get_pq_pivot_file};
 use serde::{Deserialize, Serialize};
@@ -68,6 +70,10 @@ pub(crate) struct DiskIndexBuild {
     pub(crate) num_pq_chunks: NonZeroUsize,
     #[cfg(feature = "disk-index")]
     pub(crate) quantization_type: QuantizationType,
+    /// Build algorithm: "Vamana" (default) or "PiPNN" with config params.
+    #[cfg(feature = "disk-index")]
+    #[serde(default)]
+    pub(crate) build_algorithm: BuildAlgorithm,
     pub(crate) save_path: String,
 }
 
@@ -257,6 +263,8 @@ impl Example for DiskIndexOperation {
             num_pq_chunks: NonZeroUsize::new(16).unwrap(),
             #[cfg(feature = "disk-index")]
             quantization_type: QuantizationType::PQ { num_chunks: 16 },
+            #[cfg(feature = "disk-index")]
+            build_algorithm: BuildAlgorithm::default(),
             save_path: "sample_index_l50_r32".to_string(),
         };
 
@@ -351,6 +359,8 @@ impl DiskIndexBuild {
                 }
             }
         }
+        #[cfg(feature = "disk-index")]
+        write_field!(f, "Build Algorithm", self.build_algorithm)?;
         write_field!(f, "Save Path", self.save_path)?;
         Ok(())
     }

diskann-disk/Cargo.toml

@@ -45,6 +45,7 @@ vfs = { workspace = true }
 
 # Optional dependencies
 opentelemetry = { workspace = true, optional = true }
+diskann-pipnn = { workspace = true, optional = true }
 
 [target.'cfg(target_os = "linux")'.dependencies]
 io-uring = "0.6.4"
@@ -54,6 +55,7 @@ libc = "0.2.148"
 rstest.workspace = true
 tempfile.workspace = true
 vfs.workspace = true
+serde_json.workspace = true
 diskann-providers = { workspace = true, default-features = false, features = [
     "testing",
     "virtual_storage",
@@ -66,6 +68,7 @@ diskann = { workspace = true }
 [features]
 default = []
 perf_test = ["dep:opentelemetry"]
+pipnn = ["dep:diskann-pipnn"]
 virtual_storage = ["diskann-providers/virtual_storage"]
 experimental_diversity_search = [
     "diskann/experimental_diversity_search",
@@ -82,4 +85,4 @@ harness = false
 
 # Some 'cfg's in the source tree will be flagged by `cargo clippy -j 2 --workspace --no-deps --all-targets -- -D warnings`
 [lints.rust]
-unexpected_cfgs = { level = "warn", check-cfg = ['cfg(coverage)'] }
+unexpected_cfgs = { level = "warn", check-cfg = ['cfg(coverage)', 'cfg(feature, values("pipnn"))'] }

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

@@ -53,6 +53,7 @@ use crate::{
             },
             continuation::{process_while_resource_is_available_async, ChunkingConfig},
         },
+        configuration::build_algorithm::BuildAlgorithm,
     },
     storage::{
         quant::{GeneratorContext, PQGeneration, PQGenerationContext, QuantDataGenerator},
@@ -235,15 +236,35 @@ where
             self.index_configuration.num_threads
         );
 
+        let t_pq = std::time::Instant::now();
         self.generate_compressed_data(&pool).await?;
         logger.log_checkpoint(DiskIndexBuildCheckpoint::PqConstruction);
+        let pq_secs = t_pq.elapsed().as_secs_f64();
 
+        let t_index = std::time::Instant::now();
         self.build_inmem_index(&pool).await?;
         logger.log_checkpoint(DiskIndexBuildCheckpoint::InmemIndexBuild);
+        let index_secs = t_index.elapsed().as_secs_f64();
+
+        // Return freed memory (f32 data, graph, PiPNN internals) to the OS
+        // before disk layout starts. Without this, ~1.7 GB of freed-but-retained
+        // memory inflates peak RSS during the disk layout phase.
+        #[cfg(target_os = "linux")]
+        unsafe {
+            extern "C" { fn malloc_trim(pad: usize) -> i32; }
+            malloc_trim(0);
+        }
 
         // Use physical file to pass the memory index to the disk writer
+        let t_layout = std::time::Instant::now();
         self.create_disk_layout()?;
         logger.log_checkpoint(DiskIndexBuildCheckpoint::DiskLayout);
+        let layout_secs = t_layout.elapsed().as_secs_f64();
+
+        println!("Disk Index Build Phases");
+        println!("  PQ compression: {:.3}s", pq_secs);
+        println!("  Graph build:    {:.3}s", index_secs);
+        println!("  Disk layout:    {:.3}s", layout_secs);
 
         Ok(())
     }
@@ -313,6 +334,22 @@ where
     }
 
     async fn build_inmem_index(&mut self, pool: &RayonThreadPool) -> ANNResult<()> {
+        // Check for PiPNN algorithm
+        #[cfg(feature = "pipnn")]
+        if let BuildAlgorithm::PiPNN { .. } = self.disk_build_param.build_algorithm() {
+            return self.build_pipnn_index().await;
+        }
+
+        #[cfg(not(feature = "pipnn"))]
+        if !matches!(
+            self.disk_build_param.build_algorithm(),
+            BuildAlgorithm::Vamana
+        ) {
+            return Err(ANNError::log_index_error(
+                "PiPNN build algorithm requires the 'pipnn' feature to be enabled",
+            ));
+        }
+
         match determine_build_strategy::<Data>(
             &self.index_configuration,
             self.disk_build_param.build_memory_limit().in_bytes() as f64,
@@ -326,6 +363,91 @@ where
         }
     }
 
+    #[cfg(feature = "pipnn")]
+    async fn build_pipnn_index(&mut self) -> ANNResult<()> {
+        use diskann_pipnn::builder;
+
+        let config = self.disk_build_param.build_algorithm()
+            .to_pipnn_config(
+                self.index_configuration.config.pruned_degree().get(),
+                self.index_configuration.dist_metric,
+                self.index_configuration.config.alpha(),
+            )
+            .ok_or_else(|| ANNError::log_index_error(
+                "build_pipnn_index called but build algorithm is not PiPNN"
+            ))?;
+
+        config.validate().map_err(|e| {
+            ANNError::log_index_error(format!("PiPNN config error: {}", e))
+        })?;
+
+        info!(
+            "Building PiPNN index: max_degree={}",
+            config.max_degree
+        );
+
+        let data_path = self.index_writer.get_dataset_file();
+
+        // Build the PiPNN graph, using pre-trained SQ if available.
+        let graph = match &self.build_quantizer {
+            BuildQuantizer::Scalar1Bit(with_bits) => {
+                // SQ path needs f32 data for quantize_1bit.
+                let (npoints, ndims, data) = load_data_as_f32::<Data::VectorDataType, _>(
+                    &data_path,
+                    self.storage_provider,
+                )?;
+                // Use the DiskANN-trained ScalarQuantizer for 1-bit quantization.
+                // This ensures identical quantization between Vamana and PiPNN builds.
+                let sq = with_bits.quantizer();
+                let scale = sq.scale();
+                let inverse_scale = if scale == 0.0 { 1.0 } else { 1.0 / scale };
+                let sq_params = builder::SQParams {
+                    shift: sq.shift().to_vec(),
+                    inverse_scale,
+                };
+                info!("Using pre-trained SQ quantizer for PiPNN 1-bit build");
+                builder::build_with_sq(&data, npoints, ndims, &config, &sq_params)
+                    .map_err(|e| ANNError::log_index_error(format!("PiPNN build failed: {}", e)))?
+            }
+            _ => {
+                // Full precision or PQ build quantization — load data in native type
+                // and use build_typed to avoid upfront f32 conversion (saves ~793 MB
+                // peak RSS for f16 data).
+                let (npoints, ndims, data) = load_data_typed::<Data::VectorDataType, _>(
+                    &data_path,
+                    self.storage_provider,
+                )?;
+                builder::build_typed(&data, npoints, ndims, &config)
+                    .map_err(|e| ANNError::log_index_error(format!("PiPNN build failed: {}", e)))?
+            }
+        };
+
+        let save_path = self.index_writer.get_mem_index_file();
+        graph.save_graph(std::path::Path::new(&save_path))
+            .map_err(|e| ANNError::log_index_error(format!("PiPNN graph save failed: {}", e)))?;
+
+        info!(
+            "PiPNN build complete: avg_degree={:.1}, max_degree={}, isolated={}, total={:.3}s",
+            graph.avg_degree(),
+            graph.max_degree(),
+            graph.num_isolated(),
+            graph.build_stats.total_secs
+        );
+        // Print timing breakdown to stdout (tracing goes to OpenTelemetry spans,
+        // not stdout, so use print! for user-visible output like Vamana does).
+        print!("{}", graph.build_stats);
+
+        // Mark checkpoint stages as complete so the checkpoint system is consistent.
+        self.checkpoint_record_manager.execute_stage(
+            WorkStage::InMemIndexBuild,
+            WorkStage::WriteDiskLayout,
+            || Ok(()),
+            || Ok(()),
+        )?;
+
+        Ok(())
+    }
+
     async fn build_merged_vamana_index(&mut self, pool: &RayonThreadPool) -> ANNResult<()> {
         let mut logger = PerfLogger::new_disk_index_build_logger();
         let mut workflow = MergedVamanaIndexWorkflow::new(self, pool);
@@ -480,6 +602,51 @@ where
     }
 }
 
+#[cfg(feature = "pipnn")]
+fn load_data_as_f32<T, SP>(
+    data_path: &str,
+    storage_provider: &SP,
+) -> ANNResult<(usize, usize, Vec<f32>)>
+where
+    T: VectorRepr,
+    SP: StorageReadProvider,
+{
+    let matrix = read_bin::<T>(&mut storage_provider.open_reader(data_path)?)?;
+    let npoints = matrix.nrows();
+    let ndims = matrix.ncols();
+
+    // Convert to f32
+    let mut f32_data = vec![0.0f32; npoints * ndims];
+    for i in 0..npoints {
+        let src = matrix.row(i);
+        let dst = &mut f32_data[i * ndims..(i + 1) * ndims];
+        T::as_f32_into(src, dst)
+            .map_err(|e| ANNError::log_index_error(format!("Data conversion error: {}", e)))?;
+    }
+
+    Ok((npoints, ndims, f32_data))
+}
+
+/// Load data in its native type T without converting to f32.
+/// This avoids doubling memory for f16 data by keeping it as f16 in memory
+/// and converting to f32 on-the-fly at each access point inside PiPNN.
+#[cfg(feature = "pipnn")]
+fn load_data_typed<T, SP>(
+    data_path: &str,
+    storage_provider: &SP,
+) -> ANNResult<(usize, usize, Vec<T>)>
+where
+    T: VectorRepr,
+    SP: StorageReadProvider,
+{
+    let matrix = read_bin::<T>(&mut storage_provider.open_reader(data_path)?)?;
+    let npoints = matrix.nrows();
+    let ndims = matrix.ncols();
+    let data: Vec<T> = matrix.into_inner().into_vec();
+
+    Ok((npoints, ndims, data))
+}
+
 #[allow(clippy::too_many_arguments)]
 async fn build_inmem_index<T, StorageProvider>(
     config: IndexConfiguration,