diskann_providers/storage/
api.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use std::{future::Future, num::NonZeroUsize};
7
8use super::{StorageReadProvider, StorageWriteProvider};
9use diskann_vector::distance::Metric;
10
11use super::get_mem_index_data_file;
12use crate::model::graph::provider::async_::PrefetchCacheLineLevel;
13
14/// A general trait for saving `self` to disk.
15///
16/// The generic parameter `T` is used to allow types to require arbitrary associated
17/// metadata in order to successfully save themselves.
18///
19/// Additionally, types may overload the auxiliary state to customize semantics.
20///
21/// See also: [`LoadWith`].
22pub trait SaveWith<T> {
23    /// The return type upon successful saving. This is often (but is not required to be)
24    /// the number of bytes written to disk.
25    type Ok: Send;
26
27    /// The error type if serialization is unsuccessful.
28    type Error: std::error::Error + Send;
29
30    /// Safe `self` to disk using `provider` for IO-related needs. Argument `auxiliary`
31    /// can be arbitrary metadata required for a successful operation, such as file paths.
32    fn save_with<P>(
33        &self,
34        provider: &P,
35        auxiliary: &T,
36    ) -> impl Future<Output = Result<Self::Ok, Self::Error>> + Send
37    where
38        P: StorageWriteProvider;
39}
40
41/// A general trait for loading `self` from disk.
42///
43/// The generic parameter `T` is used to allow types to require arbitrary associated
44/// metadata in order to successfully load themselves.
45///
46/// Additionally, types may overload the auxiliary state to customize semantics.
47///
48/// See also: [`SaveWith`].
49pub trait LoadWith<T>: Sized {
50    /// The error type if deserialization is unsuccessful.
51    type Error: std::error::Error + Send;
52
53    /// Load `self` form disk using `provider` for IO-related needs. Argument `auxiliary`
54    /// can be arbitrary metadata required for a successful operation, such as file paths.
55    fn load_with<P>(
56        provider: &P,
57        auxiliary: &T,
58    ) -> impl Future<Output = Result<Self, Self::Error>> + Send
59    where
60        P: StorageReadProvider;
61}
62
63/// The file-path prefix for saving and loading an async index.
64///
65/// An auxiliary type for [`SaveWith`] and [`LoadWith`] to indicate that the object being
66/// saved or loaded is part of an async in-memory index.
67///
68/// This mainly controls how file-paths are generated.
69///
70/// For example, graph data is located at the raw-prefix, while the full-precision data
71/// is saved using the `.data` suffix.
72#[derive(Debug, Clone, PartialEq, Eq)]
73pub struct AsyncIndexMetadata {
74    prefix: String,
75}
76
77impl AsyncIndexMetadata {
78    /// Construct a new `AsyncIndexPrefix` from `pathlike`.
79    pub fn new<T>(pathlike: T) -> Self
80    where
81        String: From<T>,
82    {
83        Self {
84            prefix: pathlike.into(),
85        }
86    }
87
88    /// Return the file path contained in `self` as a `&str`.
89    pub fn prefix(&self) -> &str {
90        &self.prefix
91    }
92
93    /// Obtain the file path for full-precision data using `self` as the file path prefix.
94    pub fn data_path(&self) -> String {
95        get_mem_index_data_file(&self.prefix)
96    }
97
98    /// Obtain the file path for additional points file
99    pub fn additional_points_id_path(&self) -> String {
100        format!("{}.additional_points_id", self.prefix)
101    }
102}
103
104/// The file-path prefix for saving only graph data during disk index construction.
105///
106/// An auxiliary type for [`SaveWith`] to specify graph-only serialization.
107#[derive(Debug, Clone)]
108pub struct DiskGraphOnly {
109    prefix: String,
110}
111
112impl DiskGraphOnly {
113    /// Constructs a new `DiskGraphOnly` from a path-like object.
114    pub fn new<T>(pathlike: T) -> Self
115    where
116        String: From<T>,
117    {
118        Self {
119            prefix: pathlike.into(),
120        }
121    }
122
123    /// Return the file path contained in `self` as a `&str`.
124    pub fn prefix(&self) -> &str {
125        &self.prefix
126    }
127}
128
129/// Indicates that the canonical layout for a Quant index is expected for deserialization.
130///
131/// For a file-path prefix `prefix`, this layout includes the following files:
132///
133/// **Common:**
134/// * `prefix`: The file that contains the saved graph.
135/// * `prefix.data`: The serialized full-precision data in `.bin` form.
136///
137/// Depending on the quantization method used, one of the following sets of files will also be present:
138///
139/// **If Product Quantization (PQ) is used:**
140/// * `prefix_build_pq_pivots.bin`: The saved PQ pivot table.
141/// * `prefix_build_pq_compressed.bin`: The saved PQ codes.
142///
143/// **If Scalar Quantization (SQ) is used:**
144/// * `prefix_sq_compressed.bin`: The saved scalar quantized codes.
145/// * `prefix_scalar_quantizer_proto.bin`: The saved scalar quantizer metadata.
146pub struct AsyncQuantLoadContext {
147    /// The file path prefix of the index.
148    pub metadata: AsyncIndexMetadata,
149    /// The number of frozen points stored in the datasets.
150    pub num_frozen_points: NonZeroUsize,
151    /// The metric to use for this index.
152    pub metric: Metric,
153    /// The number of iterations to prefetch when performing bulk-retrievals.
154    pub prefetch_lookahead: Option<usize>,
155    /// Temporary parameter to indicate if the index is a disk index to load right file names.
156    /// This can be removed once disk index uses same file names as async index.
157    pub is_disk_index: bool,
158    /// controls the prefetch cache line level for the index.
159    pub prefetch_cache_line_level: Option<PrefetchCacheLineLevel>,
160}
diskann_providers/storage/api.rs

diskann_providers/storage/
api.rs
