Skip to content

Remove k_value from Knn #1151

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
magdalendobson wants to merge 24 commits into
base: main
Choose a base branch
from
Open

Remove k_value from Knn #1151

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
97b36ef
finish up recall computation patch
May 14, 2026
07b3671
Merge branch 'main' of github.com:microsoft/DiskANN
May 15, 2026
eac3ffb
Merge branch 'main' of github.com:microsoft/DiskANN
May 19, 2026
17780f8
fix conflict
May 22, 2026
43eb517
fix conflict
May 22, 2026
1d3a52b
Merge branch 'main' of github.com:microsoft/DiskANN
May 25, 2026
17eac62
Merge branch 'main' of github.com:microsoft/DiskANN
May 26, 2026
0ab4baa
Merge branch 'main' of github.com:microsoft/DiskANN
May 27, 2026
4ddca60
Merge branch 'main' of github.com:microsoft/DiskANN
Jun 2, 2026
54ee01b
fix conflict
Jun 2, 2026
9e1743f
Merge branch 'main' of github.com:microsoft/DiskANN
Jun 5, 2026
93504e6
Merge branch 'main' of github.com:microsoft/DiskANN
Jun 8, 2026
b7c27ce
Merge branch 'main' of github.com:microsoft/DiskANN
Jun 11, 2026
acd1f8c
remove k_value from Knn within diskann crate
Jun 11, 2026
0202348
integrate changes into benchmark
Jun 12, 2026
a7b696f
fix merge conflict
Jun 12, 2026
0191fee
Potential fix for pull request finding
magdalendobson Jun 12, 2026
2f28f50
Potential fix for pull request finding
magdalendobson Jun 12, 2026
7ee7938
move changes to bftree and disk
Jun 12, 2026
10940ad
address changes in more crates
Jun 12, 2026
e0ca6f6
fmt
Jun 12, 2026
d664e44
fix merge conflict
Jul 24, 2026
dc0145e
changes requested in review
Jul 24, 2026
a5fa47a
fix problems, respond to other PR comments
Jul 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions diskann-benchmark-core/src/search/graph/inline.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use diskann::{
};
use diskann_utils::{future::AsyncFriendly, views::Matrix};

use crate::search::{self, Search, graph::Strategy};
use crate::search::{self, Search, graph::KnnParams, graph::Strategy};

