Skip to content

Commit 1d723fe

Browse files
usamoiusamoi
authored
refactor: make SQ build & preprocess faster (#596)
* chore: rename inverted_index to sparse_inverted_index Signed-off-by: usamoi <[email protected]> * refactor: make scalar quantization faster Signed-off-by: usamoi <[email protected]> --------- Signed-off-by: usamoi <[email protected]>
1 parent 22904a6 commit 1d723fe

File tree

25 files changed

+293
-218
lines changed

25 files changed

+293
-218
lines changed

‎Cargo.lock

Lines changed: 11 additions & 11 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/base/src/index.rs

Lines changed: 31 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -104,51 +104,49 @@ pub struct IndexOptions {
104104
}
105105

106106
impl IndexOptions {
107-
fn validate_self_quantization(
108-
&self,
109-
quantization: &Option<QuantizationOptions>,
110-
) -> Result<(), ValidationError> {
111-
match quantization {
112-
None => Ok(()),
113-
Some(
114-
QuantizationOptions::Scalar(_)
115-
| QuantizationOptions::Product(_)
116-
| QuantizationOptions::Rabitq(_),
117-
) => {
118-
if !matches!(self.vector.v, VectorKind::Vecf32 | VectorKind::Vecf16) {
119-
return Err(ValidationError::new(
120-
"quantization is not support for vectors that are not dense vectors",
121-
));
122-
}
123-
Ok(())
124-
}
125-
}
126-
}
127107
fn validate_self(&self) -> Result<(), ValidationError> {
128108
match &self.indexing {
129109
IndexingOptions::Flat(FlatIndexingOptions { quantization }) => {
130-
self.validate_self_quantization(quantization)?;
110+
if quantization.is_some()
111+
&& !matches!(self.vector.v, VectorKind::Vecf32 | VectorKind::Vecf16)
112+
{
113+
return Err(ValidationError::new(
114+
"quantization is only supported for dense vectors",
115+
));
116+
}
131117
}
132118
IndexingOptions::Ivf(IvfIndexingOptions { quantization, .. }) => {
133119
if !matches!(self.vector.v, VectorKind::Vecf32 | VectorKind::Vecf16) {
134120
return Err(ValidationError::new(
135-
"ivf is not support for vectors that are not dense vectors",
121+
"ivf is only supported for dense vectors",
122+
));
123+
}
124+
if quantization.is_some()
125+
&& !matches!(self.vector.v, VectorKind::Vecf32 | VectorKind::Vecf16)
126+
{
127+
return Err(ValidationError::new(
128+
"quantization is only supported for dense vectors",
136129
));
137130
}
138-
self.validate_self_quantization(quantization)?;
139131
}
140132
IndexingOptions::Hnsw(HnswIndexingOptions { quantization, .. }) => {
141-
self.validate_self_quantization(quantization)?;
142-
}
143-
IndexingOptions::InvertedIndex(_) => {
144-
if !matches!(self.vector.d, DistanceKind::Dot) {
133+
if quantization.is_some()
134+
&& !matches!(self.vector.v, VectorKind::Vecf32 | VectorKind::Vecf16)
135+
{
145136
return Err(ValidationError::new(
146-
"inverted_index is not support for distance that is not negative dot product",
137+
"quantization is only supported for dense vectors",
147138
));
148139
}
140+
}
141+
IndexingOptions::SparseInvertedIndex(_) => {
149142
if !matches!(self.vector.v, VectorKind::SVecf32) {
150143
return Err(ValidationError::new(
151-
"inverted_index is not support for vectors that are not sparse vectors",
144+
"sparse_inverted_index is only supported for sparse vectors",
145+
));
146+
}
147+
if !matches!(self.vector.d, DistanceKind::Dot) {
148+
return Err(ValidationError::new(
149+
"sparse_inverted_index is only supported for dot distance",
152150
));
153151
}
154152
}
@@ -284,7 +282,7 @@ pub enum IndexingOptions {
284282
Flat(FlatIndexingOptions),
285283
Ivf(IvfIndexingOptions),
286284
Hnsw(HnswIndexingOptions),
287-
InvertedIndex(InvertedIndexingOptions),
285+
SparseInvertedIndex(SparseInvertedIndexIndexingOptions),
288286
}
289287

290288
impl IndexingOptions {
@@ -320,16 +318,16 @@ impl Validate for IndexingOptions {
320318
Self::Flat(x) => x.validate(),
321319
Self::Ivf(x) => x.validate(),
322320
Self::Hnsw(x) => x.validate(),
323-
Self::InvertedIndex(x) => x.validate(),
321+
Self::SparseInvertedIndex(x) => x.validate(),
324322
}
325323
}
326324
}
327325

328326
#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
329327
#[serde(deny_unknown_fields)]
330-
pub struct InvertedIndexingOptions {}
328+
pub struct SparseInvertedIndexIndexingOptions {}
331329

332-
impl Default for InvertedIndexingOptions {
330+
impl Default for SparseInvertedIndexIndexingOptions {
333331
fn default() -> Self {
334332
Self {}
335333
}

‎crates/base/src/operator/mod.rs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,5 +23,4 @@ pub trait Operator: Copy + 'static + Send + Sync {
2323
fn distance(lhs: Borrowed<'_, Self>, rhs: Borrowed<'_, Self>) -> Distance;
2424
}
2525

26-
pub type Owned<T> = <T as Operator>::Vector;
2726
pub type Borrowed<'a, T> = <<T as Operator>::Vector as VectorOwned>::Borrowed<'a>;

‎crates/flat/src/lib.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ impl<O: OperatorFlat, Q: Quantizer<O>> Flat<O, Q> {
2929
pub fn create(
3030
path: impl AsRef<Path>,
3131
options: IndexOptions,
32-
source: &(impl Vectors<Owned<O>> + Collection + Source + Sync),
32+
source: &(impl Vectors<O::Vector> + Collection + Source + Sync),
3333
) -> Self {
3434
let remapped = RemappedCollection::from_source(source);
3535
from_nothing(path, options, &remapped)
@@ -83,7 +83,7 @@ impl<O: OperatorFlat, Q: Quantizer<O>> Flat<O, Q> {
8383
fn from_nothing<O: OperatorFlat, Q: Quantizer<O>>(
8484
path: impl AsRef<Path>,
8585
options: IndexOptions,
86-
collection: &(impl Vectors<Owned<O>> + Collection + Sync),
86+
collection: &(impl Vectors<O::Vector> + Collection + Sync),
8787
) -> Flat<O, Q> {
8888
create_dir(path.as_ref()).unwrap();
8989
let flat_indexing_options = options.indexing.clone().unwrap_flat();

‎crates/hnsw/src/lib.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ impl<O: OperatorHnsw, Q: Quantizer<O>> Hnsw<O, Q> {
4646
pub fn create(
4747
path: impl AsRef<Path>,
4848
options: IndexOptions,
49-
source: &(impl Vectors<Owned<O>> + Collection + Source + Sync),
49+
source: &(impl Vectors<O::Vector> + Collection + Source + Sync),
5050
) -> Self {
5151
let remapped = RemappedCollection::from_source(source);
5252
if let Some(main) = source.get_main::<Self>() {
@@ -116,7 +116,7 @@ impl<O: OperatorHnsw, Q: Quantizer<O>> Hnsw<O, Q> {
116116
fn from_nothing<O: OperatorHnsw, Q: Quantizer<O>>(
117117
path: impl AsRef<Path>,
118118
options: IndexOptions,
119-
collection: &(impl Vectors<Owned<O>> + Collection + Sync),
119+
collection: &(impl Vectors<O::Vector> + Collection + Sync),
120120
) -> Hnsw<O, Q> {
121121
create_dir(path.as_ref()).unwrap();
122122
let HnswIndexingOptions {
@@ -198,7 +198,7 @@ fn from_nothing<O: OperatorHnsw, Q: Quantizer<O>>(
198198
fn from_main<O: OperatorHnsw, Q: Quantizer<O>>(
199199
path: impl AsRef<Path>,
200200
options: IndexOptions,
201-
remapped: &RemappedCollection<Owned<O>, impl Vectors<Owned<O>> + Collection + Sync>,
201+
remapped: &RemappedCollection<O::Vector, impl Vectors<O::Vector> + Collection + Sync>,
202202
main: &Hnsw<O, Q>,
203203
) -> Hnsw<O, Q> {
204204
create_dir(path.as_ref()).unwrap();

‎crates/index/src/lib.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -316,7 +316,7 @@ impl<O: Op> Index<O> {
316316
}
317317
pub fn create_sealed_segment(
318318
&self,
319-
source: &(impl Vectors<Owned<O>> + Collection + Source + Sync),
319+
source: &(impl Vectors<O::Vector> + Collection + Source + Sync),
320320
sealed_segment_ids: &[NonZeroU128],
321321
growing_segment_ids: &[NonZeroU128],
322322
) -> Option<Arc<SealedSegment<O>>> {
@@ -444,7 +444,7 @@ impl<O: Op> IndexView<O> {
444444
}
445445
pub fn insert(
446446
&self,
447-
vector: Owned<O>,
447+
vector: O::Vector,
448448
pointer: Pointer,
449449
) -> Result<Result<(), OutdatedError>, InsertError> {
450450
if self.options.vector.dims != vector.as_borrowed().dims() {

‎crates/index/src/optimizing/index_source.rs

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use crate::delete::Delete;
22
use crate::Op;
33
use crate::{GrowingSegment, SealedSegment};
44
use base::index::IndexOptions;
5-
use base::operator::{Borrowed, Owned};
5+
use base::operator::Borrowed;
66
use base::search::*;
77
use std::any::Any;
88
use std::fmt::Debug;
@@ -17,7 +17,7 @@ pub struct IndexSource<V, O: Op> {
1717
_phantom: PhantomData<fn(V) -> V>,
1818
}
1919

20-
impl<O: Op> IndexSource<Owned<O>, O> {
20+
impl<O: Op> IndexSource<O::Vector, O> {
2121
pub fn new(
2222
options: IndexOptions,
2323
sealed: Option<Arc<SealedSegment<O>>>,
@@ -34,7 +34,7 @@ impl<O: Op> IndexSource<Owned<O>, O> {
3434
}
3535
}
3636

37-
impl<O: Op> Vectors<Owned<O>> for IndexSource<Owned<O>, O> {
37+
impl<O: Op> Vectors<O::Vector> for IndexSource<O::Vector, O> {
3838
fn dims(&self) -> u32 {
3939
self.dims
4040
}
@@ -61,7 +61,7 @@ impl<O: Op> Vectors<Owned<O>> for IndexSource<Owned<O>, O> {
6161
}
6262
}
6363

64-
impl<O: Op> Collection for IndexSource<Owned<O>, O> {
64+
impl<O: Op> Collection for IndexSource<O::Vector, O> {
6565
fn payload(&self, mut index: u32) -> Payload {
6666
for x in self.sealed.iter() {
6767
if index < x.len() {
@@ -79,7 +79,7 @@ impl<O: Op> Collection for IndexSource<Owned<O>, O> {
7979
}
8080
}
8181

82-
impl<O: Op> Source for IndexSource<Owned<O>, O> {
82+
impl<O: Op> Source for IndexSource<O::Vector, O> {
8383
fn get_main<T: Any>(&self) -> Option<&T> {
8484
let x = self.sealed.as_ref()?;
8585
Some(
@@ -104,7 +104,7 @@ pub struct RoGrowingCollection<V, O: Op> {
104104
_phantom: PhantomData<fn(V) -> V>,
105105
}
106106

107-
impl<O: Op> Debug for RoGrowingCollection<Owned<O>, O> {
107+
impl<O: Op> Debug for RoGrowingCollection<O::Vector, O> {
108108
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109109
f.debug_struct("RoGrowingCollection")
110110
.field("growing", &self.growing)
@@ -113,7 +113,7 @@ impl<O: Op> Debug for RoGrowingCollection<Owned<O>, O> {
113113
}
114114
}
115115

116-
impl<O: Op> Vectors<Owned<O>> for RoGrowingCollection<Owned<O>, O> {
116+
impl<O: Op> Vectors<O::Vector> for RoGrowingCollection<O::Vector, O> {
117117
fn dims(&self) -> u32 {
118118
self.dims
119119
}
@@ -133,7 +133,7 @@ impl<O: Op> Vectors<Owned<O>> for RoGrowingCollection<Owned<O>, O> {
133133
}
134134
}
135135

136-
impl<O: Op> Collection for RoGrowingCollection<Owned<O>, O> {
136+
impl<O: Op> Collection for RoGrowingCollection<O::Vector, O> {
137137
fn payload(&self, mut index: u32) -> Payload {
138138
for x in self.growing.iter() {
139139
if index < x.len() {

‎crates/index/src/optimizing/indexing.rs

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,13 @@
11
use crate::optimizing::index_source::IndexSource;
22
use crate::Index;
33
use crate::Op;
4-
use base::operator::Owned;
54
use std::sync::Arc;
65

76
pub fn scan<O: Op>(
87
index: Arc<Index<O>>,
98
capacity: u32,
109
delete_threshold: f64,
11-
) -> Option<IndexSource<Owned<O>, O>> {
10+
) -> Option<IndexSource<O::Vector, O>> {
1211
let (sealed, growing) = 'a: {
1312
let protect = index.protect.lock();
1413
// approach 1: merge small segments to a big segment
@@ -87,7 +86,7 @@ pub fn scan<O: Op>(
8786
))
8887
}
8988

90-
pub fn make<O: Op>(index: Arc<Index<O>>, source: IndexSource<Owned<O>, O>) {
89+
pub fn make<O: Op>(index: Arc<Index<O>>, source: IndexSource<O::Vector, O>) {
9190
let _ = index.create_sealed_segment(
9291
&source,
9392
&source.sealed.iter().map(|x| x.id()).collect::<Vec<_>>(),

‎crates/index/src/segment/sealed.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ impl<O: Op> SealedSegment<O> {
3737
path: PathBuf,
3838
id: NonZeroU128,
3939
options: IndexOptions,
40-
source: &(impl Vectors<Owned<O>> + Collection + Source + Sync),
40+
source: &(impl Vectors<O::Vector> + Collection + Source + Sync),
4141
) -> Arc<Self> {
4242
let indexing = SealedIndexing::create(&path, options, source);
4343
Arc::new(Self {

‎crates/indexing/Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,9 @@ base = { path = "../base" }
1111
# algorithms
1212
flat = { path = "../flat" }
1313
hnsw = { path = "../hnsw" }
14-
inverted = { path = "../inverted" }
1514
ivf = { path = "../ivf" }
1615
quantization = { path = "../quantization" }
16+
sparse_inverted_index = { path = "../sparse_inverted_index" }
1717

1818
[lints]
1919
workspace = true

‎crates/indexing/src/lib.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,16 @@ use quantization::rabitq::OperatorRabitqQuantization;
44
pub use sealed::SealedIndexing;
55

66
use base::operator::Operator;
7-
use inverted::operator::OperatorInvertedIndex;
87
use ivf::operator::OperatorIvf;
98
use quantization::product::OperatorProductQuantization;
109
use quantization::scalar::OperatorScalarQuantization;
10+
use sparse_inverted_index::operator::OperatorSparseInvertedIndex;
1111

1212
pub trait OperatorIndexing
1313
where
1414
Self: Operator,
1515
Self: OperatorIvf,
16-
Self: OperatorInvertedIndex,
16+
Self: OperatorSparseInvertedIndex,
1717
Self: OperatorScalarQuantization,
1818
Self: OperatorProductQuantization,
1919
Self: OperatorRabitqQuantization,
@@ -24,7 +24,7 @@ impl<T> OperatorIndexing for T
2424
where
2525
Self: Operator,
2626
Self: OperatorIvf,
27-
Self: OperatorInvertedIndex,
27+
Self: OperatorSparseInvertedIndex,
2828
Self: OperatorScalarQuantization,
2929
Self: OperatorProductQuantization,
3030
Self: OperatorRabitqQuantization,

0 commit comments

Comments
 (0)