PiPNN 5/6: add dedicated benchmark pipelines

diskann-benchmark/Cargo.toml

@@ -83,3 +83,6 @@ disk-index = [
     "dep:opentelemetry_sdk",
     "dep:scopeguard",
 ]
+
+# Enable PiPNN graph construction.
+pipnn = ["diskann/pipnn", "diskann-disk/pipnn"]

diskann-benchmark/example/pipnn-disk-index.json

@@ -0,0 +1,45 @@
+{
+  "search_directories": [
+    "test_data/disk_index_search"
+  ],
+  "jobs": [
+    {
+      "type": "disk-index",
+      "content": {
+        "source": {
+          "disk-index-source": "Build",
+          "data_type": "float32",
+          "data": "disk_index_siftsmall_learn_256pts_data.fbin",
+          "distance": "squared_l2",
+          "dim": 128,
+          "max_degree": 32,
+          "l_build": 50,
+          "alpha": 1.2,
+          "num_threads": 1,
+          "build_ram_limit_gb": 2.0,
+          "num_pq_chunks": 128,
+          "build_algorithm": {
+            "algorithm": "PiPNN",
+            "c_max": 64,
+            "p_samp": 0.05,
+            "fanout": [4, 2],
+            "k": 2,
+            "replicas": 1
+          },
+          "save_path": "pipnn_disk_index"
+        },
+        "search_phase": {
+          "queries": "disk_index_sample_query_10pts.fbin",
+          "groundtruth": "disk_index_10pts_idx_uint32_exact_ground_truth.bin",
+          "search_list": [10, 20, 40],
+          "beam_width": 4,
+          "recall_at": 10,
+          "num_threads": 1,
+          "is_flat_search": false,
+          "distance": "squared_l2",
+          "vector_filters_file": null
+        }
+      }
+    }
+  ]
+}

diskann-benchmark/example/pipnn-graph-index.json

@@ -0,0 +1,46 @@
+{
+  "search_directories": [
+    "test_data/disk_index_search"
+  ],
+  "jobs": [
+    {
+      "type": "graph-index-build",
+      "content": {
+        "source": {
+          "index-source": "Build",
+          "data_type": "float32",
+          "data": "disk_index_siftsmall_learn_256pts_data.fbin",
+          "distance": "squared_l2",
+          "max_degree": 32,
+          "l_build": 50,
+          "alpha": 1.2,
+          "backedge_ratio": 1.0,
+          "num_threads": 1,
+          "start_point_strategy": "first_vector",
+          "build_algorithm": {
+            "algorithm": "PiPNN",
+            "c_max": 64,
+            "p_samp": 0.05,
+            "fanout": [4, 2],
+            "k": 2,
+            "replicas": 1
+          }
+        },
+        "search_phase": {
+          "search-type": "topk",
+          "queries": "disk_index_sample_query_10pts.fbin",
+          "groundtruth": "disk_index_10pts_idx_uint32_exact_ground_truth.bin",
+          "reps": 1,
+          "num_threads": [1],
+          "runs": [
+            {
+              "search_n": 10,
+              "search_l": [10, 20, 40],
+              "recall_k": 10
+            }
+          ]
+        }
+      }
+    }
+  ]
+}

diskann-benchmark/src/disk_index/build.rs

@@ -84,6 +84,7 @@ where
         params.l_build,
         metric.into(),
         |b| {
+            b.alpha(params.alpha);
             b.saturate_after_prune(true);
         },
     )
@@ -93,11 +94,23 @@ where
 
     let metadata = load_metadata_from_file(storage_provider, &data_path)?;
 
-    let build_parameters = DiskIndexBuildParameters::new(
-        MemoryBudget::try_from_gb(params.build_ram_limit_gb)?,
-        params.quantization_type,
-        NumPQChunks::new_with(params.num_pq_chunks.get(), metadata.ndims())?,
-    );
+    let search_pq_chunks = NumPQChunks::new_with(params.num_pq_chunks.get(), metadata.ndims())?;
+    let build_parameters = match &params.build_algorithm {
+        diskann_disk::BuildAlgorithm::Vamana => DiskIndexBuildParameters::new(
+            MemoryBudget::try_from_gb(params.build_ram_limit_gb)?,
+            params.quantization_type.ok_or_else(|| {
+                anyhow::anyhow!("quantization_type is required for Vamana builds")
+            })?,
+            search_pq_chunks,
+        ),
+        #[cfg(feature = "pipnn")]
+        diskann_disk::BuildAlgorithm::PiPNN(config) => DiskIndexBuildParameters::new_pipnn(
+            MemoryBudget::try_from_gb(params.build_ram_limit_gb)?,
+            search_pq_chunks,
+            config.clone(),
+        ),
+        _ => anyhow::bail!("unsupported disk graph-build algorithm"),
+    };
 
     let index_configuration = IndexConfiguration::new(
         metric,

diskann-benchmark/src/index/benchmarks.rs

@@ -212,26 +212,42 @@ where
         writeln!(output, "{}", input)?;
         let (index, build_stats) = match &input.source {
             IndexSource::Build(build) => {
-                let (index, build_stats) = run_build(
-                    build,
-                    common::FullPrecision,
-                    None,
-                    output,
-                    |data| {
-                        let index = diskann_async::new_index::<T, _>(
-                            build.try_as_config()?.build()?,
-                            build.inmem_parameters(data.nrows(), data.ncols()),
-                            common::NoDeletes,
-                        )?;
-                        build::set_start_points(
-                            index.provider(),
-                            data.as_view(),
-                            *build.start_point_strategy(),
-                        )?;
-                        Ok(index)
-                    },
-                    single_or_multi_insert,
-                )?;
+                let mut incremental = || {
+                    run_build(
+                        build,
+                        common::FullPrecision,
+                        None,
+                        output,
+                        |data| {
+                            let index = diskann_async::new_index::<T, _>(
+                                build.try_as_config()?.build()?,
+                                build.inmem_parameters(data.nrows(), data.ncols()),
+                                common::NoDeletes,
+                            )?;
+                            build::set_start_points(
+                                index.provider(),
+                                data.as_view(),
+                                *build.start_point_strategy(),
+                            )?;
+                            Ok(index)
+                        },
+                        single_or_multi_insert,
+                    )
+                };
+                // PiPNN uses the batch build path. Do not create an incremental
+                // provider or run Vamana insertion.
+                #[cfg(feature = "pipnn")]
+                let result = match build.build_algorithm() {
+                    diskann_disk::BuildAlgorithm::PiPNN(parameters) => {
+                        let data =
+                            Arc::new(datafiles::load_dataset(datafiles::BinFile(build.data()))?);
+                        build::pipnn_build(data, build, parameters)
+                    }
+                    _ => incremental(),
+                };
+                #[cfg(not(feature = "pipnn"))]
+                let result = incremental();
+                let (index, build_stats) = result?;
 
                 // save the index if requested
                 if let Some(save_path) = build.save_path() {