[Experiment do not review - for collab with Magdalen] Decoupled yield_k knob for Design B attribute-diversity search

assets/diverse-search/chart_data.json

@@ -0,0 +1,52 @@
+{
+  "_comment": "QPS-vs-recall data at L=100 for the diverse-search benchmark charts. Recall is deterministic and authoritative; QPS is the best (least-contended) of several single-threaded runs (see diverse-search-results.md). Update these values when rerunning benchmarks, then regenerate the PNGs with plot_qps_vs_recall.py.",
+  "methods": ["Standard", "Queue", "Design A", "Design B"],
+  "colors": {
+    "Standard": "#7f7f7f",
+    "Queue": "#1f77b4",
+    "Design A": "#2ca02c",
+    "Design B": "#d62728"
+  },
+  "charts": [
+    {
+      "output": "enron-qps-vs-recall.png",
+      "title": "Enron (concentrated attribute)",
+      "points": {
+        "Standard": {"recall": 70.72, "qps": 142.6},
+        "Queue": {"recall": 77.92, "qps": 100.0},
+        "Design A": {"recall": 86.05, "qps": 159.2},
+        "Design B": {"recall": 93.01, "qps": 53.0}
+      }
+    },
+    {
+      "output": "caselaw-qps-vs-recall.png",
+      "title": "Caselaw (diffuse attribute \u2014 doc_id)",
+      "points": {
+        "Standard": {"recall": 96.91, "qps": 160.3},
+        "Queue": {"recall": 98.32, "qps": 84.3},
+        "Design A": {"recall": 99.25, "qps": 128.8},
+        "Design B": {"recall": 99.25, "qps": 105.9}
+      }
+    },
+    {
+      "output": "caselaw-jurisdiction-qps-vs-recall.png",
+      "title": "Caselaw (concentrated + distance-correlated \u2014 court_jurisdiction)",
+      "points": {
+        "Standard": {"recall": 35.91, "qps": 165.6},
+        "Queue": {"recall": 79.12, "qps": 152.1},
+        "Design A": {"recall": 72.59, "qps": 174.5},
+        "Design B": {"recall": 78.97, "qps": 96.0}
+      }
+    },
+    {
+      "output": "yfcc-qps-vs-recall.png",
+      "title": "YFCC (concentrated attribute \u2014 camera)",
+      "points": {
+        "Standard": {"recall": 49.73, "qps": 227.1},
+        "Queue": {"recall": 81.78, "qps": 186.1},
+        "Design A": {"recall": 93.98, "qps": 192.6},
+        "Design B": {"recall": 96.92, "qps": 111.9}
+      }
+    }
+  ]
+}

