Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
1 change: 0 additions & 1 deletion Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 2 additions & 2 deletions diskann-benchmark/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,7 @@ mod tests {

use diskann_benchmark_runner::{app::Commands, output::Memory};
use diskann_providers::storage::FileStorageProvider;
use diskann_tools::utils::{compute_ground_truth_from_datafiles, GraphDataF32Vector};
use diskann_tools::utils::compute_ground_truth_from_datafiles;
use diskann_vector::distance::Metric;

// Add these structs to deserialize the benchmark results
Expand Down Expand Up @@ -733,7 +733,7 @@ mod tests {

let disk_index_search_path = root_directory().join("test_data/disk_index_search");

let result = compute_ground_truth_from_datafiles::<GraphDataF32Vector, FileStorageProvider>(
let result = compute_ground_truth_from_datafiles::<f32, (), FileStorageProvider>(
&storage_provider,
Metric::L2, // distance function
disk_index_search_path
Expand Down
15 changes: 0 additions & 15 deletions diskann-disk/src/build/configuration/filter_parameter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,6 @@

use crate::data_model::GraphDataType;

pub type AssociatedDataFilter<Data> =
Box<dyn Fn(&<Data as GraphDataType>::AssociatedDataType) -> bool>;

pub fn default_associated_data_filter<Data: GraphDataType>() -> AssociatedDataFilter<Data> {
Box::new(|_| true)
}

pub type VectorFilter<'a, Data> =
Box<dyn Fn(&<Data as GraphDataType>::VectorIdType) -> bool + Send + Sync + 'a>;

Expand All @@ -26,14 +19,6 @@ mod tests {

type TestGraphData = GraphDataF32VectorUnitData;

#[test]
fn test_default_associated_data_filter_returns_true_for_all() {
let filter = default_associated_data_filter::<TestGraphData>();
// Test that the default filter always returns true
assert!(filter(&()));
assert!(filter(&()));
}

#[test]
fn test_default_vector_filter_returns_true_for_all() {
let filter = default_vector_filter::<TestGraphData>();
Expand Down
3 changes: 1 addition & 2 deletions diskann-tools/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@ byteorder.workspace = true
clap = { workspace = true, features = ["derive"] }
diskann-providers = { workspace = true, default-features = false } # see `linalg/Cargo.toml`
diskann-vector = { workspace = true }
diskann-disk = { workspace = true }
diskann-utils = { workspace = true }
bytemuck.workspace = true
num_cpus.workspace = true
Expand Down Expand Up @@ -47,7 +46,7 @@ diskann-utils = { workspace = true, features = ["testing"] }
[features]
default = []
# When a feature is enabled in tools, enable it in dependencies
perf_test = ["diskann-providers/perf_test", "diskann-disk/perf_test"]
perf_test = ["diskann-providers/perf_test"]

[lints]
clippy.undocumented_unsafe_blocks = "warn"
Expand Down
124 changes: 58 additions & 66 deletions diskann-tools/src/bin/compute_groundtruth.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,10 @@
use clap::Parser;
use diskann_providers::{storage::FileStorageProvider, utils::Timer};
use diskann_tools::utils::{
compute_ground_truth_from_datafiles, init_subscriber, CMDResult, DataType, GraphDataF32Vector,
GraphDataHalfVector, GraphDataInt8Vector, GraphDataU8Vector,
compute_ground_truth_from_datafiles, init_subscriber, CMDResult, DataType,
};
use diskann_vector::distance::Metric;
use diskann_vector::Half;

fn main() -> CMDResult<()> {
init_subscriber();
Expand All @@ -24,70 +24,62 @@ fn main() -> CMDResult<()> {
let storage_provider = FileStorageProvider;

let err = match args.data_type {
DataType::Float => {
compute_ground_truth_from_datafiles::<GraphDataF32Vector, FileStorageProvider>(
&storage_provider,
args.distance_function,
&args.base_file,
&args.query_file,
&args.ground_truth_file,
args.filter_bitmap_file.as_deref(),
args.recall_at,
insert_file,
skip_base,
args.associated_data_file,
args.base_file_labels.as_deref(),
args.query_file_labels.as_deref(),
)
}
DataType::Fp16 => {
compute_ground_truth_from_datafiles::<GraphDataHalfVector, FileStorageProvider>(
&storage_provider,
args.distance_function,
&args.base_file,
&args.query_file,
&args.ground_truth_file,
args.filter_bitmap_file.as_deref(),
args.recall_at,
insert_file,
skip_base,
args.associated_data_file,
args.base_file_labels.as_deref(),
args.query_file_labels.as_deref(),
)
}
DataType::Uint8 => {
compute_ground_truth_from_datafiles::<GraphDataU8Vector, FileStorageProvider>(
&storage_provider,
args.distance_function,
&args.base_file,
&args.query_file,
&args.ground_truth_file,
args.filter_bitmap_file.as_deref(),
args.recall_at,
insert_file,
skip_base,
args.associated_data_file,
args.base_file_labels.as_deref(),
args.query_file_labels.as_deref(),
)
}
DataType::Int8 => {
compute_ground_truth_from_datafiles::<GraphDataInt8Vector, FileStorageProvider>(
&storage_provider,
args.distance_function,
&args.base_file,
&args.query_file,
&args.ground_truth_file,
args.filter_bitmap_file.as_deref(),
args.recall_at,
insert_file,
skip_base,
args.associated_data_file,
args.base_file_labels.as_deref(),
args.query_file_labels.as_deref(),
)
}
DataType::Float => compute_ground_truth_from_datafiles::<f32, (), FileStorageProvider>(
&storage_provider,
args.distance_function,
&args.base_file,
&args.query_file,
&args.ground_truth_file,
args.filter_bitmap_file.as_deref(),
args.recall_at,
insert_file,
skip_base,
args.associated_data_file,
args.base_file_labels.as_deref(),
args.query_file_labels.as_deref(),
),
DataType::Fp16 => compute_ground_truth_from_datafiles::<Half, (), FileStorageProvider>(
&storage_provider,
args.distance_function,
&args.base_file,
&args.query_file,
&args.ground_truth_file,
args.filter_bitmap_file.as_deref(),
args.recall_at,
insert_file,
skip_base,
args.associated_data_file,
args.base_file_labels.as_deref(),
args.query_file_labels.as_deref(),
),
DataType::Uint8 => compute_ground_truth_from_datafiles::<u8, (), FileStorageProvider>(
&storage_provider,
args.distance_function,
&args.base_file,
&args.query_file,
&args.ground_truth_file,
args.filter_bitmap_file.as_deref(),
args.recall_at,
insert_file,
skip_base,
args.associated_data_file,
args.base_file_labels.as_deref(),
args.query_file_labels.as_deref(),
),
DataType::Int8 => compute_ground_truth_from_datafiles::<i8, (), FileStorageProvider>(
&storage_provider,
args.distance_function,
&args.base_file,
&args.query_file,
&args.ground_truth_file,
args.filter_bitmap_file.as_deref(),
args.recall_at,
insert_file,
skip_base,
args.associated_data_file,
args.base_file_labels.as_deref(),
args.query_file_labels.as_deref(),
),
};

match err {
Expand Down
54 changes: 25 additions & 29 deletions diskann-tools/src/bin/compute_multivec_groundtruth.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,10 @@ use clap::Parser;
use diskann_providers::{storage::FileStorageProvider, utils::Timer};
use diskann_tools::utils::{
compute_multivec_ground_truth_from_datafiles, init_subscriber, CMDResult, DataType,
GraphDataF32Vector, GraphDataHalfVector, GraphDataInt8Vector, GraphDataU8Vector,
MultivecAggregationMethod,
};
use diskann_vector::distance::Metric;
use diskann_vector::Half;

fn main() -> CMDResult<()> {
init_subscriber();
Expand All @@ -23,7 +23,7 @@ fn main() -> CMDResult<()> {

let err = match args.data_type {
DataType::Float => {
compute_multivec_ground_truth_from_datafiles::<GraphDataF32Vector, FileStorageProvider>(
compute_multivec_ground_truth_from_datafiles::<f32, FileStorageProvider>(
&storage_provider,
args.distance_function,
args.aggregation,
Expand All @@ -36,33 +36,7 @@ fn main() -> CMDResult<()> {
)
}
DataType::Fp16 => {
compute_multivec_ground_truth_from_datafiles::<GraphDataHalfVector, FileStorageProvider>(
&storage_provider,
args.distance_function,
args.aggregation,
&args.base_file,
&args.query_file,
&args.ground_truth_file,
args.recall_at,
args.base_file_labels.as_deref(),
args.query_file_labels.as_deref(),
)
}
DataType::Uint8 => {
compute_multivec_ground_truth_from_datafiles::<GraphDataU8Vector, FileStorageProvider>(
&storage_provider,
args.distance_function,
args.aggregation,
&args.base_file,
&args.query_file,
&args.ground_truth_file,
args.recall_at,
args.base_file_labels.as_deref(),
args.query_file_labels.as_deref(),
)
}
DataType::Int8 => {
compute_multivec_ground_truth_from_datafiles::<GraphDataInt8Vector, FileStorageProvider>(
compute_multivec_ground_truth_from_datafiles::<Half, FileStorageProvider>(
&storage_provider,
args.distance_function,
args.aggregation,
Expand All @@ -74,6 +48,28 @@ fn main() -> CMDResult<()> {
args.query_file_labels.as_deref(),
)
}
DataType::Uint8 => compute_multivec_ground_truth_from_datafiles::<u8, FileStorageProvider>(
&storage_provider,
args.distance_function,
args.aggregation,
&args.base_file,
&args.query_file,
&args.ground_truth_file,
args.recall_at,
args.base_file_labels.as_deref(),
args.query_file_labels.as_deref(),
),
DataType::Int8 => compute_multivec_ground_truth_from_datafiles::<i8, FileStorageProvider>(
&storage_provider,
args.distance_function,
args.aggregation,
&args.base_file,
&args.query_file,
&args.ground_truth_file,
args.recall_at,
args.base_file_labels.as_deref(),
args.query_file_labels.as_deref(),
),
};

match err {
Expand Down
32 changes: 15 additions & 17 deletions diskann-tools/src/bin/compute_range_groundtruth.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,9 @@ use clap::Parser;
use diskann_providers::{storage::FileStorageProvider, utils::Timer};
use diskann_tools::utils::{
compute_range_ground_truth_from_datafiles, init_subscriber, CMDResult, DataType,
GraphDataF32Vector, GraphDataHalfVector, GraphDataInt8Vector, GraphDataU8Vector,
};
use diskann_vector::distance::Metric;
use diskann_vector::Half;

fn main() -> CMDResult<()> {
init_subscriber();
Expand All @@ -22,7 +22,7 @@ fn main() -> CMDResult<()> {

let err = match args.data_type {
DataType::Float => {
compute_range_ground_truth_from_datafiles::<GraphDataF32Vector, FileStorageProvider>(
compute_range_ground_truth_from_datafiles::<f32, (), FileStorageProvider>(
&storage_provider,
args.distance_function,
&args.base_file,
Expand All @@ -35,7 +35,7 @@ fn main() -> CMDResult<()> {
)
}
DataType::Fp16 => {
compute_range_ground_truth_from_datafiles::<GraphDataHalfVector, FileStorageProvider>(
compute_range_ground_truth_from_datafiles::<Half, (), FileStorageProvider>(
&storage_provider,
args.distance_function,
&args.base_file,
Expand All @@ -48,20 +48,7 @@ fn main() -> CMDResult<()> {
)
}
DataType::Uint8 => {
compute_range_ground_truth_from_datafiles::<GraphDataU8Vector, FileStorageProvider>(
&storage_provider,
args.distance_function,
&args.base_file,
&args.query_file,
&args.ground_truth_file,
args.radius,
args.filter_bitmap_file.as_deref(),
args.base_file_labels.as_deref(),
args.query_file_labels.as_deref(),
)
}
DataType::Int8 => {
compute_range_ground_truth_from_datafiles::<GraphDataInt8Vector, FileStorageProvider>(
compute_range_ground_truth_from_datafiles::<u8, (), FileStorageProvider>(
&storage_provider,
args.distance_function,
&args.base_file,
Expand All @@ -73,6 +60,17 @@ fn main() -> CMDResult<()> {
args.query_file_labels.as_deref(),
)
}
DataType::Int8 => compute_range_ground_truth_from_datafiles::<i8, (), FileStorageProvider>(
&storage_provider,
args.distance_function,
&args.base_file,
&args.query_file,
&args.ground_truth_file,
args.radius,
args.filter_bitmap_file.as_deref(),
args.base_file_labels.as_deref(),
args.query_file_labels.as_deref(),
),
};

match err {
Expand Down
10 changes: 5 additions & 5 deletions diskann-tools/src/bin/generate_pq.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,9 @@ use clap::Parser;
use diskann_providers::storage::FileStorageProvider;
use diskann_tools::utils::{
build_pq, get_num_threads, init_subscriber, BuildPQParameters, CMDToolError, DataType,
GraphDataF32Vector, GraphDataHalfVector, GraphDataInt8Vector, GraphDataU8Vector,
};
use diskann_vector::distance::Metric;
use diskann_vector::Half;
use tracing::{error, info};

fn main() -> Result<(), CMDToolError> {
Expand All @@ -31,10 +31,10 @@ fn main() -> Result<(), CMDToolError> {
};

let err = match args.data_type {
DataType::Int8 => build_pq::<GraphDataInt8Vector>(&storage_provider, parameters),
DataType::Uint8 => build_pq::<GraphDataU8Vector>(&storage_provider, parameters),
DataType::Float => build_pq::<GraphDataF32Vector>(&storage_provider, parameters),
DataType::Fp16 => build_pq::<GraphDataHalfVector>(&storage_provider, parameters),
DataType::Int8 => build_pq::<i8>(&storage_provider, parameters),
DataType::Uint8 => build_pq::<u8>(&storage_provider, parameters),
DataType::Float => build_pq::<f32>(&storage_provider, parameters),
DataType::Fp16 => build_pq::<Half>(&storage_provider, parameters),
};

match err {
Expand Down
Loading
Loading