mirror of
https://github.com/meilisearch/meilisearch.git
synced 2025-10-21 11:06:27 +00:00
Switch from version to backend selector
This commit is contained in:
@@ -31,7 +31,7 @@ use crate::prompt::PromptData;
|
|||||||
use crate::proximity::ProximityPrecision;
|
use crate::proximity::ProximityPrecision;
|
||||||
use crate::update::new::StdResult;
|
use crate::update::new::StdResult;
|
||||||
use crate::vector::db::IndexEmbeddingConfigs;
|
use crate::vector::db::IndexEmbeddingConfigs;
|
||||||
use crate::vector::{Embedding, HannoyStats, VectorStore};
|
use crate::vector::{Embedding, HannoyStats, VectorStore, VectorStoreBackend};
|
||||||
use crate::{
|
use crate::{
|
||||||
default_criteria, CboRoaringBitmapCodec, Criterion, DocumentId, ExternalDocumentsIds,
|
default_criteria, CboRoaringBitmapCodec, Criterion, DocumentId, ExternalDocumentsIds,
|
||||||
FacetDistribution, FieldDistribution, FieldId, FieldIdMapMissingEntry, FieldIdWordCountCodec,
|
FacetDistribution, FieldDistribution, FieldId, FieldIdMapMissingEntry, FieldIdWordCountCodec,
|
||||||
@@ -87,6 +87,7 @@ pub mod main_key {
|
|||||||
pub const DOCUMENTS_STATS: &str = "documents_stats";
|
pub const DOCUMENTS_STATS: &str = "documents_stats";
|
||||||
pub const DISABLED_TYPOS_TERMS: &str = "disabled_typos_terms";
|
pub const DISABLED_TYPOS_TERMS: &str = "disabled_typos_terms";
|
||||||
pub const CHAT: &str = "chat";
|
pub const CHAT: &str = "chat";
|
||||||
|
pub const VECTOR_STORE_BACKEND: &str = "vector_store_backend";
|
||||||
}
|
}
|
||||||
|
|
||||||
pub mod db_name {
|
pub mod db_name {
|
||||||
@@ -454,6 +455,35 @@ impl Index {
|
|||||||
self.main.remap_types::<Str, VersionCodec>().get(rtxn, main_key::VERSION_KEY)
|
self.main.remap_types::<Str, VersionCodec>().get(rtxn, main_key::VERSION_KEY)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/* vector store */
|
||||||
|
/// Writes the vector store
|
||||||
|
pub(crate) fn put_vector_store(
|
||||||
|
&self,
|
||||||
|
wtxn: &mut RwTxn<'_>,
|
||||||
|
backend: VectorStoreBackend,
|
||||||
|
) -> Result<()> {
|
||||||
|
Ok(self.main.remap_types::<Str, SerdeJson<VectorStoreBackend>>().put(
|
||||||
|
wtxn,
|
||||||
|
main_key::VECTOR_STORE_BACKEND,
|
||||||
|
&backend,
|
||||||
|
)?)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn get_vector_store(&self, rtxn: &RoTxn<'_>) -> Result<VectorStoreBackend> {
|
||||||
|
Ok(self
|
||||||
|
.main
|
||||||
|
.remap_types::<Str, SerdeJson<VectorStoreBackend>>()
|
||||||
|
.get(rtxn, main_key::VECTOR_STORE_BACKEND)?
|
||||||
|
.unwrap_or_default())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn delete_vector_store(&self, wtxn: &mut RwTxn<'_>) -> Result<bool> {
|
||||||
|
Ok(self
|
||||||
|
.main
|
||||||
|
.remap_types::<Str, SerdeJson<VectorStoreBackend>>()
|
||||||
|
.delete(wtxn, main_key::VECTOR_STORE_BACKEND)?)
|
||||||
|
}
|
||||||
|
|
||||||
/* documents ids */
|
/* documents ids */
|
||||||
|
|
||||||
/// Writes the documents ids that corresponds to the user-ids-documents-ids FST.
|
/// Writes the documents ids that corresponds to the user-ids-documents-ids FST.
|
||||||
@@ -1769,12 +1799,13 @@ impl Index {
|
|||||||
) -> Result<BTreeMap<String, EmbeddingsWithMetadata>> {
|
) -> Result<BTreeMap<String, EmbeddingsWithMetadata>> {
|
||||||
let mut res = BTreeMap::new();
|
let mut res = BTreeMap::new();
|
||||||
let embedders = self.embedding_configs();
|
let embedders = self.embedding_configs();
|
||||||
let index_version = self.get_version(rtxn)?.unwrap();
|
let backend = self.get_vector_store(rtxn)?;
|
||||||
|
|
||||||
for config in embedders.embedding_configs(rtxn)? {
|
for config in embedders.embedding_configs(rtxn)? {
|
||||||
let embedder_info = embedders.embedder_info(rtxn, &config.name)?.unwrap();
|
let embedder_info = embedders.embedder_info(rtxn, &config.name)?.unwrap();
|
||||||
let has_fragments = config.config.embedder_options.has_fragments();
|
let has_fragments = config.config.embedder_options.has_fragments();
|
||||||
let reader = VectorStore::new(
|
let reader = VectorStore::new(
|
||||||
index_version,
|
backend,
|
||||||
self.vector_store,
|
self.vector_store,
|
||||||
embedder_info.embedder_id,
|
embedder_info.embedder_id,
|
||||||
config.config.quantized(),
|
config.config.quantized(),
|
||||||
@@ -1797,11 +1828,12 @@ impl Index {
|
|||||||
pub fn hannoy_stats(&self, rtxn: &RoTxn<'_>) -> Result<HannoyStats> {
|
pub fn hannoy_stats(&self, rtxn: &RoTxn<'_>) -> Result<HannoyStats> {
|
||||||
let mut stats = HannoyStats::default();
|
let mut stats = HannoyStats::default();
|
||||||
let embedding_configs = self.embedding_configs();
|
let embedding_configs = self.embedding_configs();
|
||||||
let index_version = self.get_version(rtxn)?.unwrap();
|
let backend = self.get_vector_store(rtxn)?;
|
||||||
|
|
||||||
for config in embedding_configs.embedding_configs(rtxn)? {
|
for config in embedding_configs.embedding_configs(rtxn)? {
|
||||||
let embedder_id = embedding_configs.embedder_id(rtxn, &config.name)?.unwrap();
|
let embedder_id = embedding_configs.embedder_id(rtxn, &config.name)?.unwrap();
|
||||||
let reader = VectorStore::new(
|
let reader = VectorStore::new(
|
||||||
index_version,
|
backend,
|
||||||
self.vector_store,
|
self.vector_store,
|
||||||
embedder_id,
|
embedder_id,
|
||||||
config.config.quantized(),
|
config.config.quantized(),
|
||||||
|
@@ -82,7 +82,7 @@ fn evaluate_inner(
|
|||||||
embedding_configs: &[IndexEmbeddingConfig],
|
embedding_configs: &[IndexEmbeddingConfig],
|
||||||
filter: &VectorFilter<'_>,
|
filter: &VectorFilter<'_>,
|
||||||
) -> crate::Result<RoaringBitmap> {
|
) -> crate::Result<RoaringBitmap> {
|
||||||
let index_version = index.get_version(rtxn)?.unwrap();
|
let backend = index.get_vector_store(rtxn)?;
|
||||||
let embedder_name = embedder.value();
|
let embedder_name = embedder.value();
|
||||||
let available_embedders =
|
let available_embedders =
|
||||||
|| embedding_configs.iter().map(|c| c.name.clone()).collect::<Vec<_>>();
|
|| embedding_configs.iter().map(|c| c.name.clone()).collect::<Vec<_>>();
|
||||||
@@ -98,7 +98,7 @@ fn evaluate_inner(
|
|||||||
.ok_or_else(|| EmbedderDoesNotExist { embedder, available: available_embedders() })?;
|
.ok_or_else(|| EmbedderDoesNotExist { embedder, available: available_embedders() })?;
|
||||||
|
|
||||||
let vector_store = VectorStore::new(
|
let vector_store = VectorStore::new(
|
||||||
index_version,
|
backend,
|
||||||
index.vector_store,
|
index.vector_store,
|
||||||
embedder_info.embedder_id,
|
embedder_info.embedder_id,
|
||||||
embedding_config.config.quantized(),
|
embedding_config.config.quantized(),
|
||||||
|
@@ -54,14 +54,11 @@ impl<Q: RankingRuleQueryTrait> VectorSort<Q> {
|
|||||||
vector_candidates: &RoaringBitmap,
|
vector_candidates: &RoaringBitmap,
|
||||||
) -> Result<()> {
|
) -> Result<()> {
|
||||||
let target = &self.target;
|
let target = &self.target;
|
||||||
|
let backend = ctx.index.get_vector_store(ctx.txn)?;
|
||||||
|
|
||||||
let before = Instant::now();
|
let before = Instant::now();
|
||||||
let reader = VectorStore::new(
|
let reader =
|
||||||
ctx.index.get_version(ctx.txn)?.unwrap(),
|
VectorStore::new(backend, ctx.index.vector_store, self.embedder_index, self.quantized);
|
||||||
ctx.index.vector_store,
|
|
||||||
self.embedder_index,
|
|
||||||
self.quantized,
|
|
||||||
);
|
|
||||||
let results = reader.nns_by_vector(ctx.txn, target, self.limit, Some(vector_candidates))?;
|
let results = reader.nns_by_vector(ctx.txn, target, self.limit, Some(vector_candidates))?;
|
||||||
self.cached_sorted_docids = results.into_iter();
|
self.cached_sorted_docids = results.into_iter();
|
||||||
*ctx.vector_store_stats.get_or_insert_default() += VectorStoreStats {
|
*ctx.vector_store_stats.get_or_insert_default() += VectorStoreStats {
|
||||||
|
@@ -72,12 +72,10 @@ impl<'a> Similar<'a> {
|
|||||||
crate::UserError::InvalidSimilarEmbedder(self.embedder_name.to_owned())
|
crate::UserError::InvalidSimilarEmbedder(self.embedder_name.to_owned())
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
let reader = VectorStore::new(
|
let backend = self.index.get_vector_store(self.rtxn)?;
|
||||||
self.index.get_version(self.rtxn)?.unwrap(),
|
|
||||||
self.index.vector_store,
|
let reader =
|
||||||
embedder_index,
|
VectorStore::new(backend, self.index.vector_store, embedder_index, self.quantized);
|
||||||
self.quantized,
|
|
||||||
);
|
|
||||||
let results = reader.nns_by_item(
|
let results = reader.nns_by_item(
|
||||||
self.rtxn,
|
self.rtxn,
|
||||||
self.id,
|
self.id,
|
||||||
|
@@ -485,7 +485,7 @@ where
|
|||||||
|
|
||||||
// If an embedder wasn't used in the typedchunk but must be binary quantized
|
// If an embedder wasn't used in the typedchunk but must be binary quantized
|
||||||
// we should insert it in `dimension`
|
// we should insert it in `dimension`
|
||||||
let index_version = self.index.get_version(self.wtxn)?.unwrap();
|
let backend = self.index.get_vector_store(self.wtxn)?;
|
||||||
for (name, action) in settings_diff.embedding_config_updates.iter() {
|
for (name, action) in settings_diff.embedding_config_updates.iter() {
|
||||||
if action.is_being_quantized && !dimension.contains_key(name.as_str()) {
|
if action.is_being_quantized && !dimension.contains_key(name.as_str()) {
|
||||||
let index = self.index.embedding_configs().embedder_id(self.wtxn, name)?.ok_or(
|
let index = self.index.embedding_configs().embedder_id(self.wtxn, name)?.ok_or(
|
||||||
@@ -494,12 +494,8 @@ where
|
|||||||
key: None,
|
key: None,
|
||||||
},
|
},
|
||||||
)?;
|
)?;
|
||||||
let reader = VectorStore::new(
|
let reader =
|
||||||
index_version,
|
VectorStore::new(backend, self.index.vector_store, index, action.was_quantized);
|
||||||
self.index.vector_store,
|
|
||||||
index,
|
|
||||||
action.was_quantized,
|
|
||||||
);
|
|
||||||
let Some(dim) = reader.dimensions(self.wtxn)? else {
|
let Some(dim) = reader.dimensions(self.wtxn)? else {
|
||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
@@ -529,7 +525,7 @@ where
|
|||||||
|
|
||||||
pool.install(|| {
|
pool.install(|| {
|
||||||
let mut writer =
|
let mut writer =
|
||||||
VectorStore::new(index_version, vector_hannoy, embedder_index, was_quantized);
|
VectorStore::new(backend, vector_hannoy, embedder_index, was_quantized);
|
||||||
writer.build_and_quantize(
|
writer.build_and_quantize(
|
||||||
wtxn,
|
wtxn,
|
||||||
// In the settings we don't have any progress to share
|
// In the settings we don't have any progress to share
|
||||||
|
@@ -834,7 +834,7 @@ impl<'a, 'i> Transform<'a, 'i> {
|
|||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
|
||||||
let index_version = self.index.get_version(wtxn)?.unwrap();
|
let backend = self.index.get_vector_store(wtxn)?;
|
||||||
let readers: BTreeMap<&str, (VectorStore, &RoaringBitmap)> = settings_diff
|
let readers: BTreeMap<&str, (VectorStore, &RoaringBitmap)> = settings_diff
|
||||||
.embedding_config_updates
|
.embedding_config_updates
|
||||||
.iter()
|
.iter()
|
||||||
@@ -843,7 +843,7 @@ impl<'a, 'i> Transform<'a, 'i> {
|
|||||||
action.write_back()
|
action.write_back()
|
||||||
{
|
{
|
||||||
let reader = VectorStore::new(
|
let reader = VectorStore::new(
|
||||||
index_version,
|
backend,
|
||||||
self.index.vector_store,
|
self.index.vector_store,
|
||||||
*embedder_id,
|
*embedder_id,
|
||||||
action.was_quantized,
|
action.was_quantized,
|
||||||
@@ -949,7 +949,7 @@ impl<'a, 'i> Transform<'a, 'i> {
|
|||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
let hannoy = VectorStore::new(
|
let hannoy = VectorStore::new(
|
||||||
index_version,
|
backend,
|
||||||
self.index.vector_store,
|
self.index.vector_store,
|
||||||
infos.embedder_id,
|
infos.embedder_id,
|
||||||
was_quantized,
|
was_quantized,
|
||||||
|
@@ -619,7 +619,7 @@ pub(crate) fn write_typed_chunk_into_index(
|
|||||||
let _entered = span.enter();
|
let _entered = span.enter();
|
||||||
|
|
||||||
let embedders = index.embedding_configs();
|
let embedders = index.embedding_configs();
|
||||||
let index_version = index.get_version(wtxn)?.unwrap();
|
let backend = index.get_vector_store(wtxn)?;
|
||||||
|
|
||||||
let mut remove_vectors_builder = MergerBuilder::new(KeepFirst);
|
let mut remove_vectors_builder = MergerBuilder::new(KeepFirst);
|
||||||
let mut manual_vectors_builder = MergerBuilder::new(KeepFirst);
|
let mut manual_vectors_builder = MergerBuilder::new(KeepFirst);
|
||||||
@@ -678,12 +678,8 @@ pub(crate) fn write_typed_chunk_into_index(
|
|||||||
.get(&embedder_name)
|
.get(&embedder_name)
|
||||||
.is_some_and(|conf| conf.is_quantized);
|
.is_some_and(|conf| conf.is_quantized);
|
||||||
// FIXME: allow customizing distance
|
// FIXME: allow customizing distance
|
||||||
let writer = VectorStore::new(
|
let writer =
|
||||||
index_version,
|
VectorStore::new(backend, index.vector_store, infos.embedder_id, binary_quantized);
|
||||||
index.vector_store,
|
|
||||||
infos.embedder_id,
|
|
||||||
binary_quantized,
|
|
||||||
);
|
|
||||||
|
|
||||||
// remove vectors for docids we want them removed
|
// remove vectors for docids we want them removed
|
||||||
let merger = remove_vectors_builder.build();
|
let merger = remove_vectors_builder.build();
|
||||||
|
@@ -131,7 +131,7 @@ where
|
|||||||
let global_fields_ids_map = GlobalFieldsIdsMap::new(&new_fields_ids_map);
|
let global_fields_ids_map = GlobalFieldsIdsMap::new(&new_fields_ids_map);
|
||||||
|
|
||||||
let vector_arroy = index.vector_store;
|
let vector_arroy = index.vector_store;
|
||||||
let index_version = index.get_version(wtxn)?.unwrap();
|
let backend = index.get_vector_store(wtxn)?;
|
||||||
let hannoy_writers: Result<HashMap<_, _>> = embedders
|
let hannoy_writers: Result<HashMap<_, _>> = embedders
|
||||||
.inner_as_ref()
|
.inner_as_ref()
|
||||||
.iter()
|
.iter()
|
||||||
@@ -145,12 +145,8 @@ where
|
|||||||
})?;
|
})?;
|
||||||
|
|
||||||
let dimensions = runtime.embedder.dimensions();
|
let dimensions = runtime.embedder.dimensions();
|
||||||
let writer = VectorStore::new(
|
let writer =
|
||||||
index_version,
|
VectorStore::new(backend, vector_arroy, embedder_index, runtime.is_quantized);
|
||||||
vector_arroy,
|
|
||||||
embedder_index,
|
|
||||||
runtime.is_quantized,
|
|
||||||
);
|
|
||||||
|
|
||||||
Ok((
|
Ok((
|
||||||
embedder_index,
|
embedder_index,
|
||||||
@@ -352,7 +348,7 @@ fn hannoy_writers_from_embedder_actions<'indexer>(
|
|||||||
index_embedder_category_ids: &'indexer std::collections::HashMap<String, u8>,
|
index_embedder_category_ids: &'indexer std::collections::HashMap<String, u8>,
|
||||||
) -> Result<HashMap<u8, (&'indexer str, &'indexer Embedder, VectorStore, usize)>> {
|
) -> Result<HashMap<u8, (&'indexer str, &'indexer Embedder, VectorStore, usize)>> {
|
||||||
let vector_arroy = index.vector_store;
|
let vector_arroy = index.vector_store;
|
||||||
let index_version = index.get_version(rtxn)?.unwrap();
|
let backend = index.get_vector_store(rtxn)?;
|
||||||
|
|
||||||
embedders
|
embedders
|
||||||
.inner_as_ref()
|
.inner_as_ref()
|
||||||
@@ -371,7 +367,7 @@ fn hannoy_writers_from_embedder_actions<'indexer>(
|
|||||||
)));
|
)));
|
||||||
};
|
};
|
||||||
let writer = VectorStore::new(
|
let writer = VectorStore::new(
|
||||||
index_version,
|
backend,
|
||||||
vector_arroy,
|
vector_arroy,
|
||||||
embedder_category_id,
|
embedder_category_id,
|
||||||
action.was_quantized,
|
action.was_quantized,
|
||||||
@@ -394,16 +390,13 @@ fn delete_old_embedders_and_fragments<SD>(
|
|||||||
where
|
where
|
||||||
SD: SettingsDelta,
|
SD: SettingsDelta,
|
||||||
{
|
{
|
||||||
|
let backend = index.get_vector_store(wtxn)?;
|
||||||
for action in settings_delta.embedder_actions().values() {
|
for action in settings_delta.embedder_actions().values() {
|
||||||
let Some(WriteBackToDocuments { embedder_id, .. }) = action.write_back() else {
|
let Some(WriteBackToDocuments { embedder_id, .. }) = action.write_back() else {
|
||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
let reader = VectorStore::new(
|
let reader =
|
||||||
index.get_version(wtxn)?.unwrap(),
|
VectorStore::new(backend, index.vector_store, *embedder_id, action.was_quantized);
|
||||||
index.vector_store,
|
|
||||||
*embedder_id,
|
|
||||||
action.was_quantized,
|
|
||||||
);
|
|
||||||
let Some(dimensions) = reader.dimensions(wtxn)? else {
|
let Some(dimensions) = reader.dimensions(wtxn)? else {
|
||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
@@ -419,12 +412,7 @@ where
|
|||||||
let Some(infos) = index.embedding_configs().embedder_info(wtxn, embedder_name)? else {
|
let Some(infos) = index.embedding_configs().embedder_info(wtxn, embedder_name)? else {
|
||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
let arroy = VectorStore::new(
|
let arroy = VectorStore::new(backend, index.vector_store, infos.embedder_id, was_quantized);
|
||||||
index.get_version(wtxn)?.unwrap(),
|
|
||||||
index.vector_store,
|
|
||||||
infos.embedder_id,
|
|
||||||
was_quantized,
|
|
||||||
);
|
|
||||||
let Some(dimensions) = arroy.dimensions(wtxn)? else {
|
let Some(dimensions) = arroy.dimensions(wtxn)? else {
|
||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
|
@@ -120,9 +120,9 @@ impl<'t> VectorDocumentFromDb<'t> {
|
|||||||
config: &IndexEmbeddingConfig,
|
config: &IndexEmbeddingConfig,
|
||||||
status: &EmbeddingStatus,
|
status: &EmbeddingStatus,
|
||||||
) -> Result<VectorEntry<'t>> {
|
) -> Result<VectorEntry<'t>> {
|
||||||
let index_version = self.index.get_version(self.rtxn)?.unwrap();
|
let backend = self.index.get_vector_store(self.rtxn)?;
|
||||||
let reader = VectorStore::new(
|
let reader = VectorStore::new(
|
||||||
index_version,
|
backend,
|
||||||
self.index.vector_store,
|
self.index.vector_store,
|
||||||
embedder_id,
|
embedder_id,
|
||||||
config.config.quantized(),
|
config.config.quantized(),
|
||||||
|
@@ -17,13 +17,14 @@ impl UpgradeIndex for Latest_V1_18_New_Hannoy {
|
|||||||
progress: Progress,
|
progress: Progress,
|
||||||
) -> Result<bool> {
|
) -> Result<bool> {
|
||||||
let embedding_configs = index.embedding_configs();
|
let embedding_configs = index.embedding_configs();
|
||||||
let index_version = index.get_version(wtxn)?.unwrap();
|
let backend = index.get_vector_store(wtxn)?;
|
||||||
for config in embedding_configs.embedding_configs(wtxn)? {
|
for config in embedding_configs.embedding_configs(wtxn)? {
|
||||||
// TODO use the embedder name to display progress
|
// TODO use the embedder name to display progress
|
||||||
|
/// REMOVE THIS FILE, IMPLEMENT CONVERSION AS A SETTING CHANGE
|
||||||
let quantized = config.config.quantized();
|
let quantized = config.config.quantized();
|
||||||
let embedder_id = embedding_configs.embedder_id(wtxn, &config.name)?.unwrap();
|
let embedder_id = embedding_configs.embedder_id(wtxn, &config.name)?.unwrap();
|
||||||
let vector_store =
|
let vector_store =
|
||||||
VectorStore::new(index_version, index.vector_store, embedder_id, quantized);
|
VectorStore::new(backend, index.vector_store, embedder_id, quantized);
|
||||||
vector_store.convert_from_arroy(wtxn, progress.clone())?;
|
vector_store.convert_from_arroy(wtxn, progress.clone())?;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@@ -19,7 +19,7 @@ pub use distribution::DistributionShift;
|
|||||||
pub use embedder::{Embedder, EmbedderOptions, EmbeddingConfig, SearchQuery};
|
pub use embedder::{Embedder, EmbedderOptions, EmbeddingConfig, SearchQuery};
|
||||||
pub use embeddings::Embeddings;
|
pub use embeddings::Embeddings;
|
||||||
pub use runtime::{RuntimeEmbedder, RuntimeEmbedders, RuntimeFragment};
|
pub use runtime::{RuntimeEmbedder, RuntimeEmbedders, RuntimeFragment};
|
||||||
pub use store::{HannoyStats, VectorStore};
|
pub use store::{HannoyStats, VectorStore, VectorStoreBackend};
|
||||||
|
|
||||||
pub const REQUEST_PARALLELISM: usize = 40;
|
pub const REQUEST_PARALLELISM: usize = 40;
|
||||||
|
|
||||||
|
@@ -4,6 +4,7 @@ use heed::{RoTxn, RwTxn, Unspecified};
|
|||||||
use ordered_float::OrderedFloat;
|
use ordered_float::OrderedFloat;
|
||||||
use rand::SeedableRng as _;
|
use rand::SeedableRng as _;
|
||||||
use roaring::RoaringBitmap;
|
use roaring::RoaringBitmap;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
use crate::progress::Progress;
|
use crate::progress::Progress;
|
||||||
use crate::vector::Embeddings;
|
use crate::vector::Embeddings;
|
||||||
@@ -12,8 +13,15 @@ const HANNOY_EF_CONSTRUCTION: usize = 125;
|
|||||||
const HANNOY_M: usize = 16;
|
const HANNOY_M: usize = 16;
|
||||||
const HANNOY_M0: usize = 32;
|
const HANNOY_M0: usize = 32;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
|
||||||
|
pub enum VectorStoreBackend {
|
||||||
|
#[default]
|
||||||
|
Arroy,
|
||||||
|
Hannoy,
|
||||||
|
}
|
||||||
|
|
||||||
pub struct VectorStore {
|
pub struct VectorStore {
|
||||||
version: (u32, u32, u32),
|
backend: VectorStoreBackend,
|
||||||
database: hannoy::Database<Unspecified>,
|
database: hannoy::Database<Unspecified>,
|
||||||
embedder_index: u8,
|
embedder_index: u8,
|
||||||
quantized: bool,
|
quantized: bool,
|
||||||
@@ -21,24 +29,18 @@ pub struct VectorStore {
|
|||||||
|
|
||||||
impl VectorStore {
|
impl VectorStore {
|
||||||
pub fn new(
|
pub fn new(
|
||||||
version: (u32, u32, u32),
|
backend: VectorStoreBackend,
|
||||||
database: hannoy::Database<Unspecified>,
|
database: hannoy::Database<Unspecified>,
|
||||||
embedder_index: u8,
|
embedder_index: u8,
|
||||||
quantized: bool,
|
quantized: bool,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
Self { version, database, embedder_index, quantized }
|
Self { backend, database, embedder_index, quantized }
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn embedder_index(&self) -> u8 {
|
pub fn embedder_index(&self) -> u8 {
|
||||||
self.embedder_index
|
self.embedder_index
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Whether we must use the arroy to read the vector store.
|
|
||||||
pub fn version_uses_arroy(&self) -> bool {
|
|
||||||
let (major, minor, _patch) = self.version;
|
|
||||||
major == 1 && minor < 18
|
|
||||||
}
|
|
||||||
|
|
||||||
fn arroy_readers<'a, D: arroy::Distance>(
|
fn arroy_readers<'a, D: arroy::Distance>(
|
||||||
&'a self,
|
&'a self,
|
||||||
rtxn: &'a RoTxn<'a>,
|
rtxn: &'a RoTxn<'a>,
|
||||||
@@ -87,7 +89,7 @@ impl VectorStore {
|
|||||||
where
|
where
|
||||||
F: FnOnce(&RoaringBitmap) -> O,
|
F: FnOnce(&RoaringBitmap) -> O,
|
||||||
{
|
{
|
||||||
if self.version_uses_arroy() {
|
if self.backend == VectorStoreBackend::Arroy {
|
||||||
if self.quantized {
|
if self.quantized {
|
||||||
self._arroy_items_in_store(rtxn, self.arroy_quantized_db(), store_id, with_items)
|
self._arroy_items_in_store(rtxn, self.arroy_quantized_db(), store_id, with_items)
|
||||||
.map_err(Into::into)
|
.map_err(Into::into)
|
||||||
@@ -142,7 +144,7 @@ impl VectorStore {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn dimensions(&self, rtxn: &RoTxn) -> crate::Result<Option<usize>> {
|
pub fn dimensions(&self, rtxn: &RoTxn) -> crate::Result<Option<usize>> {
|
||||||
if self.version_uses_arroy() {
|
if self.backend == VectorStoreBackend::Arroy {
|
||||||
if self.quantized {
|
if self.quantized {
|
||||||
Ok(self
|
Ok(self
|
||||||
.arroy_readers(rtxn, self.arroy_quantized_db())
|
.arroy_readers(rtxn, self.arroy_quantized_db())
|
||||||
@@ -497,7 +499,7 @@ impl VectorStore {
|
|||||||
item: hannoy::ItemId,
|
item: hannoy::ItemId,
|
||||||
) -> crate::Result<bool> {
|
) -> crate::Result<bool> {
|
||||||
for index in vector_store_range_for_embedder(self.embedder_index) {
|
for index in vector_store_range_for_embedder(self.embedder_index) {
|
||||||
let contains = if self.version_uses_arroy() {
|
let contains = if self.backend == VectorStoreBackend::Arroy {
|
||||||
if self.quantized {
|
if self.quantized {
|
||||||
let writer = arroy::Writer::new(self.arroy_quantized_db(), index, dimension);
|
let writer = arroy::Writer::new(self.arroy_quantized_db(), index, dimension);
|
||||||
if writer.is_empty(rtxn)? {
|
if writer.is_empty(rtxn)? {
|
||||||
@@ -538,7 +540,7 @@ impl VectorStore {
|
|||||||
limit: usize,
|
limit: usize,
|
||||||
filter: Option<&RoaringBitmap>,
|
filter: Option<&RoaringBitmap>,
|
||||||
) -> crate::Result<Vec<(ItemId, f32)>> {
|
) -> crate::Result<Vec<(ItemId, f32)>> {
|
||||||
if self.version_uses_arroy() {
|
if self.backend == VectorStoreBackend::Arroy {
|
||||||
if self.quantized {
|
if self.quantized {
|
||||||
self._arroy_nns_by_item(rtxn, self.arroy_quantized_db(), item, limit, filter)
|
self._arroy_nns_by_item(rtxn, self.arroy_quantized_db(), item, limit, filter)
|
||||||
.map_err(Into::into)
|
.map_err(Into::into)
|
||||||
@@ -614,7 +616,7 @@ impl VectorStore {
|
|||||||
limit: usize,
|
limit: usize,
|
||||||
filter: Option<&RoaringBitmap>,
|
filter: Option<&RoaringBitmap>,
|
||||||
) -> crate::Result<Vec<(ItemId, f32)>> {
|
) -> crate::Result<Vec<(ItemId, f32)>> {
|
||||||
if self.version_uses_arroy() {
|
if self.backend == VectorStoreBackend::Arroy {
|
||||||
if self.quantized {
|
if self.quantized {
|
||||||
self._arroy_nns_by_vector(rtxn, self.arroy_quantized_db(), vector, limit, filter)
|
self._arroy_nns_by_vector(rtxn, self.arroy_quantized_db(), vector, limit, filter)
|
||||||
.map_err(Into::into)
|
.map_err(Into::into)
|
||||||
@@ -687,7 +689,7 @@ impl VectorStore {
|
|||||||
pub fn item_vectors(&self, rtxn: &RoTxn, item_id: u32) -> crate::Result<Vec<Vec<f32>>> {
|
pub fn item_vectors(&self, rtxn: &RoTxn, item_id: u32) -> crate::Result<Vec<Vec<f32>>> {
|
||||||
let mut vectors = Vec::new();
|
let mut vectors = Vec::new();
|
||||||
|
|
||||||
if self.version_uses_arroy() {
|
if self.backend == VectorStoreBackend::Arroy {
|
||||||
if self.quantized {
|
if self.quantized {
|
||||||
for reader in self.arroy_readers(rtxn, self.arroy_quantized_db()) {
|
for reader in self.arroy_readers(rtxn, self.arroy_quantized_db()) {
|
||||||
if let Some(vec) = reader?.item_vector(rtxn, item_id)? {
|
if let Some(vec) = reader?.item_vector(rtxn, item_id)? {
|
||||||
|
Reference in New Issue
Block a user