assets/diverse-search/plot_qps_vs_recall.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+"""Generate the QPS-vs-recall scatter charts for the diverse-search benchmark doc.
+
+Reads chart_data.json (recall/QPS at L=100 for each method and dataset) and
+writes one PNG per dataset next to this script. Kept in the repo so the charts
+can be regenerated whenever the benchmark numbers change.
+
+Usage:
+    python plot_qps_vs_recall.py
+
+Requires matplotlib (pip install matplotlib).
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+
+HERE = Path(__file__).resolve().parent
+DATA_FILE = HERE / "chart_data.json"
+
+
+def main() -> None:
+    data = json.loads(DATA_FILE.read_text(encoding="utf-8"))
+    methods = data["methods"]
+    colors = data["colors"]
+
+    for chart in data["charts"]:
+        fig, ax = plt.subplots(figsize=(9, 6))
+
+        for method in methods:
+            point = chart["points"][method]
+            recall = point["recall"]
+            qps = point["qps"]
+            ax.scatter(
+                recall,
+                qps,
+                s=160,
+                color=colors[method],
+                label=method,
+                zorder=3,
+            )
+            ax.annotate(
+                method,
+                (recall, qps),
+                textcoords="offset points",
+                xytext=(8, 8),
+                color=colors[method],
+                fontsize=10,
+            )
+
+        ax.set_title(f"{chart['title']}\nQPS vs recall at L=100")
+        ax.set_xlabel("Recall @10 (%)  -  higher is better")
+        ax.set_ylabel("QPS  -  higher is better")
+        ax.grid(True, color="0.9", zorder=0)
+        ax.legend(title="Approach", loc="lower left")
+
+        out_path = HERE / chart["output"]
+        fig.tight_layout()
+        fig.savefig(out_path, dpi=100)
+        plt.close(fig)
+        print(f"wrote {out_path}")
+
+
+if __name__ == "__main__":
+    main()

diskann-benchmark/src/disk_index/search.rs

@@ -37,7 +37,7 @@ use serde::{Deserialize, Serialize};
 use crate::{
     disk_index::json_spancollector::JsonSpanCollector,
     inputs::disk::{DiskIndexLoad, DiskSearchPhase},
-    utils::{datafiles, SimilarityMeasure},
+    utils::{attributes::FileAttributeProvider, datafiles, SimilarityMeasure},
 };
 
 #[derive(Serialize, Deserialize, Debug)]
@@ -236,6 +236,35 @@ where
 
     logger.log_checkpoint("index_loaded");
 
+    // Optional attribute-bucket diversity (Design A): loaded once and shared
+    // across all queries and search-list values.
+    let diverse_attribute_provider = match &search_params.attributes {
+        Some(attributes_file) => Some(std::sync::Arc::new(FileAttributeProvider::load(
+            attributes_file,
+        )?)),
+        None => None,
+    };
+    let diverse_results_k = search_params.diverse_results_k;
+    let diverse_attribute_id = search_params.diverse_attribute_id;
+    // Effective output cap (per-bucket cap on the final result set).
+    let effective_diverse_results_k =
+        diverse_results_k.unwrap_or(search_params.recall_at as usize);
+    // Design B: when the diverse path is active and `adaptive_l` is configured,
+    // grow `L` per query from the observed bucket concentration.
+    let diverse_adaptive_l = search_params.search_mode.adaptive_l.as_ref().map(|a| {
+        diskann::graph::search::AdaptiveL::new(a.sample_count.into(), a.scale_factor)
+            .expect("validated adaptive L must construct")
+    });
+    // Per-bucket cap for the yield sampler; defaults to the output cap when the
+    // config does not override it.
+    let diverse_yield_k = search_params
+        .search_mode
+        .adaptive_l
+        .as_ref()
+        .and_then(|a| a.yield_k)
+        .map(usize::from)
+        .unwrap_or(effective_diverse_results_k);
+
     let pool = create_thread_pool(search_params.num_threads)?;
     let mut search_results_per_l = Vec::with_capacity(search_params.search_list.len());
     let has_any_search_failed = AtomicBool::new(false);
@@ -271,12 +300,23 @@ where
                 // Construct the SearchMode from the JSON-driven
                 // `adaptive_l` is now encapsulated in `DiskSearchMode`, so the
                 // benchmark only supplies the per-query filter and post-processor.
-                let has_filter = search_params.vector_filters_file.is_some();
-                let mode: SearchMode<'_> = search_params.search_mode.search_mode(
-                    has_filter,
-                    vf,
-                    search_params.post_processor.as_ref(),
-                );
+                let mode: SearchMode<'_> = match diverse_attribute_provider.as_ref() {
+                    Some(provider) => SearchMode::diverse_attribute(
+                        provider.clone(),
+                        diverse_attribute_id,
+                        effective_diverse_results_k,
+                        diverse_yield_k,
+                        diverse_adaptive_l.clone(),
+                    ),
+                    None => {
+                        let has_filter = search_params.vector_filters_file.is_some();
+                        search_params.search_mode.search_mode(
+                            has_filter,
+                            vf,
+                            search_params.post_processor.as_ref(),
+                        )
+                    }
+                };
 
                 match searcher.search(
                     q,

diskann-benchmark/src/inputs/disk.rs

@@ -170,6 +170,20 @@ pub(crate) struct DiskSearchPhase {
     pub(crate) num_nodes_to_cache: Option<usize>,
     pub(crate) search_io_limit: Option<usize>,
     pub(crate) post_processor: Option<TopkPostProcessor>,
+    /// Plaintext file with one integer attribute per line; line `N` is vector `N`'s attribute.
+    ///
+    /// When present, the search phase runs attribute-bucket diversity via
+    /// post-processing (Design A): a plain greedy search over the top-`L` pool
+    /// followed by bucket selection keeping at most `diverse_results_k` results
+    /// per distinct attribute value, instead of the mode selected by
+    /// `search_mode`/`post_processor`.
+    pub(crate) attributes: Option<InputFile>,
+    /// The attribute dimension used for diversity (currently only a single attribute is used).
+    #[serde(default)]
+    pub(crate) diverse_attribute_id: usize,
+    /// The maximum number of results to keep per distinct attribute value.
+    #[serde(default)]
+    pub(crate) diverse_results_k: Option<usize>,
 }
 
 /////////
@@ -320,6 +334,18 @@ impl DiskSearchPhase {
                 .context("invalid disk search post processor")?;
         }
 
+        if let Some(attributes) = self.attributes.as_mut() {
+            attributes
+                .resolve(checker)
+                .context("invalid attributes file")?;
+            match self.diverse_results_k {
+                Some(k) if k > 0 => {}
+                _ => anyhow::bail!(
+                    "diverse_results_k must be a positive integer when attributes is set"
+                ),
+            }
+        }
+
         Ok(())
     }
 }
@@ -367,6 +393,9 @@ impl Example for DiskIndexOperation {
             num_nodes_to_cache: None,
             search_io_limit: None,
             post_processor: None,
+            attributes: None,
+            diverse_attribute_id: 0,
+            diverse_results_k: None,
         };
 
         Self {

diskann-benchmark/src/inputs/graph_index.rs

@@ -248,6 +248,12 @@ impl MultihopFilterSearchPhase {
 pub(crate) struct AdaptiveL {
     pub(crate) sample_count: NonZeroUsize,
     pub(crate) scale_factor: f64,
+    /// Per-bucket cap used by the Design-B diversity sampler when estimating
+    /// yield. Decoupled from `diverse_results_k` (the output cap); when unset it
+    /// defaults to the effective `diverse_results_k`. Only used by attribute
+    /// diversity search.
+    #[serde(default)]
+    pub(crate) yield_k: Option<NonZeroUsize>,
 }
 
 impl AdaptiveL {

diskann-benchmark/src/utils/attributes.rs

@@ -0,0 +1,91 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+//! A file-backed [`AttributeValueProvider`] for diversity-aware benchmarks.
+
+use std::path::Path;
+
+use anyhow::Context;
+use diskann::{neighbor::AttributeValueProvider, provider::HasId};
+
+/// The reserved attribute bucket assigned to graph navigation nodes (frozen start points)
+/// that fall outside the range of loaded per-vector attributes.
+///
+/// Greedy graph search seeds traversal from the index's start points. Those points must
+/// therefore carry an attribute, otherwise the diverse queue would drop them and search
+/// would never expand beyond the entry node. Assigning them a dedicated bucket keeps them
+/// traversable without merging them into any real attribute group.
+const NAVIGATION_BUCKET: u32 = u32::MAX;
+
+/// An attribute value provider backed by a plaintext attribute file.
+///
+/// The file is expected to contain one unsigned integer per line, where the value on the
+/// `N`-th line (0-indexed) is the diversity attribute of the `N`-th vector. This matches the
+/// on-disk layout produced by the labelling tools used elsewhere in the pipeline.
+///
+/// Ids outside the range of loaded attributes (for example the graph's frozen start points)
+/// are mapped to a reserved [`NAVIGATION_BUCKET`] so that greedy search can still traverse
+/// the graph through them.
+#[derive(Debug, Clone)]
+pub(crate) struct FileAttributeProvider {
+    attributes: Vec<u32>,
+    num_buckets: usize,
+}
+
+impl FileAttributeProvider {
+    /// Load attributes from the plaintext file at `path`.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the file cannot be read or if any line fails to parse as a `u32`.
+    pub(crate) fn load(path: &Path) -> anyhow::Result<Self> {
+        let contents = std::fs::read_to_string(path)
+            .with_context(|| format!("while reading attribute file {}", path.display()))?;
+
+        let attributes = contents
+            .lines()
+            .map(str::trim)
+            .filter(|line| !line.is_empty())
+            .enumerate()
+            .map(|(line, value)| {
+                value.parse::<u32>().with_context(|| {
+                    format!("invalid attribute value {value:?} on line {}", line + 1)
+                })
+            })
+            .collect::<anyhow::Result<Vec<u32>>>()?;
+
+        let num_buckets = attributes
+            .iter()
+            .copied()
+            .collect::<std::collections::HashSet<u32>>()
+            .len();
+
+        Ok(Self {
+            attributes,
+            num_buckets,
+        })
+    }
+}
+
+impl HasId for FileAttributeProvider {
+    type Id = u32;
+}
+
+impl AttributeValueProvider for FileAttributeProvider {
+    type Value = u32;
+
+    fn get(&self, id: Self::Id) -> Option<Self::Value> {
+        Some(
+            self.attributes
+                .get(id as usize)
+                .copied()
+                .unwrap_or(NAVIGATION_BUCKET),
+        )
+    }
+
+    fn num_buckets(&self) -> Option<usize> {
+        Some(self.num_buckets)
+    }
+}

diskann-benchmark/src/utils/mod.rs

@@ -9,6 +9,8 @@ use diskann_benchmark_runner::{
 };
 use serde::{Deserialize, Serialize};
 
+#[cfg(feature = "disk-index")]
+pub(crate) mod attributes;
 pub(crate) mod datafiles;
 pub(crate) mod filters;
 pub(crate) mod recall;

diskann-bftree/Cargo.toml

@@ -35,7 +35,6 @@ tokio = { workspace = true, features = ["full"] }
 
 [features]
 default = []
-experimental_diversity_search = ["diskann/experimental_diversity_search"]
 
 [lints.clippy]
 undocumented_unsafe_blocks = "warn"

diskann-disk/Cargo.toml

@@ -68,10 +68,6 @@ proptest.workspace = true
 default = []
 perf_test = ["dep:opentelemetry"]
 virtual_storage = ["diskann-providers/virtual_storage"]
-experimental_diversity_search = [
-    "diskann/experimental_diversity_search",
-    "diskann-providers/experimental_diversity_search",
-]
 
 [[bench]]
 name = "bench_main"