/// A built-in helper for benchmarking filtered K-nearest neighbors search
/// using the inline search method.
Expand All @@ -22,7 +22,7 @@ use crate::search::{self, Search, graph::Strategy};
/// [`search::search_all`] is provided by the [`search::graph::knn::Aggregator`] type (same
/// aggregator as [`search::graph::knn::KNN`]).
///
/// The provided implementation of [`Search`] accepts [`graph::search::Knn`]
/// The provided implementation of [`Search`] accepts [`KnnParams`]
/// and returns [`search::graph::knn::Metrics`] as additional output.
#[derive(Debug)]
pub struct InlineFilterSearch<DP, T, S>
Expand Down Expand Up @@ -100,7 +100,7 @@ where
T: AsyncFriendly + Clone,
{
type Id = DP::ExternalId;
type Parameters = graph::search::Knn;
type Parameters = KnnParams;
type Output = super::knn::Metrics;

fn num_queries(&self) -> usize {
Expand All @@ -121,8 +121,8 @@ where
O: graph::SearchOutputBuffer<DP::ExternalId> + Send,
{
let context = DP::Context::default();
let inline_search =
graph::search::InlineFilterSearch::new(*parameters, self.adaptive_l.clone());
let knn = parameters.knn;
let inline_search = graph::search::InlineFilterSearch::new(knn, self.adaptive_l.clone());
let strategy =
labeled::Filtered::new(self.strategy.get(index)?.clone(), &*self.labels[index]);

Expand Down Expand Up @@ -199,7 +199,7 @@ mod tests {
let rt = crate::tokio::runtime(2).unwrap();
let results = search::search(
inline.clone(),
graph::search::Knn::new(nearest_neighbors, 10, None).unwrap(),
KnnParams::new(nearest_neighbors, 10).unwrap(),
NonZeroUsize::new(2).unwrap(),
&rt,
)
Expand Down Expand Up @@ -227,11 +227,11 @@ mod tests {
// Try the aggregated strategy.
let parameters = [
search::Run::new(
graph::search::Knn::new(nearest_neighbors, 10, None).unwrap(),
KnnParams::new(nearest_neighbors, 10).unwrap(),
setup.clone(),
),
search::Run::new(
graph::search::Knn::new(nearest_neighbors, 15, None).unwrap(),
KnnParams::new(nearest_neighbors, 15).unwrap(),
setup.clone(),
),
];
Expand Down
68 changes: 57 additions & 11 deletions diskann-benchmark-core/src/search/graph/knn.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,11 @@

//! A built-in helper for benchmarking K-nearest neighbors.

use std::sync::Arc;
use std::{num::NonZeroUsize, sync::Arc};
use thiserror::Error;

use diskann::{
ANNResult,
ANNError, ANNResult,
graph::{self, glue},
provider,
};
Expand All @@ -30,7 +31,7 @@ use crate::{
/// the latter. Result aggregation for [`search::search_all`] is provided
/// by the [`Aggregator`] type.
///
/// The provided implementation of [`Search`] accepts [`graph::search::Knn`]
/// The provided implementation of [`Search`] accepts [`KnnParams`]
/// and returns [`Metrics`] as additional output.
///
/// # Type Parameters
Expand Down Expand Up @@ -194,7 +195,7 @@ where
T: AsyncFriendly + Clone,
{
type Id = DP::ExternalId;
type Parameters = graph::search::Knn;
type Parameters = KnnParams;
type Output = Metrics;

fn num_queries(&self) -> usize {
Expand All @@ -215,7 +216,7 @@ where
O: graph::SearchOutputBuffer<DP::ExternalId> + Send,
{
let context = DP::Context::default();
let knn_search = *parameters;
let knn_search = parameters.knn;
let strategy = self.strategy.get(index)?;
let processor = self.post_processor.as_post_processor(strategy);

Expand All @@ -238,6 +239,51 @@ where
}
}

#[derive(Debug, Error)]
pub enum KnnParamsError {
#[error("k_value cannot be zero")]
KZero,
#[error("l_value ({l_value}) must be at least k_value ({k_value})")]
LLessThanK { l_value: usize, k_value: usize },
#[error("invalid KNN parameters")]
InvalidKnnParameters,
}
Comment thread
Copilot marked this conversation as resolved.

impl From<KnnParamsError> for ANNError {
#[track_caller]
fn from(err: KnnParamsError) -> Self {
ANNError::opaque(err)
}
}

/// A wrapper for the [`graph::search::Knn`] struct that also includes the `k` value.
#[derive(Debug, Copy, Clone)]
pub struct KnnParams {
k_value: NonZeroUsize,
pub knn: graph::search::Knn,
}

@hildebrandmw hildebrandmw Jul 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we change this to

pub struct KnnParams {
    num_neighbors: NonZeroUsize,
    knn: graph::search::Knn,
}

impl KnnParams {
  pub fn new(num_neighbors: NonZeroUsize, knn: graph::search::Knn) -> Result<Self, Error>;

  pub fn num_neighbors(&self) -> NonZeroUsize;
  pub fn knn(&self) -> &graph::search::Knn;

Reasons:

  • A pub field for knn undoes the validation check in the constructor.
  • Taking graph::search::Knn directly allows us to use a configurable beam width and simplifies the error.
  • num_neighbors is IMO slightly clearer.

@hildebrandmw hildebrandmw Jul 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Alternatively, the name Knn could work too since it's in a separate module, but in this case I think KnnParams is a little less overloaded. Maybe we could call it Params and access it as knn::Params? Any one of these works, but I'm not a fan of KnnWrapper.


impl KnnParams {
/// Construct a new [`KnnParams`].
pub fn new(k_value: usize, l_value: usize) -> Result<Self, KnnParamsError> {
let k_value = NonZeroUsize::new(k_value).ok_or(KnnParamsError::KZero)?;
if l_value < k_value.get() {
return Err(KnnParamsError::LLessThanK {
l_value,
k_value: k_value.get(),
});
}
Comment thread
Copilot marked this conversation as resolved.

let knn = graph::search::Knn::new(l_value, None)
.map_err(|_| KnnParamsError::InvalidKnnParameters)?;
Ok(Self { k_value, knn })
}

pub fn k_value(&self) -> NonZeroUsize {
self.k_value
}
}

/// An [`search::Aggregate`]d summary of multiple [`KNN`] search runs
/// returned by the provided [`Aggregator`].
///
Expand All @@ -249,7 +295,7 @@ pub struct Summary {
pub setup: search::Setup,

/// The [`Search::Parameters`] used for the batch of runs.
pub parameters: graph::search::Knn,
pub parameters: KnnParams,

/// The end-to-end latency for each repetition in the batch.
pub end_to_end_latencies: Vec<MicroSeconds>,
Expand Down Expand Up @@ -317,15 +363,15 @@ impl<'a, I> Aggregator<'a, I> {
}
}

impl<I> search::Aggregate<graph::search::Knn, I, Metrics> for Aggregator<'_, I>
impl<I> search::Aggregate<KnnParams, I, Metrics> for Aggregator<'_, I>
where
I: crate::recall::RecallCompatible,
{
type Output = Summary;

fn aggregate(
&mut self,
run: search::Run<graph::search::Knn>,
run: search::Run<KnnParams>,
mut results: Vec<search::SearchResults<I, Metrics>>,
) -> anyhow::Result<Summary> {
// Compute the recall using just the first result.
Expand Down Expand Up @@ -422,7 +468,7 @@ mod tests {
let rt = crate::tokio::runtime(2).unwrap();
let results = search::search(
knn.clone(),
graph::search::Knn::new(nearest_neighbors, 10, None).unwrap(),
KnnParams::new(nearest_neighbors, 10).unwrap(),
NonZeroUsize::new(2).unwrap(),
&rt,
)
Expand All @@ -446,11 +492,11 @@ mod tests {
// Try the aggregated strategy.
let parameters = [
search::Run::new(
graph::search::Knn::new(nearest_neighbors, 10, None).unwrap(),
KnnParams::new(nearest_neighbors, 10).unwrap(),
setup.clone(),
),
search::Run::new(
graph::search::Knn::new(nearest_neighbors, 15, None).unwrap(),
KnnParams::new(nearest_neighbors, 15).unwrap(),
setup.clone(),
),
];
Expand Down
2 changes: 1 addition & 1 deletion diskann-benchmark-core/src/search/graph/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ pub mod range;
pub mod strategy;

pub use inline::InlineFilterSearch;
pub use knn::KNN;
pub use knn::{KNN, KnnParams};
pub use multihop::MultiHop;
pub use range::Range;

Expand Down
15 changes: 8 additions & 7 deletions diskann-benchmark-core/src/search/graph/multihop.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use diskann::{
};
use diskann_utils::{future::AsyncFriendly, views::Matrix};

use crate::search::{self, Search, graph::Strategy};
use crate::search::{self, Search, graph::KnnParams, graph::Strategy};

/// A built-in helper for benchmarking filtered K-nearest neighbors search
/// using the multi-hop search method.
Expand All @@ -22,7 +22,7 @@ use crate::search::{self, Search, graph::Strategy};
/// [`search::search_all`] is provided by the [`search::graph::knn::Aggregator`] type (same
/// aggregator as [`search::graph::knn::KNN`]).
///
/// The provided implementation of [`Search`] accepts [`graph::search::Knn`]
/// The provided implementation of [`Search`] accepts [`KnnParams`]
/// and returns [`search::graph::knn::Metrics`] as additional output.
#[derive(Debug)]
pub struct MultiHop<DP, T, S>
Expand Down Expand Up @@ -97,7 +97,7 @@ where
T: AsyncFriendly + Clone,
{
type Id = DP::ExternalId;
type Parameters = graph::search::Knn;
type Parameters = KnnParams;
type Output = super::knn::Metrics;

fn num_queries(&self) -> usize {
Expand All @@ -118,7 +118,8 @@ where
O: graph::SearchOutputBuffer<DP::ExternalId> + Send,
{
let context = DP::Context::default();
let multihop_search = graph::search::MultihopFilterSearch::new(*parameters);
let knn = parameters.knn;
let multihop_search = graph::search::MultihopFilterSearch::new(knn);
let strategy =
labeled::Filtered::new(self.strategy.get(index)?.clone(), &*self.labels[index]);
let stats = self
Expand Down Expand Up @@ -191,7 +192,7 @@ mod tests {
let rt = crate::tokio::runtime(2).unwrap();
let results = search::search(
multihop.clone(),
graph::search::Knn::new(nearest_neighbors, 10, None).unwrap(),
KnnParams::new(nearest_neighbors, 10).unwrap(),
NonZeroUsize::new(2).unwrap(),
&rt,
)
Expand Down Expand Up @@ -219,11 +220,11 @@ mod tests {
// Try the aggregated strategy.
let parameters = [
search::Run::new(
graph::search::Knn::new(nearest_neighbors, 10, None).unwrap(),
KnnParams::new(nearest_neighbors, 10).unwrap(),
setup.clone(),
),
search::Run::new(
graph::search::Knn::new(nearest_neighbors, 15, None).unwrap(),
KnnParams::new(nearest_neighbors, 15).unwrap(),
setup.clone(),
),
];
Expand Down
2 changes: 1 addition & 1 deletion diskann-benchmark/src/index/result.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,7 @@ impl SearchResults {
Self {
num_tasks: setup.tasks.into(),
search_n: parameters.k_value().get(),
search_l: parameters.l_value().get(),
search_l: parameters.knn.l_value().get(),
qps,
search_latencies: end_to_end_latencies,
mean_latencies,
Expand Down
19 changes: 9 additions & 10 deletions diskann-benchmark/src/index/search/knn.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@

use std::{num::NonZeroUsize, sync::Arc};

use diskann_benchmark_core::recall::GroundTruthMode;
use diskann_benchmark_core::{self as benchmark_core, search as core_search};
use diskann_benchmark_core::{recall::GroundTruthMode, search::graph::KnnParams};

use crate::{index::result::SearchResults, inputs::graph_index::GraphSearch};

Expand Down Expand Up @@ -53,8 +53,7 @@ pub(crate) fn run<I>(
.search_l
.iter()
.map(|search_l| {
let search_params =
diskann::graph::search::Knn::new(run.search_n, *search_l, None).unwrap();
let search_params = KnnParams::new(run.search_n, *search_l).unwrap();

core_search::Run::new(search_params, setup.clone())
})
Expand All @@ -73,7 +72,7 @@ pub(crate) fn run<I>(
Ok(all)
}

type Run = core_search::Run<diskann::graph::search::Knn>;
type Run = core_search::Run<KnnParams>;
pub(crate) trait Knn<I> {
fn search_all(
&self,
Expand All @@ -94,13 +93,13 @@ where
DP: diskann::provider::DataProvider,
core_search::graph::KNN<DP, T, S, PP>: core_search::Search<
Id = DP::InternalId,
Parameters = diskann::graph::search::Knn,
Parameters = KnnParams,
Output = core_search::graph::knn::Metrics,
>,
{
fn search_all(
&self,
parameters: Vec<core_search::Run<diskann::graph::search::Knn>>,
parameters: Vec<core_search::Run<KnnParams>>,
groundtruth: &dyn benchmark_core::recall::Rows<DP::InternalId>,
recall_k: usize,
recall_n: usize,
Expand All @@ -126,13 +125,13 @@ where
DP: diskann::provider::DataProvider,
core_search::graph::MultiHop<DP, T, S>: core_search::Search<
Id = DP::InternalId,
Parameters = diskann::graph::search::Knn,
Parameters = KnnParams,
Output = core_search::graph::knn::Metrics,
>,
{
fn search_all(
&self,
parameters: Vec<core_search::Run<diskann::graph::search::Knn>>,
parameters: Vec<core_search::Run<KnnParams>>,
groundtruth: &dyn benchmark_core::recall::Rows<DP::InternalId>,
recall_k: usize,
recall_n: usize,
Expand All @@ -158,13 +157,13 @@ where
DP: diskann::provider::DataProvider,
core_search::graph::InlineFilterSearch<DP, T, S>: core_search::Search<
Id = DP::InternalId,
Parameters = diskann::graph::search::Knn,
Parameters = KnnParams,
Output = core_search::graph::knn::Metrics,
>,
{
fn search_all(
&self,
parameters: Vec<core_search::Run<diskann::graph::search::Knn>>,
parameters: Vec<core_search::Run<KnnParams>>,
groundtruth: &dyn benchmark_core::recall::Rows<DP::InternalId>,
recall_k: usize,
recall_n: usize,
Expand Down
Loading
Loading