fix the cellulite integration

This commit is contained in:
Tamo
2025-07-23 15:25:38 +02:00
parent 610e62c0e1
commit fde155ed22
7 changed files with 839 additions and 4 deletions

View File

@ -5,6 +5,7 @@ use std::fmt;
use std::fs::File;
use std::path::Path;
use cellulite::Cellulite;
use deserr::Deserr;
use heed::types::*;
use heed::{CompactionOption, Database, DatabaseStat, RoTxn, RwTxn, Unspecified, WithoutTls};
@ -118,7 +119,7 @@ pub mod db_name {
pub const CELLULITE: &str = "cellulite";
pub const DOCUMENTS: &str = "documents";
}
const NUMBER_OF_DBS: u32 = 25;
const NUMBER_OF_DBS: u32 = 25 + Cellulite::nb_dbs();
#[derive(Clone)]
pub struct Index {
@ -185,7 +186,7 @@ pub struct Index {
pub vector_store: hannoy::Database<Unspecified>,
/// Geo store based on cellulite™.
pub cellulite: cellulite::Cellulite,
pub cellulite: Cellulite,
/// Maps the document id to the document as an obkv store.
pub(crate) documents: Database<BEU32, ObkvCodec>,

View File

@ -166,8 +166,6 @@ where
indexing_context.progress.update_progress(IndexingStep::WritingEmbeddingsToDatabase);
index.cellulite.build(wtxn, indexing_context.progress)?;
pool.install(|| {
build_vectors(
index,
@ -189,6 +187,10 @@ where
facet_field_ids_delta,
)?;
println!("Building geojson");
indexing_context.progress.update_progress(IndexingStep::BuildingGeoJson);
index.cellulite.build(wtxn, indexing_context.progress)?;
indexing_context.progress.update_progress(IndexingStep::Finalizing);
Ok(congestion) as Result<_>
@ -319,6 +321,9 @@ where
})
.unwrap()?;
indexing_context.progress.update_progress(IndexingStep::BuildingGeoJson);
index.cellulite.build(wtxn, indexing_context.progress)?;
indexing_context.progress.update_progress(IndexingStep::Finalizing);
Ok(congestion) as Result<_>

View File

@ -75,9 +75,11 @@ pub fn write_to_db(
ReceiverAction::GeoJson(docid, geojson) => {
match geojson {
Some(geojson) => {
println!("Adding geojson {docid}");
index.cellulite.add(wtxn, docid, &geojson).map_err(InternalError::CelluliteError)?;
}
None => {
println!("Deleting geojson {docid}");
index.cellulite.delete(wtxn, docid).map_err(InternalError::CelluliteError)?;
}
}

View File

@ -18,6 +18,7 @@ make_enum_progress! {
WritingEmbeddingsToDatabase,
PostProcessingFacets,
PostProcessingWords,
BuildingGeoJson,
Finalizing,
}
}