Switch from version to backend selector

This commit is contained in:
Louis Dureuil
2025-08-26 17:49:56 +02:00
parent b5f0c19406
commit da6fffdf6d
12 changed files with 88 additions and 78 deletions

View File

@ -31,7 +31,7 @@ use crate::prompt::PromptData;
use crate::proximity::ProximityPrecision;
use crate::update::new::StdResult;
use crate::vector::db::IndexEmbeddingConfigs;
use crate::vector::{Embedding, HannoyStats, VectorStore};
use crate::vector::{Embedding, HannoyStats, VectorStore, VectorStoreBackend};
use crate::{
default_criteria, CboRoaringBitmapCodec, Criterion, DocumentId, ExternalDocumentsIds,
FacetDistribution, FieldDistribution, FieldId, FieldIdMapMissingEntry, FieldIdWordCountCodec,
@ -87,6 +87,7 @@ pub mod main_key {
pub const DOCUMENTS_STATS: &str = "documents_stats";
pub const DISABLED_TYPOS_TERMS: &str = "disabled_typos_terms";
pub const CHAT: &str = "chat";
pub const VECTOR_STORE_BACKEND: &str = "vector_store_backend";
}
pub mod db_name {
@ -454,6 +455,35 @@ impl Index {
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 */
/// Writes the documents ids that corresponds to the user-ids-documents-ids FST.
@ -1769,12 +1799,13 @@ impl Index {
) -> Result<BTreeMap<String, EmbeddingsWithMetadata>> {
let mut res = BTreeMap::new();
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)? {
let embedder_info = embedders.embedder_info(rtxn, &config.name)?.unwrap();
let has_fragments = config.config.embedder_options.has_fragments();
let reader = VectorStore::new(
index_version,
backend,
self.vector_store,
embedder_info.embedder_id,
config.config.quantized(),
@ -1797,11 +1828,12 @@ impl Index {
pub fn hannoy_stats(&self, rtxn: &RoTxn<'_>) -> Result<HannoyStats> {
let mut stats = HannoyStats::default();
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)? {
let embedder_id = embedding_configs.embedder_id(rtxn, &config.name)?.unwrap();
let reader = VectorStore::new(
index_version,
backend,
self.vector_store,
embedder_id,
config.config.quantized(),

View File

@ -82,7 +82,7 @@ fn evaluate_inner(
embedding_configs: &[IndexEmbeddingConfig],
filter: &VectorFilter<'_>,
) -> crate::Result<RoaringBitmap> {
let index_version = index.get_version(rtxn)?.unwrap();
let backend = index.get_vector_store(rtxn)?;
let embedder_name = embedder.value();
let available_embedders =
|| 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() })?;
let vector_store = VectorStore::new(
index_version,
backend,
index.vector_store,
embedder_info.embedder_id,
embedding_config.config.quantized(),

View File

@ -54,14 +54,11 @@ impl<Q: RankingRuleQueryTrait> VectorSort<Q> {
vector_candidates: &RoaringBitmap,
) -> Result<()> {
let target = &self.target;
let backend = ctx.index.get_vector_store(ctx.txn)?;
let before = Instant::now();
let reader = VectorStore::new(
ctx.index.get_version(ctx.txn)?.unwrap(),
ctx.index.vector_store,
self.embedder_index,
self.quantized,
);
let reader =
VectorStore::new(backend, ctx.index.vector_store, self.embedder_index, self.quantized);
let results = reader.nns_by_vector(ctx.txn, target, self.limit, Some(vector_candidates))?;
self.cached_sorted_docids = results.into_iter();
*ctx.vector_store_stats.get_or_insert_default() += VectorStoreStats {

View File

@ -72,12 +72,10 @@ impl<'a> Similar<'a> {
crate::UserError::InvalidSimilarEmbedder(self.embedder_name.to_owned())
})?;
let reader = VectorStore::new(
self.index.get_version(self.rtxn)?.unwrap(),
self.index.vector_store,
embedder_index,
self.quantized,
);
let backend = self.index.get_vector_store(self.rtxn)?;
let reader =
VectorStore::new(backend, self.index.vector_store, embedder_index, self.quantized);
let results = reader.nns_by_item(
self.rtxn,
self.id,

View File

@ -485,7 +485,7 @@ where
// If an embedder wasn't used in the typedchunk but must be binary quantized
// 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() {
if action.is_being_quantized && !dimension.contains_key(name.as_str()) {
let index = self.index.embedding_configs().embedder_id(self.wtxn, name)?.ok_or(
@ -494,12 +494,8 @@ where
key: None,
},
)?;
let reader = VectorStore::new(
index_version,
self.index.vector_store,
index,
action.was_quantized,
);
let reader =
VectorStore::new(backend, self.index.vector_store, index, action.was_quantized);
let Some(dim) = reader.dimensions(self.wtxn)? else {
continue;
};
@ -529,7 +525,7 @@ where
pool.install(|| {
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(
wtxn,
// In the settings we don't have any progress to share

View File

@ -834,7 +834,7 @@ impl<'a, 'i> Transform<'a, 'i> {
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
.embedding_config_updates
.iter()
@ -843,7 +843,7 @@ impl<'a, 'i> Transform<'a, 'i> {
action.write_back()
{
let reader = VectorStore::new(
index_version,
backend,
self.index.vector_store,
*embedder_id,
action.was_quantized,
@ -949,7 +949,7 @@ impl<'a, 'i> Transform<'a, 'i> {
continue;
};
let hannoy = VectorStore::new(
index_version,
backend,
self.index.vector_store,
infos.embedder_id,
was_quantized,

View File

@ -619,7 +619,7 @@ pub(crate) fn write_typed_chunk_into_index(
let _entered = span.enter();
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 manual_vectors_builder = MergerBuilder::new(KeepFirst);
@ -678,12 +678,8 @@ pub(crate) fn write_typed_chunk_into_index(
.get(&embedder_name)
.is_some_and(|conf| conf.is_quantized);
// FIXME: allow customizing distance
let writer = VectorStore::new(
index_version,
index.vector_store,
infos.embedder_id,
binary_quantized,
);
let writer =
VectorStore::new(backend, index.vector_store, infos.embedder_id, binary_quantized);
// remove vectors for docids we want them removed
let merger = remove_vectors_builder.build();

View File

@ -131,7 +131,7 @@ where
let global_fields_ids_map = GlobalFieldsIdsMap::new(&new_fields_ids_map);
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
.inner_as_ref()
.iter()
@ -145,12 +145,8 @@ where
})?;
let dimensions = runtime.embedder.dimensions();
let writer = VectorStore::new(
index_version,
vector_arroy,
embedder_index,
runtime.is_quantized,
);
let writer =
VectorStore::new(backend, vector_arroy, embedder_index, runtime.is_quantized);
Ok((
embedder_index,
@ -352,7 +348,7 @@ fn hannoy_writers_from_embedder_actions<'indexer>(
index_embedder_category_ids: &'indexer std::collections::HashMap<String, u8>,
) -> Result<HashMap<u8, (&'indexer str, &'indexer Embedder, VectorStore, usize)>> {
let vector_arroy = index.vector_store;
let index_version = index.get_version(rtxn)?.unwrap();
let backend = index.get_vector_store(rtxn)?;
embedders
.inner_as_ref()
@ -371,7 +367,7 @@ fn hannoy_writers_from_embedder_actions<'indexer>(
)));
};
let writer = VectorStore::new(
index_version,
backend,
vector_arroy,
embedder_category_id,
action.was_quantized,
@ -394,16 +390,13 @@ fn delete_old_embedders_and_fragments<SD>(
where
SD: SettingsDelta,
{
let backend = index.get_vector_store(wtxn)?;
for action in settings_delta.embedder_actions().values() {
let Some(WriteBackToDocuments { embedder_id, .. }) = action.write_back() else {
continue;
};
let reader = VectorStore::new(
index.get_version(wtxn)?.unwrap(),
index.vector_store,
*embedder_id,
action.was_quantized,
);
let reader =
VectorStore::new(backend, index.vector_store, *embedder_id, action.was_quantized);
let Some(dimensions) = reader.dimensions(wtxn)? else {
continue;
};
@ -419,12 +412,7 @@ where
let Some(infos) = index.embedding_configs().embedder_info(wtxn, embedder_name)? else {
continue;
};
let arroy = VectorStore::new(
index.get_version(wtxn)?.unwrap(),
index.vector_store,
infos.embedder_id,
was_quantized,
);
let arroy = VectorStore::new(backend, index.vector_store, infos.embedder_id, was_quantized);
let Some(dimensions) = arroy.dimensions(wtxn)? else {
continue;
};

View File

@ -120,9 +120,9 @@ impl<'t> VectorDocumentFromDb<'t> {
config: &IndexEmbeddingConfig,
status: &EmbeddingStatus,
) -> 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(
index_version,
backend,
self.index.vector_store,
embedder_id,
config.config.quantized(),

View File

@ -17,13 +17,14 @@ impl UpgradeIndex for Latest_V1_18_New_Hannoy {
progress: Progress,
) -> Result<bool> {
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)? {
// TODO use the embedder name to display progress
/// REMOVE THIS FILE, IMPLEMENT CONVERSION AS A SETTING CHANGE
let quantized = config.config.quantized();
let embedder_id = embedding_configs.embedder_id(wtxn, &config.name)?.unwrap();
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())?;
}

View File

@ -19,7 +19,7 @@ pub use distribution::DistributionShift;
pub use embedder::{Embedder, EmbedderOptions, EmbeddingConfig, SearchQuery};
pub use embeddings::Embeddings;
pub use runtime::{RuntimeEmbedder, RuntimeEmbedders, RuntimeFragment};
pub use store::{HannoyStats, VectorStore};
pub use store::{HannoyStats, VectorStore, VectorStoreBackend};
pub const REQUEST_PARALLELISM: usize = 40;

View File

@ -4,6 +4,7 @@ use heed::{RoTxn, RwTxn, Unspecified};
use ordered_float::OrderedFloat;
use rand::SeedableRng as _;
use roaring::RoaringBitmap;
use serde::{Deserialize, Serialize};
use crate::progress::Progress;
use crate::vector::Embeddings;
@ -12,8 +13,15 @@ const HANNOY_EF_CONSTRUCTION: usize = 125;
const HANNOY_M: usize = 16;
const HANNOY_M0: usize = 32;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub enum VectorStoreBackend {
#[default]
Arroy,
Hannoy,
}
pub struct VectorStore {
version: (u32, u32, u32),
backend: VectorStoreBackend,
database: hannoy::Database<Unspecified>,
embedder_index: u8,
quantized: bool,
@ -21,24 +29,18 @@ pub struct VectorStore {
impl VectorStore {
pub fn new(
version: (u32, u32, u32),
backend: VectorStoreBackend,
database: hannoy::Database<Unspecified>,
embedder_index: u8,
quantized: bool,
) -> Self {
Self { version, database, embedder_index, quantized }
Self { backend, database, embedder_index, quantized }
}
pub fn embedder_index(&self) -> u8 {
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>(
&'a self,
rtxn: &'a RoTxn<'a>,
@ -87,7 +89,7 @@ impl VectorStore {
where
F: FnOnce(&RoaringBitmap) -> O,
{
if self.version_uses_arroy() {
if self.backend == VectorStoreBackend::Arroy {
if self.quantized {
self._arroy_items_in_store(rtxn, self.arroy_quantized_db(), store_id, with_items)
.map_err(Into::into)
@ -142,7 +144,7 @@ impl VectorStore {
}
pub fn dimensions(&self, rtxn: &RoTxn) -> crate::Result<Option<usize>> {
if self.version_uses_arroy() {
if self.backend == VectorStoreBackend::Arroy {
if self.quantized {
Ok(self
.arroy_readers(rtxn, self.arroy_quantized_db())
@ -497,7 +499,7 @@ impl VectorStore {
item: hannoy::ItemId,
) -> crate::Result<bool> {
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 {
let writer = arroy::Writer::new(self.arroy_quantized_db(), index, dimension);
if writer.is_empty(rtxn)? {
@ -538,7 +540,7 @@ impl VectorStore {
limit: usize,
filter: Option<&RoaringBitmap>,
) -> crate::Result<Vec<(ItemId, f32)>> {
if self.version_uses_arroy() {
if self.backend == VectorStoreBackend::Arroy {
if self.quantized {
self._arroy_nns_by_item(rtxn, self.arroy_quantized_db(), item, limit, filter)
.map_err(Into::into)
@ -614,7 +616,7 @@ impl VectorStore {
limit: usize,
filter: Option<&RoaringBitmap>,
) -> crate::Result<Vec<(ItemId, f32)>> {
if self.version_uses_arroy() {
if self.backend == VectorStoreBackend::Arroy {
if self.quantized {
self._arroy_nns_by_vector(rtxn, self.arroy_quantized_db(), vector, limit, filter)
.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>>> {
let mut vectors = Vec::new();
if self.version_uses_arroy() {
if self.backend == VectorStoreBackend::Arroy {
if self.quantized {
for reader in self.arroy_readers(rtxn, self.arroy_quantized_db()) {
if let Some(vec) = reader?.item_vector(rtxn, item_id)? {