PiPNN 6/6: add HashPrune candidate merging

.github/workflows/ci.yml

@@ -70,6 +70,7 @@ jobs:
       - test-workspace
       - test-workspace-features
       - coverage
+      - miri-pipnn
       - vectorset-clippy
       - vectorset-fmt
       - vectorset-build
@@ -93,6 +94,35 @@ jobs:
     steps:
       - run: exit 0
 
+  miri-pipnn:
+    needs: basics
+    name: strict-provenance PiPNN kernels
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install Rust nightly with Miri
+        run: rustup toolchain install nightly --component miri
+
+      - uses: Swatinem/rust-cache@v2
+
+      - name: Run HashPrune pointer boundaries
+        env:
+          MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance
+        run: |
+          cargo +nightly miri test --locked -p diskann --features pipnn --lib \
+            graph::pipnn::hash_prune::tests::find_hash_handles_padded_boundaries_and_all_bit_patterns \
+            -- --exact
+          cargo +nightly miri test --locked -p diskann --features pipnn --lib \
+            graph::pipnn::hash_prune::tests::relative_hash_matches_numeric_reference \
+            -- --exact
+          cargo +nightly miri test --locked -p diskann --features pipnn --lib \
+            graph::pipnn::hash_prune::tests::full_reservoir_evicts_the_farthest_candidate \
+            -- --exact
+          cargo +nightly miri test --locked -p diskann --features pipnn --lib \
+            graph::pipnn::hash_prune::tests::hot_slot_contention_serializes_state_mutation \
+            -- --exact
+
   fmt:
     name: format check
     runs-on: ubuntu-latest

diskann-benchmark/src/index/build.rs

@@ -142,12 +142,15 @@ where
 
     let started = std::time::Instant::now();
     let adjacency = {
-        let context = diskann::graph::pipnn::PiPNNBuildContext::new(
+        let mut context = diskann::graph::pipnn::PiPNNBuildContext::new(
             parameters.into(),
             &graph,
             metric,
             &pool,
         )?;
+        if let Some(hash_prune) = &parameters.hash_prune {
+            context = context.with_hash_prune(hash_prune.into())?;
+        }
         diskann::graph::pipnn::build_graph(data.as_view(), &context)?
     };
     let start_points = input

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

@@ -77,8 +77,12 @@ where
         index_writer: DiskIndexWriter,
     ) -> ANNResult<Self> {
         #[cfg(feature = "pipnn")]
-        if let Some(config) = disk_build_param.pipnn_config() {
-            config.validate()?;
+        if let Some(parameters) = disk_build_param.pipnn_parameters() {
+            diskann::graph::pipnn::PiPNNConfig::from(parameters).validate()?;
+            if let Some(hash_prune) = &parameters.hash_prune {
+                diskann::graph::pipnn::HashPruneConfig::from(hash_prune)
+                    .validate_for_degree(index_configuration.config.pruned_degree().get())?;
+            }
         }
 
         let pq_storage = PQStorage::new(
@@ -182,8 +186,8 @@ where
 
     async fn build_graph(&mut self, pool: RayonThreadPoolRef<'_>) -> ANNResult<()> {
         #[cfg(feature = "pipnn")]
-        if let Some(config) = self.disk_build_param.pipnn_config() {
-            return pipnn::build_graph(self, pool, config);
+        if let Some(parameters) = self.disk_build_param.pipnn_parameters().cloned() {
+            return pipnn::build_graph(self, pool, &parameters);
         }
 
         match determine_build_strategy::<Data>(

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

@@ -11,7 +11,7 @@
 //!
 //! PiPNN and Vamana use the same disk graph format.
 
-use diskann::graph::pipnn::{PiPNNBuildContext, PiPNNConfig};
+use diskann::graph::pipnn::PiPNNBuildContext;
 use diskann::{utils::VectorRepr, ANNError, ANNResult};
 use diskann_providers::{
     storage::{save_adjacency_graph, StorageReadProvider, StorageWriteProvider},
@@ -20,13 +20,13 @@ use diskann_providers::{
 use diskann_utils::io::{read_bin, Metadata};
 
 use super::{u32_try_from, DiskIndexBuilder};
-use crate::data_model::GraphDataType;
+use crate::{data_model::GraphDataType, PiPNNParameters};
 
 /// Build PiPNN adjacency and persist it through the canonical disk graph writer.
 pub(super) fn build_graph<Data, StorageProvider>(
     builder: &DiskIndexBuilder<'_, Data, StorageProvider>,
     pool: RayonThreadPoolRef<'_>,
-    config: PiPNNConfig,
+    parameters: &PiPNNParameters,
 ) -> ANNResult<()>
 where
     Data: GraphDataType<VectorIdType = u32>,
@@ -55,12 +55,15 @@ where
     // supplied Rayon pool.
     let data =
         read_bin::<Data::VectorDataType>(&mut builder.storage_provider.open_reader(&data_path)?)?;
-    let context = PiPNNBuildContext::new(
-        config,
+    let mut context = PiPNNBuildContext::new(
+        parameters.into(),
         &builder.index_configuration.config,
         builder.index_configuration.dist_metric,
         pool.as_rayon(),
     )?;
+    if let Some(hash_prune) = &parameters.hash_prune {
+        context = context.with_hash_prune(hash_prune.into())?;
+    }
     let adjacency = diskann::graph::pipnn::build_graph(data.as_view(), &context)?;
 
     // The disk header requires a start point. Use the same sampled medoid policy
@@ -116,6 +119,7 @@ mod tests {
             fanout: vec![10, 3],
             k: 2,
             replicas: 1,
+            hash_prune: Some(crate::HashPruneParameters::default()),
         }
     }
 
@@ -199,7 +203,7 @@ mod tests {
         let builder = builder(&storage, 3, 8, 1.0, 1.2, parameters.clone());
         let pool = create_thread_pool(1).unwrap();
 
-        let error = super::build_graph(&builder, pool.as_ref(), (&parameters).into()).unwrap_err();
+        let error = super::build_graph(&builder, pool.as_ref(), &parameters).unwrap_err();
         assert!(format!("{error:?}").contains("configured point count 3"));
         assert!(!storage.exists(&builder.index_writer.get_mem_index_file()));
     }
@@ -213,7 +217,7 @@ mod tests {
         let builder = builder(&storage, points, dimensions, 1.0, 1.2, parameters.clone());
         let pool = create_thread_pool(1).unwrap();
 
-        super::build_graph(&builder, pool.as_ref(), (&parameters).into()).unwrap();
+        super::build_graph(&builder, pool.as_ref(), &parameters).unwrap();
 
         let mut header = [0_u8; 24];
         std::io::Read::read_exact(
@@ -274,4 +278,34 @@ mod tests {
 
         assert!(format!("{error:?}").contains("c_max must be greater than zero"));
     }
+
+    #[test]
+    fn builder_rejects_hash_prune_capacity_before_quantizer_artifacts() {
+        let storage = VirtualStorageProvider::new_memory();
+        let parameters = PiPNNParameters {
+            hash_prune: Some(crate::HashPruneParameters {
+                num_hash_planes: 12,
+                l_max: 16,
+                final_prune: true,
+            }),
+            ..PiPNNParameters::default()
+        };
+        let params = DiskIndexBuildParameters::new_pipnn(
+            MemoryBudget::try_from_gb(1.0).unwrap(),
+            NumPQChunks::new_with(1, 1).unwrap(),
+            parameters,
+        );
+        let config = IndexConfiguration::new(Metric::L2, 1, 1, ONE, 1, graph_config(32, 1.2));
+        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!("HashPrune capacity below graph degree must be rejected"),
+            Err(error) => error,
+        };
+
+        assert!(format!("{error:?}").contains("must be at least the graph degree (32)"));
+        assert!(!storage.exists("/index_pq_pivots.bin"));
+        assert!(!storage.exists("/index_pq_compressed.bin"));
+    }
 }

diskann-disk/src/build/configuration/build_algorithm.rs

@@ -29,6 +29,43 @@ pub struct PiPNNParameters {
     pub k: usize,
     /// Number of independent partition passes.
     pub replicas: usize,
+    /// HashPrune policy. `None` keeps all unique direct candidates.
+    pub hash_prune: Option<HashPruneParameters>,
+}
+
+/// HashPrune parameters in the JSON build configuration.
+#[cfg(feature = "pipnn")]
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[serde(default, deny_unknown_fields)]
+pub struct HashPruneParameters {
+    /// Number of random-hyperplane sketch dimensions.
+    pub num_hash_planes: usize,
+    /// Maximum number of candidates retained per point.
+    pub l_max: usize,
+    /// Apply Vamana RobustPrune after reservoir extraction.
+    pub final_prune: bool,
+}
+
+#[cfg(feature = "pipnn")]
+impl Default for HashPruneParameters {
+    fn default() -> Self {
+        Self {
+            num_hash_planes: 12,
+            l_max: 64,
+            final_prune: true,
+        }
+    }
+}
+
+#[cfg(feature = "pipnn")]
+impl From<&HashPruneParameters> for diskann::graph::pipnn::HashPruneConfig {
+    fn from(config: &HashPruneParameters) -> Self {
+        Self {
+            num_hash_planes: config.num_hash_planes,
+            l_max: config.l_max,
+            final_prune: config.final_prune,
+        }
+    }
 }
 
 #[cfg(feature = "pipnn")]
@@ -41,6 +78,7 @@ impl Default for PiPNNParameters {
             fanout: vec![8, 3],
             k: 2,
             replicas: 1,
+            hash_prune: None,
         }
     }
 }
@@ -116,6 +154,15 @@ mod tests {
         assert_eq!(config.fanout, [10, 3]);
         assert_eq!(config.k, 3);
         assert_eq!(config.replicas, 1);
+        assert_eq!(config.hash_prune, None);
+
+        let explicit: BuildAlgorithm =
+            serde_json::from_str(r#"{"algorithm":"PiPNN","hash_prune":{}}"#).unwrap();
+        let BuildAlgorithm::PiPNN(explicit) = explicit else {
+            panic!("expected PiPNN");
+        };
+        assert_eq!(explicit.hash_prune, Some(HashPruneParameters::default()));
+
         assert!(
             serde_json::from_str::<BuildAlgorithm>(r#"{"algorithm":"PiPNN","l_max":72}"#).is_err()
         );

diskann-disk/src/build/configuration/disk_index_build_parameter.rs

@@ -198,9 +198,9 @@ impl DiskIndexBuildParameters {
     }
 
     #[cfg(feature = "pipnn")]
-    pub(crate) fn pipnn_config(&self) -> Option<diskann::graph::pipnn::PiPNNConfig> {
+    pub(crate) fn pipnn_parameters(&self) -> Option<&PiPNNParameters> {
         match &self.build_algorithm {
-            BuildAlgorithm::PiPNN(config) => Some(config.into()),
+            BuildAlgorithm::PiPNN(config) => Some(config),
             BuildAlgorithm::Vamana => None,
         }
     }

diskann-disk/src/build/configuration/mod.rs

@@ -5,7 +5,7 @@
 pub mod build_algorithm;
 pub use build_algorithm::BuildAlgorithm;
 #[cfg(feature = "pipnn")]
-pub use build_algorithm::PiPNNParameters;
+pub use build_algorithm::{HashPruneParameters, PiPNNParameters};
 
 pub mod disk_index_build_parameter;
 pub use disk_index_build_parameter::{DiskIndexBuildParameters, MemoryBudget, NumPQChunks};

diskann-disk/src/build/mod.rs

@@ -12,9 +12,9 @@ pub mod builder;
 pub mod configuration;
 
 // Re-export key types for convenience
-#[cfg(feature = "pipnn")]
-pub use configuration::PiPNNParameters;
 pub use configuration::{
     disk_index_build_parameter, filter_parameter, BuildAlgorithm, DiskIndexBuildParameters,
     QuantizationType,
 };
+#[cfg(feature = "pipnn")]
+pub use configuration::{HashPruneParameters, PiPNNParameters};

diskann-disk/src/lib.rs

@@ -14,12 +14,12 @@ pub(crate) mod test_utils;
 pub mod error;
 
 pub mod build;
-#[cfg(feature = "pipnn")]
-pub use build::PiPNNParameters;
 pub use build::{
     disk_index_build_parameter, filter_parameter, BuildAlgorithm, DiskIndexBuildParameters,
     QuantizationType,
 };
+#[cfg(feature = "pipnn")]
+pub use build::{HashPruneParameters, PiPNNParameters};
 
 pub mod data_model;
 pub mod search;