WIP

Uses the same `byte_unit` version as `meilisearch`

.github/workflows/publish-docker-images.yml

@@ -97,7 +97,7 @@ jobs:
       - name: Send CI information to Cloud team
         # Do not send if nightly build (i.e. 'schedule' or 'workflow_dispatch' event)
         if: github.event_name == 'push'
-        uses: peter-evans/repository-dispatch@v3
+        uses: peter-evans/repository-dispatch@v2
         with:
           token: ${{ secrets.MEILI_BOT_GH_PAT }}
           repository: meilisearch/meilisearch-cloud

Cargo.toml

@@ -21,7 +21,7 @@ members = [
 ]
 
 [workspace.package]
-version = "1.7.0"
+version = "1.6.0"
 authors = [
     "Quentin de Quelen <quentin@dequelen.me>",
     "Clément Renault <clement@meilisearch.com>",

README.md

@@ -41,10 +41,10 @@ Meilisearch helps you shape a delightful search experience in a snap, offering f
 ## ✨ Features
 
 - **Search-as-you-type:** find search results in less than 50 milliseconds
-- **[Typo tolerance](https://www.meilisearch.com/docs/learn/configuration/typo_tolerance?utm_campaign=oss&utm_source=github&utm_medium=meilisearch&utm_content=features):** get relevant matches even when queries contain typos and misspellings
+- **[Typo tolerance](https://www.meilisearch.com/docs/learn/getting_started/customizing_relevancy?utm_campaign=oss&utm_source=github&utm_medium=meilisearch&utm_content=features#typo-tolerance):** get relevant matches even when queries contain typos and misspellings
 - **[Filtering](https://www.meilisearch.com/docs/learn/fine_tuning_results/filtering?utm_campaign=oss&utm_source=github&utm_medium=meilisearch&utm_content=features) and [faceted search](https://www.meilisearch.com/docs/learn/fine_tuning_results/faceted_search?utm_campaign=oss&utm_source=github&utm_medium=meilisearch&utm_content=features):** enhance your users' search experience with custom filters and build a faceted search interface in a few lines of code
 - **[Sorting](https://www.meilisearch.com/docs/learn/fine_tuning_results/sorting?utm_campaign=oss&utm_source=github&utm_medium=meilisearch&utm_content=features):** sort results based on price, date, or pretty much anything else your users need
-- **[Synonym support](https://www.meilisearch.com/docs/learn/configuration/synonyms?utm_campaign=oss&utm_source=github&utm_medium=meilisearch&utm_content=features):** configure synonyms to include more relevant content in your search results
+- **[Synonym support](https://www.meilisearch.com/docs/learn/getting_started/customizing_relevancy?utm_campaign=oss&utm_source=github&utm_medium=meilisearch&utm_content=features#synonyms):** configure synonyms to include more relevant content in your search results
 - **[Geosearch](https://www.meilisearch.com/docs/learn/fine_tuning_results/geosearch?utm_campaign=oss&utm_source=github&utm_medium=meilisearch&utm_content=features):** filter and sort documents based on geographic data
 - **[Extensive language support](https://www.meilisearch.com/docs/learn/what_is_meilisearch/language?utm_campaign=oss&utm_source=github&utm_medium=meilisearch&utm_content=features):** search datasets in any language, with optimized support for Chinese, Japanese, Hebrew, and languages using the Latin alphabet
 - **[Security management](https://www.meilisearch.com/docs/learn/security/master_api_keys?utm_campaign=oss&utm_source=github&utm_medium=meilisearch&utm_content=features):** control which users can access what data with API keys that allow fine-grained permissions handling
@@ -61,6 +61,8 @@ You can consult Meilisearch's documentation at [https://www.meilisearch.com/docs
 
 For basic instructions on how to set up Meilisearch, add documents to an index, and search for documents, take a look at our [Quick Start](https://www.meilisearch.com/docs/learn/getting_started/quick_start?utm_campaign=oss&utm_source=github&utm_medium=meilisearch&utm_content=get-started) guide.
 
+You may also want to check out [Meilisearch 101](https://www.meilisearch.com/docs/learn/getting_started/filtering_and_sorting?utm_campaign=oss&utm_source=github&utm_medium=meilisearch&utm_content=get-started) for an introduction to some of Meilisearch's most popular features.
+
 ## ⚡ Supercharge your Meilisearch experience
 
 Say goodbye to server deployment and manual updates with [Meilisearch Cloud](https://www.meilisearch.com/cloud?utm_campaign=oss&utm_source=github&utm_medium=meilisearch). No credit card required.
@@ -99,7 +101,7 @@ Meilisearch is a search engine created by [Meili](https://www.welcometothejungle
 
 - For feature requests, please visit our [product repository](https://github.com/meilisearch/product/discussions)
 - Found a bug? Open an [issue](https://github.com/meilisearch/meilisearch/issues)!
-- Want to be part of our Discord community? [Join us!](https://discord.meilisearch.com/?utm_campaign=oss&utm_source=github&utm_medium=meilisearch&utm_content=contact)
+- Want to be part of our Discord community? [Join us!](https://discord.gg/meilisearch)
 
 Thank you for your support!
 

index-scheduler/src/features.rs

@@ -30,6 +30,19 @@ impl RoFeatures {
         self.runtime
     }
 
+    pub fn check_score_details(&self) -> Result<()> {
+        if self.runtime.score_details {
+            Ok(())
+        } else {
+            Err(FeatureNotEnabledError {
+                disabled_action: "Computing score details",
+                feature: "score details",
+                issue_link: "https://github.com/meilisearch/product/discussions/674",
+            }
+            .into())
+        }
+    }
+
     pub fn check_metrics(&self) -> Result<()> {
         if self.runtime.metrics {
             Ok(())
@@ -48,7 +61,7 @@ impl RoFeatures {
             Ok(())
         } else {
             Err(FeatureNotEnabledError {
-                disabled_action: "Modifying logs through the `/logs/*` routes",
+                disabled_action: "getting logs through the `/logs/stream` route",
                 feature: "logs route",
                 issue_link: "https://github.com/orgs/meilisearch/discussions/721",
             }

meilisearch-types/Cargo.toml

@@ -54,5 +54,3 @@ thai = ["milli/thai"]
 greek = ["milli/greek"]
 # allow khmer specialized tokenization
 khmer = ["milli/khmer"]
-# allow vietnamese specialized tokenization
-vietnamese = ["milli/vietnamese"]

meilisearch-types/src/error.rs

@@ -349,9 +349,6 @@ impl ErrorCode for milli::Error {
                     UserError::InvalidFieldForSource { .. }
                     | UserError::MissingFieldForSource { .. }
                     | UserError::InvalidOpenAiModel { .. }
-                    | UserError::InvalidOpenAiModelDimensions { .. }
-                    | UserError::InvalidOpenAiModelDimensionsMax { .. }
-                    | UserError::InvalidSettingsDimensions { .. }
                     | UserError::InvalidPrompt(_) => Code::InvalidSettingsEmbedders,
                     UserError::TooManyEmbedders(_) => Code::InvalidSettingsEmbedders,
                     UserError::InvalidPromptForEmbeddings(..) => Code::InvalidSettingsEmbedders,

meilisearch-types/src/features.rs

@@ -3,6 +3,7 @@ use serde::{Deserialize, Serialize};
 #[derive(Serialize, Deserialize, Debug, Clone, Copy, Default, PartialEq, Eq)]
 #[serde(rename_all = "camelCase", default)]
 pub struct RuntimeTogglableFeatures {
+    pub score_details: bool,
     pub vector_store: bool,
     pub metrics: bool,
     pub logs_route: bool,

meilisearch/Cargo.toml

@@ -104,7 +104,7 @@ serde_urlencoded = "0.7.1"
 termcolor = "1.4.1"
 url = { version = "2.5.0", features = ["serde"] }
 tracing = "0.1.40"
-tracing-subscriber = { version = "0.3.18", features = ["json"] }
+tracing-subscriber = "0.3.18"
 tracing-trace = { version = "0.1.0", path = "../tracing-trace" }
 tracing-actix-web = "0.7.9"
 
@@ -154,8 +154,7 @@ japanese = ["meilisearch-types/japanese"]
 thai = ["meilisearch-types/thai"]
 greek = ["meilisearch-types/greek"]
 khmer = ["meilisearch-types/khmer"]
-vietnamese = ["meilisearch-types/vietnamese"]
 
 [package.metadata.mini-dashboard]
-assets-url = "https://github.com/meilisearch/mini-dashboard/releases/download/v0.2.13/build.zip"
-sha1 = "e20cc9b390003c6c844f4b8bcc5c5013191a77ff"
+assets-url = "https://github.com/meilisearch/mini-dashboard/releases/download/v0.2.12/build.zip"
+sha1 = "acfe9a018c93eb0604ea87ee87bff7df5474e18e"

meilisearch/src/analytics/segment_analytics.rs

@@ -28,9 +28,7 @@ use super::{
     config_user_id_path, DocumentDeletionKind, DocumentFetchKind, MEILISEARCH_CONFIG_PATH,
 };
 use crate::analytics::Analytics;
-use crate::option::{
-    default_http_addr, IndexerOpts, LogMode, MaxMemory, MaxThreads, ScheduleSnapshot,
-};
+use crate::option::{default_http_addr, IndexerOpts, MaxMemory, MaxThreads, ScheduleSnapshot};
 use crate::routes::indexes::documents::UpdateDocumentsQuery;
 use crate::routes::indexes::facet_search::FacetSearchQuery;
 use crate::routes::tasks::TasksFilterQuery;
@@ -252,7 +250,6 @@ impl super::Analytics for SegmentAnalytics {
 struct Infos {
     env: String,
     experimental_enable_metrics: bool,
-    experimental_logs_mode: LogMode,
     experimental_enable_logs_route: bool,
     experimental_reduce_indexing_memory_usage: bool,
     experimental_max_number_of_batched_tasks: usize,
@@ -291,7 +288,6 @@ impl From<Opt> for Infos {
         let Opt {
             db_path,
             experimental_enable_metrics,
-            experimental_logs_mode,
             experimental_enable_logs_route,
             experimental_reduce_indexing_memory_usage,
             experimental_max_number_of_batched_tasks,
@@ -339,7 +335,6 @@ impl From<Opt> for Infos {
         Self {
             env,
             experimental_enable_metrics,
-            experimental_logs_mode,
             experimental_enable_logs_route,
             experimental_reduce_indexing_memory_usage,
             db_path: db_path != PathBuf::from("./data.ms"),

meilisearch/src/lib.rs

@@ -38,7 +38,7 @@ use meilisearch_types::versioning::{check_version_file, create_version_file};
 use meilisearch_types::{compression, milli, VERSION_FILE_NAME};
 pub use option::Opt;
 use option::ScheduleSnapshot;
-use tracing::{error, info_span};
+use tracing::error;
 use tracing_subscriber::filter::Targets;
 
 use crate::error::MeilisearchHttpError;
@@ -97,25 +97,11 @@ pub type LogRouteType = tracing_subscriber::filter::Filtered<
     tracing_subscriber::Registry,
 >;
 
-pub type SubscriberForSecondLayer = tracing_subscriber::layer::Layered<
-    tracing_subscriber::reload::Layer<LogRouteType, tracing_subscriber::Registry>,
-    tracing_subscriber::Registry,
->;
-
-pub type LogStderrHandle =
-    tracing_subscriber::reload::Handle<LogStderrType, SubscriberForSecondLayer>;
-
-pub type LogStderrType = tracing_subscriber::filter::Filtered<
-    Box<dyn tracing_subscriber::Layer<SubscriberForSecondLayer> + Send + Sync>,
-    Targets,
-    SubscriberForSecondLayer,
->;
-
 pub fn create_app(
     index_scheduler: Data<IndexScheduler>,
     auth_controller: Data<AuthController>,
     opt: Opt,
-    logs: (LogRouteHandle, LogStderrHandle),
+    logs: LogRouteHandle,
     analytics: Arc<dyn Analytics>,
     enable_dashboard: bool,
 ) -> actix_web::App<
@@ -150,49 +136,11 @@ pub fn create_app(
             .allow_any_method()
             .max_age(86_400), // 24h
     )
-    .wrap(tracing_actix_web::TracingLogger::<AwebTracingLogger>::new())
+    .wrap(tracing_actix_web::TracingLogger::default())
     .wrap(actix_web::middleware::Compress::default())
     .wrap(actix_web::middleware::NormalizePath::new(actix_web::middleware::TrailingSlash::Trim))
 }
 
-struct AwebTracingLogger;
-
-impl tracing_actix_web::RootSpanBuilder for AwebTracingLogger {
-    fn on_request_start(request: &actix_web::dev::ServiceRequest) -> tracing::Span {
-        use tracing::field::Empty;
-
-        let conn_info = request.connection_info();
-        let headers = request.headers();
-        let user_agent = headers
-            .get(http::header::USER_AGENT)
-            .map(|value| String::from_utf8_lossy(value.as_bytes()).into_owned())
-            .unwrap_or_default();
-        info_span!("HTTP request", method = %request.method(), host = conn_info.host(), route = %request.path(), query_parameters = %request.query_string(), %user_agent, status_code = Empty, error = Empty)
-    }
-
-    fn on_request_end<B: MessageBody>(
-        span: tracing::Span,
-        outcome: &Result<ServiceResponse<B>, actix_web::Error>,
-    ) {
-        match &outcome {
-            Ok(response) => {
-                let code: i32 = response.response().status().as_u16().into();
-                span.record("status_code", code);
-
-                if let Some(error) = response.response().error() {
-                    // use the status code already constructed for the outgoing HTTP response
-                    span.record("error", &tracing::field::display(error.as_response_error()));
-                }
-            }
-            Err(error) => {
-                let code: i32 = error.error_response().status().as_u16().into();
-                span.record("status_code", code);
-                span.record("error", &tracing::field::display(error.as_response_error()));
-            }
-        };
-    }
-}
-
 enum OnFailure {
     RemoveDb,
     KeepDb,
@@ -458,7 +406,7 @@ pub fn configure_data(
     index_scheduler: Data<IndexScheduler>,
     auth: Data<AuthController>,
     opt: &Opt,
-    (logs_route, logs_stderr): (LogRouteHandle, LogStderrHandle),
+    logs: LogRouteHandle,
     analytics: Arc<dyn Analytics>,
 ) {
     let http_payload_size_limit = opt.http_payload_size_limit.get_bytes() as usize;
@@ -466,8 +414,7 @@ pub fn configure_data(
         .app_data(index_scheduler)
         .app_data(auth)
         .app_data(web::Data::from(analytics))
-        .app_data(web::Data::new(logs_route))
-        .app_data(web::Data::new(logs_stderr))
+        .app_data(web::Data::new(logs))
         .app_data(
             web::JsonConfig::default()
                 .limit(http_payload_size_limit)

meilisearch/src/main.rs

@@ -1,5 +1,5 @@
 use std::env;
-use std::io::{stderr, LineWriter, Write};
+use std::io::{stderr, Write};
 use std::path::PathBuf;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -10,10 +10,8 @@ use actix_web::HttpServer;
 use index_scheduler::IndexScheduler;
 use is_terminal::IsTerminal;
 use meilisearch::analytics::Analytics;
-use meilisearch::option::LogMode;
 use meilisearch::{
-    analytics, create_app, prototype_name, setup_meilisearch, LogRouteHandle, LogRouteType,
-    LogStderrHandle, LogStderrType, Opt, SubscriberForSecondLayer,
+    analytics, create_app, prototype_name, setup_meilisearch, LogRouteHandle, LogRouteType, Opt,
 };
 use meilisearch_auth::{generate_master_key, AuthController, MASTER_KEY_MIN_SIZE};
 use mimalloc::MiMalloc;
@@ -25,44 +23,28 @@ use tracing_subscriber::Layer;
 #[global_allocator]
 static ALLOC: MiMalloc = MiMalloc;
 
-fn default_log_route_layer() -> LogRouteType {
+fn default_layer() -> LogRouteType {
     None.with_filter(tracing_subscriber::filter::Targets::new().with_target("", LevelFilter::OFF))
 }
 
-fn default_log_stderr_layer(opt: &Opt) -> LogStderrType {
-    let layer = tracing_subscriber::fmt::layer()
-        .with_writer(|| LineWriter::new(std::io::stderr()))
-        .with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE);
-
-    let layer = match opt.experimental_logs_mode {
-        LogMode::Human => Box::new(layer)
-            as Box<dyn tracing_subscriber::Layer<SubscriberForSecondLayer> + Send + Sync>,
-        LogMode::Json => Box::new(layer.json())
-            as Box<dyn tracing_subscriber::Layer<SubscriberForSecondLayer> + Send + Sync>,
-    };
-
-    layer.with_filter(
-        tracing_subscriber::filter::Targets::new()
-            .with_target("", LevelFilter::from_str(&opt.log_level.to_string()).unwrap()),
-    )
-}
-
 /// does all the setup before meilisearch is launched
-fn setup(opt: &Opt) -> anyhow::Result<(LogRouteHandle, LogStderrHandle)> {
-    let (route_layer, route_layer_handle) =
-        tracing_subscriber::reload::Layer::new(default_log_route_layer());
+fn setup(opt: &Opt) -> anyhow::Result<LogRouteHandle> {
+    let (route_layer, route_layer_handle) = tracing_subscriber::reload::Layer::new(default_layer());
     let route_layer: tracing_subscriber::reload::Layer<_, _> = route_layer;
 
-    let (stderr_layer, stderr_layer_handle) =
-        tracing_subscriber::reload::Layer::new(default_log_stderr_layer(opt));
-    let route_layer: tracing_subscriber::reload::Layer<_, _> = route_layer;
-
-    let subscriber = tracing_subscriber::registry().with(route_layer).with(stderr_layer);
+    let subscriber = tracing_subscriber::registry().with(route_layer).with(
+        tracing_subscriber::fmt::layer()
+            .with_span_events(tracing_subscriber::fmt::format::FmtSpan::NEW)
+            .with_filter(
+                tracing_subscriber::filter::LevelFilter::from_str(&opt.log_level.to_string())
+                    .unwrap(),
+            ),
+    );
 
     // set the subscriber as the default for the application
     tracing::subscriber::set_global_default(subscriber).unwrap();
 
-    Ok((route_layer_handle, stderr_layer_handle))
+    Ok(route_layer_handle)
 }
 
 fn on_panic(info: &std::panic::PanicInfo) {
@@ -128,7 +110,7 @@ async fn run_http(
     index_scheduler: Arc<IndexScheduler>,
     auth_controller: Arc<AuthController>,
     opt: Opt,
-    logs: (LogRouteHandle, LogStderrHandle),
+    logs: LogRouteHandle,
     analytics: Arc<dyn Analytics>,
 ) -> anyhow::Result<()> {
     let enable_dashboard = &opt.env == "development";

meilisearch/src/option.rs

@@ -51,7 +51,6 @@ const MEILI_IGNORE_MISSING_DUMP: &str = "MEILI_IGNORE_MISSING_DUMP";
 const MEILI_IGNORE_DUMP_IF_DB_EXISTS: &str = "MEILI_IGNORE_DUMP_IF_DB_EXISTS";
 const MEILI_DUMP_DIR: &str = "MEILI_DUMP_DIR";
 const MEILI_LOG_LEVEL: &str = "MEILI_LOG_LEVEL";
-const MEILI_EXPERIMENTAL_LOGS_MODE: &str = "MEILI_EXPERIMENTAL_LOGS_MODE";
 const MEILI_EXPERIMENTAL_ENABLE_LOGS_ROUTE: &str = "MEILI_EXPERIMENTAL_ENABLE_LOGS_ROUTE";
 const MEILI_EXPERIMENTAL_ENABLE_METRICS: &str = "MEILI_EXPERIMENTAL_ENABLE_METRICS";
 const MEILI_EXPERIMENTAL_REDUCE_INDEXING_MEMORY_USAGE: &str =
@@ -80,39 +79,6 @@ const DEFAULT_LOG_EVERY_N: usize = 100_000;
 pub const INDEX_SIZE: u64 = 2 * 1024 * 1024 * 1024 * 1024; // 2 TiB
 pub const TASK_DB_SIZE: u64 = 20 * 1024 * 1024 * 1024; // 20 GiB
 
-#[derive(Debug, Default, Clone, Copy, Serialize, Deserialize)]
-#[serde(rename_all = "UPPERCASE")]
-pub enum LogMode {
-    #[default]
-    Human,
-    Json,
-}
-
-impl Display for LogMode {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match self {
-            LogMode::Human => Display::fmt("HUMAN", f),
-            LogMode::Json => Display::fmt("JSON", f),
-        }
-    }
-}
-
-impl FromStr for LogMode {
-    type Err = LogModeError;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s.trim().to_lowercase().as_str() {
-            "human" => Ok(LogMode::Human),
-            "json" => Ok(LogMode::Json),
-            _ => Err(LogModeError(s.to_owned())),
-        }
-    }
-}
-
-#[derive(Debug, thiserror::Error)]
-#[error("Unsupported log mode level `{0}`. Supported values are `HUMAN` and `JSON`.")]
-pub struct LogModeError(String);
-
 #[derive(Debug, Default, Clone, Copy, Serialize, Deserialize)]
 #[serde(rename_all = "UPPERCASE")]
 pub enum LogLevel {
@@ -344,16 +310,9 @@ pub struct Opt {
     #[serde(default)]
     pub experimental_enable_metrics: bool,
 
-    /// Experimental logs mode feature. For more information, see: <https://github.com/orgs/meilisearch/discussions/723>
-    ///
-    /// Change the mode of the logs on the console.
-    #[clap(long, env = MEILI_EXPERIMENTAL_LOGS_MODE, default_value_t)]
-    #[serde(default)]
-    pub experimental_logs_mode: LogMode,
-
     /// Experimental logs route feature. For more information, see: <https://github.com/orgs/meilisearch/discussions/721>
     ///
-    /// Enables the log routes on the `POST /logs/stream`, `POST /logs/stderr` endpoints, and the `DELETE /logs/stream` to stop receiving logs.
+    /// Enables the log route on the `POST /logs/stream` endpoint and the `DELETE /logs/stream` to stop receiving logs.
     #[clap(long, env = MEILI_EXPERIMENTAL_ENABLE_LOGS_ROUTE)]
     #[serde(default)]
     pub experimental_enable_logs_route: bool,
@@ -463,7 +422,6 @@ impl Opt {
             #[cfg(feature = "analytics")]
             no_analytics,
             experimental_enable_metrics,
-            experimental_logs_mode,
             experimental_enable_logs_route,
             experimental_reduce_indexing_memory_usage,
         } = self;
@@ -521,10 +479,6 @@ impl Opt {
             MEILI_EXPERIMENTAL_ENABLE_METRICS,
             experimental_enable_metrics.to_string(),
         );
-        export_to_env_if_not_present(
-            MEILI_EXPERIMENTAL_LOGS_MODE,
-            experimental_logs_mode.to_string(),
-        );
         export_to_env_if_not_present(
             MEILI_EXPERIMENTAL_ENABLE_LOGS_ROUTE,
             experimental_enable_logs_route.to_string(),

meilisearch/src/routes/features.rs

@@ -41,6 +41,8 @@ async fn get_features(
 #[derive(Debug, Deserr)]
 #[deserr(error = DeserrJsonError, rename_all = camelCase, deny_unknown_fields)]
 pub struct RuntimeTogglableFeatures {
+    #[deserr(default)]
+    pub score_details: Option<bool>,
     #[deserr(default)]
     pub vector_store: Option<bool>,
     #[deserr(default)]
@@ -65,6 +67,7 @@ async fn patch_features(
 
     let old_features = features.runtime_features();
     let new_features = meilisearch_types::features::RuntimeTogglableFeatures {
+        score_details: new_features.0.score_details.unwrap_or(old_features.score_details),
         vector_store: new_features.0.vector_store.unwrap_or(old_features.vector_store),
         metrics: new_features.0.metrics.unwrap_or(old_features.metrics),
         logs_route: new_features.0.logs_route.unwrap_or(old_features.logs_route),
@@ -78,6 +81,7 @@ async fn patch_features(
     // the it renames to camelCase, which we don't want for analytics.
     // **Do not** ignore fields with `..` or `_` here, because we want to add them in the future.
     let meilisearch_types::features::RuntimeTogglableFeatures {
+        score_details,
         vector_store,
         metrics,
         logs_route,
@@ -87,6 +91,7 @@ async fn patch_features(
     analytics.publish(
         "Experimental features Updated".to_string(),
         json!({
+            "score_details": score_details,
             "vector_store": vector_store,
             "metrics": metrics,
             "logs_route": logs_route,

meilisearch/src/routes/logs.rs

@@ -22,23 +22,21 @@ use crate::error::MeilisearchHttpError;
 use crate::extractors::authentication::policies::*;
 use crate::extractors::authentication::GuardedData;
 use crate::extractors::sequential_extractor::SeqHandler;
-use crate::{LogRouteHandle, LogStderrHandle};
+use crate::LogRouteHandle;
 
 pub fn configure(cfg: &mut web::ServiceConfig) {
     cfg.service(
         web::resource("stream")
             .route(web::post().to(SeqHandler(get_logs)))
             .route(web::delete().to(SeqHandler(cancel_logs))),
-    )
-    .service(web::resource("stderr").route(web::post().to(SeqHandler(update_stderr_target))));
+    );
 }
 
 #[derive(Debug, Default, Clone, Copy, Deserr, PartialEq, Eq)]
-#[deserr(rename_all = camelCase)]
+#[deserr(rename_all = lowercase)]
 pub enum LogMode {
     #[default]
-    Human,
-    Json,
+    Fmt,
     Profile,
 }
 
@@ -52,7 +50,7 @@ enum MyParseError {
     #[error(transparent)]
     ParseError(#[from] tracing_subscriber::filter::ParseError),
     #[error(
-        "Empty string is not a valid target. If you want to get no logs use `OFF`. Usage: `info`, `meilisearch=info`, or you can write multiple filters in one target: `index_scheduler=info,milli=trace`"
+        "Empty string is not a valid target. If you want to get no logs use `OFF`. Usage: `info`, `info:meilisearch`, or you can write multiple filters in one target: `index_scheduler=info,milli=trace`"
     )]
     Example,
 }
@@ -162,23 +160,12 @@ fn make_layer<
 ) -> (Box<dyn Layer<S> + Send + Sync>, PinnedByteStream) {
     let guard = HandleGuard { logs: logs.into_inner() };
     match opt.mode {
-        LogMode::Human => {
+        LogMode::Fmt => {
             let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
 
             let fmt_layer = tracing_subscriber::fmt::layer()
                 .with_writer(move || LogWriter { sender: sender.clone() })
-                .with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE);
-
-            let stream = byte_stream(receiver, guard);
-            (Box::new(fmt_layer) as Box<dyn Layer<S> + Send + Sync>, Box::pin(stream))
-        }
-        LogMode::Json => {
-            let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
-
-            let fmt_layer = tracing_subscriber::fmt::layer()
-                .with_writer(move || LogWriter { sender: sender.clone() })
-                .json()
-                .with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE);
+                .with_span_events(tracing_subscriber::fmt::format::FmtSpan::ACTIVE);
 
             let stream = byte_stream(receiver, guard);
             (Box::new(fmt_layer) as Box<dyn Layer<S> + Send + Sync>, Box::pin(stream))
@@ -292,27 +279,3 @@ pub async fn cancel_logs(
 
     Ok(HttpResponse::NoContent().finish())
 }
-
-#[derive(Debug, Deserr)]
-#[deserr(error = DeserrJsonError, rename_all = camelCase, deny_unknown_fields)]
-pub struct UpdateStderrLogs {
-    #[deserr(default = "info".parse().unwrap(), try_from(&String) = MyTargets::from_str -> DeserrJsonError<BadRequest>)]
-    target: MyTargets,
-}
-
-pub async fn update_stderr_target(
-    index_scheduler: GuardedData<ActionPolicy<{ actions::METRICS_GET }>, Data<IndexScheduler>>,
-    logs: Data<LogStderrHandle>,
-    body: AwebJson<UpdateStderrLogs, DeserrJsonError>,
-) -> Result<HttpResponse, ResponseError> {
-    index_scheduler.features().check_logs_route()?;
-
-    let opt = body.into_inner();
-
-    logs.modify(|layer| {
-        *layer.filter_mut() = opt.target.0.clone();
-    })
-    .unwrap();
-
-    Ok(HttpResponse::NoContent().finish())
-}

meilisearch/src/search.rs

@@ -441,6 +441,10 @@ fn prepare_search<'t>(
         ScoringStrategy::Skip
     });
 
+    if query.show_ranking_score_details {
+        features.check_score_details()?;
+    }
+
     if let Some(HybridQuery { embedder: Some(embedder), .. }) = &query.hybrid {
         search.embedder_name(embedder);
     }

meilisearch/tests/common/server.rs

@@ -9,7 +9,7 @@ use actix_web::http::StatusCode;
 use byte_unit::{Byte, ByteUnit};
 use clap::Parser;
 use meilisearch::option::{IndexerOpts, MaxMemory, Opt};
-use meilisearch::{analytics, create_app, setup_meilisearch, SubscriberForSecondLayer};
+use meilisearch::{analytics, create_app, setup_meilisearch};
 use once_cell::sync::Lazy;
 use tempfile::TempDir;
 use tokio::time::sleep;
@@ -87,20 +87,12 @@ impl Server {
             tracing_subscriber::reload::Layer::new(None.with_filter(
                 tracing_subscriber::filter::Targets::new().with_target("", LevelFilter::OFF),
             ));
-        let (_stderr_layer, stderr_layer_handle) = tracing_subscriber::reload::Layer::new(
-            (Box::new(
-                tracing_subscriber::fmt::layer()
-                    .with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE),
-            )
-                as Box<dyn tracing_subscriber::Layer<SubscriberForSecondLayer> + Send + Sync>)
-                .with_filter(tracing_subscriber::filter::Targets::new()),
-        );
 
         actix_web::test::init_service(create_app(
             self.service.index_scheduler.clone().into(),
             self.service.auth.clone().into(),
             self.service.options.clone(),
-            (route_layer_handle, stderr_layer_handle),
+            route_layer_handle,
             analytics::MockAnalytics::new(&self.service.options),
             true,
         ))

meilisearch/tests/common/service.rs

@@ -5,7 +5,7 @@ use actix_web::http::StatusCode;
 use actix_web::test;
 use actix_web::test::TestRequest;
 use index_scheduler::IndexScheduler;
-use meilisearch::{analytics, create_app, Opt, SubscriberForSecondLayer};
+use meilisearch::{analytics, create_app, Opt};
 use meilisearch_auth::AuthController;
 use tracing::level_filters::LevelFilter;
 use tracing_subscriber::Layer;
@@ -111,20 +111,12 @@ impl Service {
             tracing_subscriber::reload::Layer::new(None.with_filter(
                 tracing_subscriber::filter::Targets::new().with_target("", LevelFilter::OFF),
             ));
-        let (_stderr_layer, stderr_layer_handle) = tracing_subscriber::reload::Layer::new(
-            (Box::new(
-                tracing_subscriber::fmt::layer()
-                    .with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE),
-            )
-                as Box<dyn tracing_subscriber::Layer<SubscriberForSecondLayer> + Send + Sync>)
-                .with_filter(tracing_subscriber::filter::Targets::new()),
-        );
 
         let app = test::init_service(create_app(
             self.index_scheduler.clone().into(),
             self.auth.clone().into(),
             self.options.clone(),
-            (route_layer_handle, stderr_layer_handle),
+            route_layer_handle,
             analytics::MockAnalytics::new(&self.options),
             true,
         ))

meilisearch/tests/dumps/mod.rs

@@ -1845,6 +1845,7 @@ async fn import_dump_v6_containing_experimental_features() {
     meili_snap::snapshot!(code, @"200 OK");
     meili_snap::snapshot!(meili_snap::json_string!(response), @r###"
     {
+      "scoreDetails": false,
       "vectorStore": false,
       "metrics": false,
       "logsRoute": false,

meilisearch/tests/features/mod.rs

@@ -18,6 +18,7 @@ async fn experimental_features() {
     meili_snap::snapshot!(code, @"200 OK");
     meili_snap::snapshot!(meili_snap::json_string!(response), @r###"
     {
+      "scoreDetails": false,
       "vectorStore": false,
       "metrics": false,
       "logsRoute": false,
@@ -30,6 +31,7 @@ async fn experimental_features() {
     meili_snap::snapshot!(code, @"200 OK");
     meili_snap::snapshot!(meili_snap::json_string!(response), @r###"
     {
+      "scoreDetails": false,
       "vectorStore": true,
       "metrics": false,
       "logsRoute": false,
@@ -42,6 +44,7 @@ async fn experimental_features() {
     meili_snap::snapshot!(code, @"200 OK");
     meili_snap::snapshot!(meili_snap::json_string!(response), @r###"
     {
+      "scoreDetails": false,
       "vectorStore": true,
       "metrics": false,
       "logsRoute": false,
@@ -55,6 +58,7 @@ async fn experimental_features() {
     meili_snap::snapshot!(code, @"200 OK");
     meili_snap::snapshot!(meili_snap::json_string!(response), @r###"
     {
+      "scoreDetails": false,
       "vectorStore": true,
       "metrics": false,
       "logsRoute": false,
@@ -68,6 +72,7 @@ async fn experimental_features() {
     meili_snap::snapshot!(code, @"200 OK");
     meili_snap::snapshot!(meili_snap::json_string!(response), @r###"
     {
+      "scoreDetails": false,
       "vectorStore": true,
       "metrics": false,
       "logsRoute": false,
@@ -88,6 +93,7 @@ async fn experimental_feature_metrics() {
     meili_snap::snapshot!(code, @"200 OK");
     meili_snap::snapshot!(meili_snap::json_string!(response), @r###"
     {
+      "scoreDetails": false,
       "vectorStore": false,
       "metrics": true,
       "logsRoute": false,
@@ -146,7 +152,7 @@ async fn errors() {
     meili_snap::snapshot!(code, @"400 Bad Request");
     meili_snap::snapshot!(meili_snap::json_string!(response), @r###"
     {
-      "message": "Unknown field `NotAFeature`: expected one of `vectorStore`, `metrics`, `logsRoute`, `exportPuffinReports`",
+      "message": "Unknown field `NotAFeature`: expected one of `scoreDetails`, `vectorStore`, `metrics`, `logsRoute`, `exportPuffinReports`",
       "code": "bad_request",
       "type": "invalid_request",
       "link": "https://docs.meilisearch.com/errors#bad_request"

meilisearch/tests/logs/error.rs

@@ -36,7 +36,7 @@ async fn logs_stream_bad_target() {
     snapshot!(code, @"400 Bad Request");
     snapshot!(response, @r###"
     {
-      "message": "Invalid value at `.target`: Empty string is not a valid target. If you want to get no logs use `OFF`. Usage: `info`, `meilisearch=info`, or you can write multiple filters in one target: `index_scheduler=info,milli=trace`",
+      "message": "Invalid value at `.target`: Empty string is not a valid target. If you want to get no logs use `OFF`. Usage: `info`, `info:meilisearch`, or you can write multiple filters in one target: `index_scheduler=info,milli=trace`",
       "code": "bad_request",
       "type": "invalid_request",
       "link": "https://docs.meilisearch.com/errors#bad_request"
@@ -89,7 +89,7 @@ async fn logs_stream_bad_mode() {
     snapshot!(code, @"400 Bad Request");
     snapshot!(response, @r###"
     {
-      "message": "Unknown value `tamo` at `.mode`: expected one of `human`, `json`, `profile`",
+      "message": "Unknown value `tamo` at `.mode`: expected one of `fmt`, `profile`",
       "code": "bad_request",
       "type": "invalid_request",
       "link": "https://docs.meilisearch.com/errors#bad_request"
@@ -133,7 +133,7 @@ async fn logs_stream_bad_profile_memory() {
     snapshot!(code, @"400 Bad Request");
     snapshot!(response, @r###"
     {
-      "message": "Invalid value: `profile_memory` can only be used while profiling code and is not compatible with the Human mode.",
+      "message": "Invalid value: `profile_memory` can only be used while profiling code and is not compatible with the Fmt mode.",
       "code": "invalid_settings_typo_tolerance",
       "type": "invalid_request",
       "link": "https://docs.meilisearch.com/errors#invalid_settings_typo_tolerance"
@@ -146,10 +146,10 @@ async fn logs_stream_bad_profile_memory() {
     snapshot!(code, @"400 Bad Request");
     snapshot!(response, @r###"
     {
-      "message": "Unknown value `fmt` at `.mode`: expected one of `human`, `json`, `profile`",
-      "code": "bad_request",
+      "message": "Invalid value: `profile_memory` can only be used while profiling code and is not compatible with the Fmt mode.",
+      "code": "invalid_settings_typo_tolerance",
       "type": "invalid_request",
-      "link": "https://docs.meilisearch.com/errors#bad_request"
+      "link": "https://docs.meilisearch.com/errors#invalid_settings_typo_tolerance"
     }
     "###);
 }
@@ -162,7 +162,7 @@ async fn logs_stream_without_enabling_the_route() {
     snapshot!(code, @"400 Bad Request");
     snapshot!(response, @r###"
     {
-      "message": "Modifying logs through the `/logs/*` routes requires enabling the `logs route` experimental feature. See https://github.com/orgs/meilisearch/discussions/721",
+      "message": "getting logs through the `/logs/stream` route requires enabling the `logs route` experimental feature. See https://github.com/orgs/meilisearch/discussions/721",
       "code": "feature_not_enabled",
       "type": "invalid_request",
       "link": "https://docs.meilisearch.com/errors#feature_not_enabled"
@@ -173,18 +173,7 @@ async fn logs_stream_without_enabling_the_route() {
     snapshot!(code, @"400 Bad Request");
     snapshot!(response, @r###"
     {
-      "message": "Modifying logs through the `/logs/*` routes requires enabling the `logs route` experimental feature. See https://github.com/orgs/meilisearch/discussions/721",
-      "code": "feature_not_enabled",
-      "type": "invalid_request",
-      "link": "https://docs.meilisearch.com/errors#feature_not_enabled"
-    }
-    "###);
-
-    let (response, code) = server.service.post("/logs/stderr", json!({})).await;
-    snapshot!(code, @"400 Bad Request");
-    snapshot!(response, @r###"
-    {
-      "message": "Modifying logs through the `/logs/*` routes requires enabling the `logs route` experimental feature. See https://github.com/orgs/meilisearch/discussions/721",
+      "message": "getting logs through the `/logs/stream` route requires enabling the `logs route` experimental feature. See https://github.com/orgs/meilisearch/discussions/721",
       "code": "feature_not_enabled",
       "type": "invalid_request",
       "link": "https://docs.meilisearch.com/errors#feature_not_enabled"

meilisearch/tests/logs/mod.rs

@@ -5,7 +5,7 @@ use std::str::FromStr;
 
 use actix_web::http::header::ContentType;
 use meili_snap::snapshot;
-use meilisearch::{analytics, create_app, Opt, SubscriberForSecondLayer};
+use meilisearch::{analytics, create_app, Opt};
 use tracing::level_filters::LevelFilter;
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::Layer;
@@ -27,25 +27,18 @@ async fn basic_test_log_stream_route() {
         tracing_subscriber::reload::Layer::new(None.with_filter(
             tracing_subscriber::filter::Targets::new().with_target("", LevelFilter::OFF),
         ));
-    let (_stderr_layer, stderr_layer_handle) = tracing_subscriber::reload::Layer::new(
-        (Box::new(
-            tracing_subscriber::fmt::layer()
-                .with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE),
-        ) as Box<dyn tracing_subscriber::Layer<SubscriberForSecondLayer> + Send + Sync>)
-            .with_filter(tracing_subscriber::filter::Targets::new()),
-    );
 
     let subscriber = tracing_subscriber::registry().with(route_layer).with(
         tracing_subscriber::fmt::layer()
             .with_span_events(tracing_subscriber::fmt::format::FmtSpan::ACTIVE)
-            .with_filter(tracing_subscriber::filter::LevelFilter::from_str("OFF").unwrap()),
+            .with_filter(tracing_subscriber::filter::LevelFilter::from_str("INFO").unwrap()),
     );
 
     let app = actix_web::test::init_service(create_app(
         server.service.index_scheduler.clone().into(),
         server.service.auth.clone().into(),
         server.service.options.clone(),
-        (route_layer_handle, stderr_layer_handle),
+        route_layer_handle,
         analytics::MockAnalytics::new(&server.service.options),
         true,
     ))
@@ -64,7 +57,7 @@ async fn basic_test_log_stream_route() {
             .insert_header(ContentType::json())
             .set_payload(
                 serde_json::to_vec(&json!({
-                    "mode": "human",
+                    "mode": "fmt",
                     "target": "info",
                 }))
                 .unwrap(),

meilisearch/tests/search/hybrid.rs

@@ -13,6 +13,7 @@ async fn index_with_documents<'a>(server: &'a Server, documents: &Value) -> Inde
     meili_snap::snapshot!(code, @"200 OK");
     meili_snap::snapshot!(meili_snap::json_string!(response), @r###"
     {
+      "scoreDetails": false,
       "vectorStore": true,
       "metrics": false,
       "logsRoute": false,
@@ -87,52 +88,6 @@ async fn simple_search() {
     snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_semanticScore":0.99029034},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_semanticScore":0.97434163},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_semanticScore":0.9472136}]"###);
 }
 
-#[actix_rt::test]
-async fn highlighter() {
-    let server = Server::new().await;
-    let index = index_with_documents(&server, &SIMPLE_SEARCH_DOCUMENTS).await;
-
-    let (response, code) = index
-        .search_post(json!({"q": "Captain Marvel", "vector": [1.0, 1.0],
-            "hybrid": {"semanticRatio": 0.2},
-            "attributesToHighlight": [
-                     "desc"
-                   ],
-                   "highlightPreTag": "**BEGIN**",
-                   "highlightPostTag": "**END**"
-        }))
-        .await;
-    snapshot!(code, @"200 OK");
-    snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_formatted":{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":["2.0","3.0"]}}},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_formatted":{"title":"Shazam!","desc":"a **BEGIN**Captain**END** **BEGIN**Marvel**END** ersatz","id":"1","_vectors":{"default":["1.0","3.0"]}}},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_formatted":{"title":"Captain Planet","desc":"He's not part of the **BEGIN**Marvel**END** Cinematic Universe","id":"2","_vectors":{"default":["1.0","2.0"]}}}]"###);
-
-    let (response, code) = index
-        .search_post(json!({"q": "Captain Marvel", "vector": [1.0, 1.0],
-            "hybrid": {"semanticRatio": 0.8},
-            "attributesToHighlight": [
-                     "desc"
-                   ],
-                   "highlightPreTag": "**BEGIN**",
-                   "highlightPostTag": "**END**"
-        }))
-        .await;
-    snapshot!(code, @"200 OK");
-    snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_formatted":{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":["2.0","3.0"]}},"_semanticScore":0.99029034},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_formatted":{"title":"Captain Planet","desc":"He's not part of the **BEGIN**Marvel**END** Cinematic Universe","id":"2","_vectors":{"default":["1.0","2.0"]}},"_semanticScore":0.97434163},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_formatted":{"title":"Shazam!","desc":"a **BEGIN**Captain**END** **BEGIN**Marvel**END** ersatz","id":"1","_vectors":{"default":["1.0","3.0"]}},"_semanticScore":0.9472136}]"###);
-
-    // no highlighting on full semantic
-    let (response, code) = index
-        .search_post(json!({"q": "Captain Marvel", "vector": [1.0, 1.0],
-            "hybrid": {"semanticRatio": 1.0},
-            "attributesToHighlight": [
-                     "desc"
-                   ],
-                   "highlightPreTag": "**BEGIN**",
-                   "highlightPostTag": "**END**"
-        }))
-        .await;
-    snapshot!(code, @"200 OK");
-    snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_formatted":{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":["2.0","3.0"]}},"_semanticScore":0.99029034},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_formatted":{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":["1.0","2.0"]}},"_semanticScore":0.97434163},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_formatted":{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":["1.0","3.0"]}}}]"###);
-}
-
 #[actix_rt::test]
 async fn invalid_semantic_ratio() {
     let server = Server::new().await;

meilisearch/tests/search/mod.rs

@@ -766,14 +766,38 @@ async fn faceting_max_values_per_facet() {
 }
 
 #[actix_rt::test]
-async fn test_score_details() {
+async fn experimental_feature_score_details() {
     let server = Server::new().await;
     let index = server.index("test");
 
     let documents = DOCUMENTS.clone();
 
-    let res = index.add_documents(json!(documents), None).await;
-    index.wait_task(res.0.uid()).await;
+    index.add_documents(json!(documents), None).await;
+    index.wait_task(0).await;
+
+    index
+        .search(
+            json!({
+                "q": "train dragon",
+                "showRankingScoreDetails": true,
+            }),
+            |response, code| {
+                meili_snap::snapshot!(code, @"400 Bad Request");
+                meili_snap::snapshot!(meili_snap::json_string!(response), @r###"
+                {
+                  "message": "Computing score details requires enabling the `score details` experimental feature. See https://github.com/meilisearch/product/discussions/674",
+                  "code": "feature_not_enabled",
+                  "type": "invalid_request",
+                  "link": "https://docs.meilisearch.com/errors#feature_not_enabled"
+                }
+                "###);
+            },
+        )
+        .await;
+
+    let (response, code) = server.set_features(json!({"scoreDetails": true})).await;
+    meili_snap::snapshot!(code, @"200 OK");
+    meili_snap::snapshot!(response["scoreDetails"], @"true");
 
     index
         .search(

milli/Cargo.toml

@@ -17,7 +17,7 @@ bincode = "1.3.3"
 bstr = "1.9.0"
 bytemuck = { version = "1.14.0", features = ["extern_crate_alloc"] }
 byteorder = "1.5.0"
-charabia = { version = "0.8.7", default-features = false }
+charabia = { version = "0.8.5", default-features = false }
 concat-arrays = "0.1.2"
 crossbeam-channel = "0.5.11"
 deserr = "0.6.1"
@@ -102,16 +102,7 @@ meili-snap = { path = "../meili-snap" }
 rand = { version = "0.8.5", features = ["small_rng"] }
 
 [features]
-all-tokenizations = [
-    "charabia/chinese",
-    "charabia/hebrew",
-    "charabia/japanese",
-    "charabia/thai",
-    "charabia/korean",
-    "charabia/greek",
-    "charabia/khmer",
-    "charabia/vietnamese",
-]
+all-tokenizations = ["charabia/chinese", "charabia/hebrew", "charabia/japanese", "charabia/thai", "charabia/korean", "charabia/greek", "charabia/khmer"]
 
 # Use POSIX semaphores instead of SysV semaphores in LMDB
 # For more information on this feature, see heed's Cargo.toml
@@ -139,7 +130,5 @@ greek = ["charabia/greek"]
 # allow khmer specialized tokenization
 khmer = ["charabia/khmer"]
 
-vietnamese = ["charabia/vietnamese"]
-
 # allow CUDA support, see <https://github.com/meilisearch/meilisearch/issues/4306>
 cuda = ["candle-core/cuda"]

milli/src/error.rs

@@ -227,22 +227,6 @@ only composed of alphanumeric characters (a-z A-Z 0-9), hyphens (-) and undersco
         source_: crate::vector::settings::EmbedderSource,
         embedder_name: String,
     },
-    #[error("`.embedders.{embedder_name}.dimensions`: Model `{model}` does not support overriding its native dimensions of {expected_dimensions}. Found {dimensions}")]
-    InvalidOpenAiModelDimensions {
-        embedder_name: String,
-        model: &'static str,
-        dimensions: usize,
-        expected_dimensions: usize,
-    },
-    #[error("`.embedders.{embedder_name}.dimensions`: Model `{model}` does not support overriding its dimensions to a value higher than {max_dimensions}. Found {dimensions}")]
-    InvalidOpenAiModelDimensionsMax {
-        embedder_name: String,
-        model: &'static str,
-        dimensions: usize,
-        max_dimensions: usize,
-    },
-    #[error("`.embedders.{embedder_name}.dimensions`: `dimensions` cannot be zero")]
-    InvalidSettingsDimensions { embedder_name: String },
 }
 
 impl From<crate::vector::Error> for Error {

milli/src/search/hybrid.rs

@@ -102,7 +102,7 @@ impl ScoreWithRatioResult {
         }
 
         SearchResult {
-            matching_words: right.matching_words,
+            matching_words: left.matching_words,
             candidates: left.candidates | right.candidates,
             documents_ids,
             document_scores,

milli/src/update/facet/bulk.rs

@@ -1,7 +1,7 @@
 use std::fs::File;
 use std::io::BufReader;
 
-use grenad::{CompressionType, Merger};
+use grenad::CompressionType;
 use heed::types::Bytes;
 use heed::{BytesDecode, BytesEncode, Error, PutFlags, RoTxn, RwTxn};
 use roaring::RoaringBitmap;
@@ -14,7 +14,6 @@ use crate::heed_codec::facet::{
 use crate::heed_codec::BytesRefCodec;
 use crate::update::del_add::{DelAdd, KvReaderDelAdd};
 use crate::update::index_documents::{create_writer, valid_lmdb_key, writer_into_reader};
-use crate::update::MergeFn;
 use crate::{CboRoaringBitmapCodec, CboRoaringBitmapLenCodec, FieldId, Index, Result};
 
 /// Algorithm to insert elememts into the `facet_id_(string/f64)_docids` databases
@@ -29,7 +28,7 @@ pub struct FacetsUpdateBulk<'i> {
     facet_type: FacetType,
     field_ids: Vec<FieldId>,
     // None if level 0 does not need to be updated
-    delta_data: Option<Merger<BufReader<File>, MergeFn>>,
+    delta_data: Option<grenad::Reader<BufReader<File>>>,
 }
 
 impl<'i> FacetsUpdateBulk<'i> {
@@ -37,7 +36,7 @@ impl<'i> FacetsUpdateBulk<'i> {
         index: &'i Index,
         field_ids: Vec<FieldId>,
         facet_type: FacetType,
-        delta_data: Merger<BufReader<File>, MergeFn>,
+        delta_data: grenad::Reader<BufReader<File>>,
         group_size: u8,
         min_level_size: u8,
     ) -> FacetsUpdateBulk<'i> {
@@ -66,7 +65,7 @@ impl<'i> FacetsUpdateBulk<'i> {
         }
     }
 
-    #[tracing::instrument(level = "trace", skip_all, target = "indexing::facets::bulk")]
+    #[logging_timer::time("FacetsUpdateBulk::{}")]
     pub fn execute(self, wtxn: &mut heed::RwTxn) -> Result<()> {
         let Self { index, field_ids, group_size, min_level_size, facet_type, delta_data } = self;
 
@@ -90,7 +89,7 @@ impl<'i> FacetsUpdateBulk<'i> {
 /// Implementation of `FacetsUpdateBulk` that is independent of milli's `Index` type
 pub(crate) struct FacetsUpdateBulkInner<R: std::io::Read + std::io::Seek> {
     pub db: heed::Database<FacetGroupKeyCodec<BytesRefCodec>, FacetGroupValueCodec>,
-    pub delta_data: Option<Merger<R, MergeFn>>,
+    pub delta_data: Option<grenad::Reader<R>>,
     pub group_size: u8,
     pub min_level_size: u8,
 }
@@ -130,8 +129,8 @@ impl<R: std::io::Read + std::io::Seek> FacetsUpdateBulkInner<R> {
         if self.db.is_empty(wtxn)? {
             let mut buffer = Vec::new();
             let mut database = self.db.iter_mut(wtxn)?.remap_types::<Bytes, Bytes>();
-            let mut iter = delta_data.into_stream_merger_iter()?;
-            while let Some((key, value)) = iter.next()? {
+            let mut cursor = delta_data.into_cursor()?;
+            while let Some((key, value)) = cursor.move_on_next()? {
                 if !valid_lmdb_key(key) {
                     continue;
                 }
@@ -155,8 +154,8 @@ impl<R: std::io::Read + std::io::Seek> FacetsUpdateBulkInner<R> {
             let mut buffer = Vec::new();
             let database = self.db.remap_types::<Bytes, Bytes>();
 
-            let mut iter = delta_data.into_stream_merger_iter()?;
-            while let Some((key, value)) = iter.next()? {
+            let mut cursor = delta_data.into_cursor()?;
+            while let Some((key, value)) = cursor.move_on_next()? {
                 if !valid_lmdb_key(key) {
                     continue;
                 }

milli/src/update/facet/incremental.rs

@@ -1,7 +1,6 @@
 use std::fs::File;
 use std::io::BufReader;
 
-use grenad::Merger;
 use heed::types::{Bytes, DecodeIgnore};
 use heed::{BytesDecode, Error, RoTxn, RwTxn};
 use obkv::KvReader;
@@ -15,7 +14,6 @@ use crate::heed_codec::BytesRefCodec;
 use crate::search::facet::get_highest_level;
 use crate::update::del_add::DelAdd;
 use crate::update::index_documents::valid_lmdb_key;
-use crate::update::MergeFn;
 use crate::{CboRoaringBitmapCodec, Index, Result};
 
 enum InsertionResult {
@@ -33,14 +31,14 @@ enum DeletionResult {
 /// `facet_id_(string/f64)_docids` databases.
 pub struct FacetsUpdateIncremental {
     inner: FacetsUpdateIncrementalInner,
-    delta_data: Merger<BufReader<File>, MergeFn>,
+    delta_data: grenad::Reader<BufReader<File>>,
 }
 
 impl FacetsUpdateIncremental {
     pub fn new(
         index: &Index,
         facet_type: FacetType,
-        delta_data: Merger<BufReader<File>, MergeFn>,
+        delta_data: grenad::Reader<BufReader<File>>,
         group_size: u8,
         min_level_size: u8,
         max_group_size: u8,
@@ -63,18 +61,16 @@ impl FacetsUpdateIncremental {
         }
     }
 
-    #[tracing::instrument(level = "trace", skip_all, target = "indexing::facets::incremental")]
     pub fn execute(self, wtxn: &mut RwTxn) -> crate::Result<()> {
-        let mut iter = self.delta_data.into_stream_merger_iter()?;
-
-        while let Some((key, value)) = iter.next()? {
+        let mut cursor = self.delta_data.into_cursor()?;
+        while let Some((key, value)) = cursor.move_on_next()? {
             if !valid_lmdb_key(key) {
                 continue;
             }
-
             let key = FacetGroupKeyCodec::<BytesRefCodec>::bytes_decode(key)
                 .map_err(heed::Error::Encoding)?;
             let value = KvReader::new(value);
+
             let docids_to_delete = value
                 .get(DelAdd::Deletion)
                 .map(CboRoaringBitmapCodec::bytes_decode)

milli/src/update/facet/mod.rs

@@ -79,9 +79,12 @@ pub const FACET_MIN_LEVEL_SIZE: u8 = 5;
 use std::collections::BTreeSet;
 use std::fs::File;
 use std::io::BufReader;
+use std::iter::FromIterator;
 
-use grenad::Merger;
-use heed::types::{Bytes, DecodeIgnore};
+use charabia::normalizer::{Normalize, NormalizerOption};
+use grenad::{CompressionType, SortAlgorithm};
+use heed::types::{Bytes, DecodeIgnore, SerdeJson};
+use heed::BytesEncode;
 use time::OffsetDateTime;
 use tracing::debug;
 
@@ -90,9 +93,9 @@ use super::FacetsUpdateBulk;
 use crate::facet::FacetType;
 use crate::heed_codec::facet::{FacetGroupKey, FacetGroupKeyCodec, FacetGroupValueCodec};
 use crate::heed_codec::BytesRefCodec;
-use crate::update::del_add::{DelAdd, KvReaderDelAdd};
-use crate::update::MergeFn;
-use crate::{try_split_array_at, FieldId, Index, Result};
+use crate::update::index_documents::create_sorter;
+use crate::update::merge_btreeset_string;
+use crate::{BEU16StrCodec, Index, Result, MAX_FACET_VALUE_LENGTH};
 
 pub mod bulk;
 pub mod incremental;
@@ -105,20 +108,16 @@ pub struct FacetsUpdate<'i> {
     index: &'i Index,
     database: heed::Database<FacetGroupKeyCodec<BytesRefCodec>, FacetGroupValueCodec>,
     facet_type: FacetType,
-    delta_data: Merger<BufReader<File>, MergeFn>,
-    normalized_delta_data: Option<Merger<BufReader<File>, MergeFn>>,
+    delta_data: grenad::Reader<BufReader<File>>,
     group_size: u8,
     max_group_size: u8,
     min_level_size: u8,
-    data_size: u64,
 }
 impl<'i> FacetsUpdate<'i> {
     pub fn new(
         index: &'i Index,
         facet_type: FacetType,
-        delta_data: Merger<BufReader<File>, MergeFn>,
-        normalized_delta_data: Option<Merger<BufReader<File>, MergeFn>>,
-        data_size: u64,
+        delta_data: grenad::Reader<BufReader<File>>,
     ) -> Self {
         let database = match facet_type {
             FacetType::String => {
@@ -136,20 +135,18 @@ impl<'i> FacetsUpdate<'i> {
             min_level_size: FACET_MIN_LEVEL_SIZE,
             facet_type,
             delta_data,
-            normalized_delta_data,
-            data_size,
         }
     }
 
     pub fn execute(self, wtxn: &mut heed::RwTxn) -> Result<()> {
-        if self.data_size == 0 {
+        if self.delta_data.is_empty() {
             return Ok(());
         }
         debug!("Computing and writing the facet values levels docids into LMDB on disk...");
         self.index.set_updated_at(wtxn, &OffsetDateTime::now_utc())?;
 
         // See self::comparison_bench::benchmark_facet_indexing
-        if self.data_size >= (self.database.len(wtxn)? / 50) {
+        if self.delta_data.len() >= (self.database.len(wtxn)? / 50) {
             let field_ids =
                 self.index.faceted_fields_ids(wtxn)?.iter().copied().collect::<Vec<_>>();
             let bulk_update = FacetsUpdateBulk::new(
@@ -173,108 +170,94 @@ impl<'i> FacetsUpdate<'i> {
             incremental_update.execute(wtxn)?;
         }
 
-        match self.normalized_delta_data {
-            Some(data) => index_facet_search(wtxn, data, self.index),
-            None => Ok(()),
-        }
-    }
-}
+        // We clear the list of normalized-for-search facets
+        // and the previous FSTs to compute everything from scratch
+        self.index.facet_id_normalized_string_strings.clear(wtxn)?;
+        self.index.facet_id_string_fst.clear(wtxn)?;
 
-fn index_facet_search(
-    wtxn: &mut heed::RwTxn,
-    normalized_delta_data: Merger<BufReader<File>, MergeFn>,
-    index: &Index,
-) -> Result<()> {
-    let mut iter = normalized_delta_data.into_stream_merger_iter()?;
-    while let Some((key_bytes, delta_bytes)) = iter.next()? {
-        let deladd_reader = KvReaderDelAdd::new(delta_bytes);
+        // As we can't use the same write transaction to read and write in two different databases
+        // we must create a temporary sorter that we will write into LMDB afterward.
+        // As multiple unnormalized facet values can become the same normalized facet value
+        // we must merge them together.
+        let mut sorter = create_sorter(
+            SortAlgorithm::Unstable,
+            merge_btreeset_string,
+            CompressionType::None,
+            None,
+            None,
+            None,
+        );
 
-        let database_set = index
-            .facet_id_normalized_string_strings
-            .remap_key_type::<Bytes>()
-            .get(wtxn, key_bytes)?
-            .unwrap_or_default();
-
-        let add_set = deladd_reader
-            .get(DelAdd::Addition)
-            .and_then(|bytes| serde_json::from_slice::<BTreeSet<String>>(bytes).ok())
-            .unwrap_or_default();
-
-        let del_set = match deladd_reader
-            .get(DelAdd::Deletion)
-            .and_then(|bytes| serde_json::from_slice::<BTreeSet<String>>(bytes).ok())
-        {
-            Some(del_set) => {
-                let (field_id_bytes, _) = try_split_array_at(key_bytes).unwrap();
-                let field_id = FieldId::from_be_bytes(field_id_bytes);
-                let mut set = BTreeSet::new();
-                for facet in del_set {
-                    let key = FacetGroupKey { field_id, level: 0, left_bound: facet.as_str() };
-                    // Check if the referenced value doesn't exist anymore before deleting it.
-                    if index
-                        .facet_id_string_docids
-                        .remap_data_type::<DecodeIgnore>()
-                        .get(wtxn, &key)?
-                        .is_none()
-                    {
-                        set.insert(facet);
-                    }
+        // We iterate on the list of original, semi-normalized, facet values
+        // and normalize them for search, inserting them in LMDB in any given order.
+        let options = NormalizerOption { lossy: true, ..Default::default() };
+        let database = self.index.facet_id_string_docids.remap_data_type::<DecodeIgnore>();
+        for result in database.iter(wtxn)? {
+            let (facet_group_key, ()) = result?;
+            if let FacetGroupKey { field_id, level: 0, left_bound } = facet_group_key {
+                let mut normalized_facet = left_bound.normalize(&options);
+                let normalized_truncated_facet: String;
+                if normalized_facet.len() > MAX_FACET_VALUE_LENGTH {
+                    normalized_truncated_facet = normalized_facet
+                        .char_indices()
+                        .take_while(|(idx, _)| *idx < MAX_FACET_VALUE_LENGTH)
+                        .map(|(_, c)| c)
+                        .collect();
+                    normalized_facet = normalized_truncated_facet.into();
                 }
-                set
+                let set = BTreeSet::from_iter(std::iter::once(left_bound));
+                let key = (field_id, normalized_facet.as_ref());
+                let key = BEU16StrCodec::bytes_encode(&key).map_err(heed::Error::Encoding)?;
+                let val = SerdeJson::bytes_encode(&set).map_err(heed::Error::Encoding)?;
+                sorter.insert(key, val)?;
             }
-            None => BTreeSet::new(),
-        };
-
-        let set: BTreeSet<_> =
-            database_set.difference(&del_set).chain(add_set.iter()).cloned().collect();
-
-        if set.is_empty() {
-            index
-                .facet_id_normalized_string_strings
-                .remap_key_type::<Bytes>()
-                .delete(wtxn, key_bytes)?;
-        } else {
-            index
-                .facet_id_normalized_string_strings
-                .remap_key_type::<Bytes>()
-                .put(wtxn, key_bytes, &set)?;
         }
-    }
 
-    // We clear the FST of normalized-for-search to compute everything from scratch.
-    index.facet_id_string_fst.clear(wtxn)?;
-    // We compute one FST by string facet
-    let mut text_fsts = vec![];
-    let mut current_fst: Option<(u16, fst::SetBuilder<Vec<u8>>)> = None;
-    let database = index.facet_id_normalized_string_strings.remap_data_type::<DecodeIgnore>();
-    for result in database.iter(wtxn)? {
-        let ((field_id, normalized_facet), _) = result?;
-        current_fst = match current_fst.take() {
-            Some((fid, fst_builder)) if fid != field_id => {
-                let fst = fst_builder.into_set();
-                text_fsts.push((fid, fst));
-                Some((field_id, fst::SetBuilder::memory()))
+        // In this loop we don't need to take care of merging bitmaps
+        // as the grenad sorter already merged them for us.
+        let mut merger_iter = sorter.into_stream_merger_iter()?;
+        while let Some((key_bytes, btreeset_bytes)) = merger_iter.next()? {
+            self.index.facet_id_normalized_string_strings.remap_types::<Bytes, Bytes>().put(
+                wtxn,
+                key_bytes,
+                btreeset_bytes,
+            )?;
+        }
+
+        // We compute one FST by string facet
+        let mut text_fsts = vec![];
+        let mut current_fst: Option<(u16, fst::SetBuilder<Vec<u8>>)> = None;
+        let database =
+            self.index.facet_id_normalized_string_strings.remap_data_type::<DecodeIgnore>();
+        for result in database.iter(wtxn)? {
+            let ((field_id, normalized_facet), _) = result?;
+            current_fst = match current_fst.take() {
+                Some((fid, fst_builder)) if fid != field_id => {
+                    let fst = fst_builder.into_set();
+                    text_fsts.push((fid, fst));
+                    Some((field_id, fst::SetBuilder::memory()))
+                }
+                Some((field_id, fst_builder)) => Some((field_id, fst_builder)),
+                None => Some((field_id, fst::SetBuilder::memory())),
+            };
+
+            if let Some((_, fst_builder)) = current_fst.as_mut() {
+                fst_builder.insert(normalized_facet)?;
             }
-            Some((field_id, fst_builder)) => Some((field_id, fst_builder)),
-            None => Some((field_id, fst::SetBuilder::memory())),
-        };
-
-        if let Some((_, fst_builder)) = current_fst.as_mut() {
-            fst_builder.insert(normalized_facet)?;
         }
-    }
 
-    if let Some((field_id, fst_builder)) = current_fst {
-        let fst = fst_builder.into_set();
-        text_fsts.push((field_id, fst));
-    }
+        if let Some((field_id, fst_builder)) = current_fst {
+            let fst = fst_builder.into_set();
+            text_fsts.push((field_id, fst));
+        }
 
-    // We write those FSTs in LMDB now
-    for (field_id, fst) in text_fsts {
-        index.facet_id_string_fst.put(wtxn, &field_id, &fst)?;
-    }
+        // We write those FSTs in LMDB now
+        for (field_id, fst) in text_fsts {
+            self.index.facet_id_string_fst.put(wtxn, &field_id, &fst)?;
+        }
 
-    Ok(())
+        Ok(())
+    }
 }
 
 #[cfg(test)]
@@ -285,7 +268,6 @@ pub(crate) mod test_helpers {
     use std::marker::PhantomData;
     use std::rc::Rc;
 
-    use grenad::MergerBuilder;
     use heed::types::Bytes;
     use heed::{BytesDecode, BytesEncode, Env, RoTxn, RwTxn};
     use roaring::RoaringBitmap;
@@ -298,8 +280,7 @@ pub(crate) mod test_helpers {
     use crate::search::facet::get_highest_level;
     use crate::snapshot_tests::display_bitmap;
     use crate::update::del_add::{DelAdd, KvWriterDelAdd};
-    use crate::update::index_documents::merge_deladd_cbo_roaring_bitmaps;
-    use crate::update::{FacetsUpdateIncrementalInner, MergeFn};
+    use crate::update::FacetsUpdateIncrementalInner;
     use crate::CboRoaringBitmapCodec;
 
     /// Utility function to generate a string whose position in a lexicographically
@@ -482,13 +463,10 @@ pub(crate) mod test_helpers {
             }
             writer.finish().unwrap();
             let reader = grenad::Reader::new(std::io::Cursor::new(new_data)).unwrap();
-            let mut builder = MergerBuilder::new(merge_deladd_cbo_roaring_bitmaps as MergeFn);
-            builder.push(reader.into_cursor().unwrap());
-            let merger = builder.build();
 
             let update = FacetsUpdateBulkInner {
                 db: self.content,
-                delta_data: Some(merger),
+                delta_data: Some(reader),
                 group_size: self.group_size.get(),
                 min_level_size: self.min_level_size.get(),
             };

milli/src/update/index_documents/extract/extract_docid_word_positions.rs

@@ -26,7 +26,7 @@ pub fn extract_docid_word_positions<R: io::Read + io::Seek>(
     obkv_documents: grenad::Reader<R>,
     indexer: GrenadParameters,
     searchable_fields: &Option<HashSet<FieldId>>,
-    stop_words: Option<&fst::Set<Vec<u8>>>,
+    stop_words: Option<&fst::Set<&[u8]>>,
     allowed_separators: Option<&[&str]>,
     dictionary: Option<&[&str]>,
     max_positions_per_attributes: Option<u32>,
@@ -181,11 +181,11 @@ fn searchable_fields_changed(
 
 /// Factorize tokenizer building.
 fn tokenizer_builder<'a>(
-    stop_words: Option<&'a fst::Set<Vec<u8>>>,
+    stop_words: Option<&'a fst::Set<&[u8]>>,
     allowed_separators: Option<&'a [&str]>,
     dictionary: Option<&'a [&str]>,
     script_language: Option<&'a HashMap<Script, Vec<Language>>>,
-) -> TokenizerBuilder<'a, Vec<u8>> {
+) -> TokenizerBuilder<'a, &'a [u8]> {
     let mut tokenizer_builder = TokenizerBuilder::new();
     if let Some(stop_words) = stop_words {
         tokenizer_builder.stop_words(stop_words);
@@ -211,7 +211,7 @@ fn lang_safe_tokens_from_document<'a>(
     obkv: &KvReader<FieldId>,
     searchable_fields: &Option<HashSet<FieldId>>,
     tokenizer: &Tokenizer,
-    stop_words: Option<&fst::Set<Vec<u8>>>,
+    stop_words: Option<&fst::Set<&[u8]>>,
     allowed_separators: Option<&[&str]>,
     dictionary: Option<&[&str]>,
     max_positions_per_attributes: u32,

milli/src/update/index_documents/extract/extract_facet_string_docids.rs

@@ -1,21 +1,15 @@
-use std::collections::BTreeSet;
 use std::fs::File;
 use std::io::BufReader;
-use std::iter::FromIterator;
 use std::{io, str};
 
-use charabia::normalizer::{Normalize, NormalizerOption};
-use heed::types::SerdeJson;
 use heed::BytesEncode;
 
 use super::helpers::{create_sorter, sorter_into_reader, try_split_array_at, GrenadParameters};
 use crate::heed_codec::facet::{FacetGroupKey, FacetGroupKeyCodec};
-use crate::heed_codec::{BEU16StrCodec, StrRefCodec};
-use crate::update::del_add::{DelAdd, KvReaderDelAdd, KvWriterDelAdd};
-use crate::update::index_documents::helpers::{
-    merge_deladd_btreeset_string, merge_deladd_cbo_roaring_bitmaps,
-};
-use crate::{FieldId, Result, MAX_FACET_VALUE_LENGTH};
+use crate::heed_codec::StrRefCodec;
+use crate::update::del_add::{KvReaderDelAdd, KvWriterDelAdd};
+use crate::update::index_documents::helpers::merge_deladd_cbo_roaring_bitmaps;
+use crate::{FieldId, Result};
 
 /// Extracts the facet string and the documents ids where this facet string appear.
 ///
@@ -25,11 +19,10 @@ use crate::{FieldId, Result, MAX_FACET_VALUE_LENGTH};
 pub fn extract_facet_string_docids<R: io::Read + io::Seek>(
     docid_fid_facet_string: grenad::Reader<R>,
     indexer: GrenadParameters,
-) -> Result<(grenad::Reader<BufReader<File>>, grenad::Reader<BufReader<File>>)> {
+) -> Result<grenad::Reader<BufReader<File>>> {
     puffin::profile_function!();
 
     let max_memory = indexer.max_memory_by_thread();
-    let options = NormalizerOption { lossy: true, ..Default::default() };
 
     let mut facet_string_docids_sorter = create_sorter(
         grenad::SortAlgorithm::Stable,
@@ -37,30 +30,12 @@ pub fn extract_facet_string_docids<R: io::Read + io::Seek>(
         indexer.chunk_compression_type,
         indexer.chunk_compression_level,
         indexer.max_nb_chunks,
-        max_memory.map(|m| m / 2),
-    );
-
-    let mut normalized_facet_string_docids_sorter = create_sorter(
-        grenad::SortAlgorithm::Stable,
-        merge_deladd_btreeset_string,
-        indexer.chunk_compression_type,
-        indexer.chunk_compression_level,
-        indexer.max_nb_chunks,
-        max_memory.map(|m| m / 2),
+        max_memory,
     );
 
     let mut buffer = Vec::new();
     let mut cursor = docid_fid_facet_string.into_cursor()?;
     while let Some((key, deladd_original_value_bytes)) = cursor.move_on_next()? {
-        let deladd_reader = KvReaderDelAdd::new(deladd_original_value_bytes);
-
-        // nothing to do if we delete and re-add the value.
-        if deladd_reader.get(DelAdd::Deletion).is_some()
-            && deladd_reader.get(DelAdd::Addition).is_some()
-        {
-            continue;
-        }
-
         let (field_id_bytes, bytes) = try_split_array_at(key).unwrap();
         let field_id = FieldId::from_be_bytes(field_id_bytes);
 
@@ -69,46 +44,17 @@ pub fn extract_facet_string_docids<R: io::Read + io::Seek>(
         let document_id = u32::from_be_bytes(document_id_bytes);
 
         let normalized_value = str::from_utf8(normalized_value_bytes)?;
-
-        // Facet search normalization
-        {
-            let mut hyper_normalized_value = normalized_value.normalize(&options);
-            let normalized_truncated_facet: String;
-            if hyper_normalized_value.len() > MAX_FACET_VALUE_LENGTH {
-                normalized_truncated_facet = hyper_normalized_value
-                    .char_indices()
-                    .take_while(|(idx, _)| *idx < MAX_FACET_VALUE_LENGTH)
-                    .map(|(_, c)| c)
-                    .collect();
-                hyper_normalized_value = normalized_truncated_facet.into();
-            }
-            let set = BTreeSet::from_iter(std::iter::once(normalized_value));
-
-            buffer.clear();
-            let mut obkv = KvWriterDelAdd::new(&mut buffer);
-            for (deladd_key, _) in deladd_reader.iter() {
-                let val = SerdeJson::bytes_encode(&set).map_err(heed::Error::Encoding)?;
-                obkv.insert(deladd_key, val)?;
-            }
-            obkv.finish()?;
-
-            let key = (field_id, hyper_normalized_value.as_ref());
-            let key_bytes = BEU16StrCodec::bytes_encode(&key).map_err(heed::Error::Encoding)?;
-            normalized_facet_string_docids_sorter.insert(key_bytes, &buffer)?;
-        }
-
         let key = FacetGroupKey { field_id, level: 0, left_bound: normalized_value };
         let key_bytes = FacetGroupKeyCodec::<StrRefCodec>::bytes_encode(&key).unwrap();
 
         buffer.clear();
         let mut obkv = KvWriterDelAdd::new(&mut buffer);
-        for (deladd_key, _) in deladd_reader.iter() {
+        for (deladd_key, _) in KvReaderDelAdd::new(deladd_original_value_bytes).iter() {
             obkv.insert(deladd_key, document_id.to_ne_bytes())?;
         }
         obkv.finish()?;
         facet_string_docids_sorter.insert(&key_bytes, &buffer)?;
     }
 
-    let normalized = sorter_into_reader(normalized_facet_string_docids_sorter, indexer)?;
-    sorter_into_reader(facet_string_docids_sorter, indexer).map(|s| (s, normalized))
+    sorter_into_reader(facet_string_docids_sorter, indexer)
 }

milli/src/update/index_documents/extract/extract_vector_points.rs

@@ -257,7 +257,6 @@ fn push_vectors_diff(
     key_buffer: &mut Vec<u8>,
     delta: VectorStateDelta,
 ) -> Result<()> {
-    puffin::profile_function!();
     let (must_remove, prompt, (mut del_vectors, mut add_vectors)) = delta.into_values();
     if must_remove {
         key_buffer.truncate(TRUNCATE_SIZE);
@@ -333,15 +332,16 @@ fn extract_vectors(
     }
 }
 
-#[tracing::instrument(level = "trace", skip_all, target = "indexing::extract")]
+#[logging_timer::time]
 pub fn extract_embeddings<R: io::Read + io::Seek>(
     // docid, prompt
     prompt_reader: grenad::Reader<R>,
     indexer: GrenadParameters,
     embedder: Arc<Embedder>,
 ) -> Result<grenad::Reader<BufReader<File>>> {
-    puffin::profile_function!();
-    let n_chunks = embedder.chunk_count_hint(); // chunk level parallelism
+    let rt = tokio::runtime::Builder::new_current_thread().enable_io().enable_time().build()?;
+
+    let n_chunks = embedder.chunk_count_hint(); // chunk level parellelism
     let n_vectors_per_chunk = embedder.prompt_count_in_chunk_hint(); // number of vectors in a single chunk
 
     // docid, state with embedding
@@ -375,8 +375,11 @@ pub fn extract_embeddings<R: io::Read + io::Seek>(
         current_chunk_ids.push(docid);
 
         if chunks.len() == chunks.capacity() {
-            let chunked_embeds = embedder
-                .embed_chunks(std::mem::replace(&mut chunks, Vec::with_capacity(n_chunks)))
+            let chunked_embeds = rt
+                .block_on(
+                    embedder
+                        .embed_chunks(std::mem::replace(&mut chunks, Vec::with_capacity(n_chunks))),
+                )
                 .map_err(crate::vector::Error::from)
                 .map_err(crate::Error::from)?;
 
@@ -393,8 +396,8 @@ pub fn extract_embeddings<R: io::Read + io::Seek>(
 
     // send last chunk
     if !chunks.is_empty() {
-        let chunked_embeds = embedder
-            .embed_chunks(std::mem::take(&mut chunks))
+        let chunked_embeds = rt
+            .block_on(embedder.embed_chunks(std::mem::take(&mut chunks)))
             .map_err(crate::vector::Error::from)
             .map_err(crate::Error::from)?;
         for (docid, embeddings) in chunks_ids
@@ -407,15 +410,13 @@ pub fn extract_embeddings<R: io::Read + io::Seek>(
     }
 
     if !current_chunk.is_empty() {
-        let embeds = embedder
-            .embed_chunks(vec![std::mem::take(&mut current_chunk)])
+        let embeds = rt
+            .block_on(embedder.embed(std::mem::take(&mut current_chunk)))
             .map_err(crate::vector::Error::from)
             .map_err(crate::Error::from)?;
 
-        if let Some(embeds) = embeds.first() {
-            for (docid, embeddings) in current_chunk_ids.iter().zip(embeds.iter()) {
-                state_writer.insert(docid.to_be_bytes(), cast_slice(embeddings.as_inner()))?;
-            }
+        for (docid, embeddings) in current_chunk_ids.iter().zip(embeds.iter()) {
+            state_writer.insert(docid.to_be_bytes(), cast_slice(embeddings.as_inner()))?;
         }
     }
 

milli/src/update/index_documents/extract/mod.rs

@@ -15,6 +15,7 @@ use std::io::BufReader;
 
 use crossbeam_channel::Sender;
 use rayon::prelude::*;
+use tracing::debug;
 
 use self::extract_docid_word_positions::extract_docid_word_positions;
 use self::extract_facet_number_docids::extract_facet_number_docids;
@@ -28,7 +29,10 @@ use self::extract_vector_points::{
 use self::extract_word_docids::extract_word_docids;
 use self::extract_word_pair_proximity_docids::extract_word_pair_proximity_docids;
 use self::extract_word_position_docids::extract_word_position_docids;
-use super::helpers::{as_cloneable_grenad, CursorClonableMmap, GrenadParameters};
+use super::helpers::{
+    as_cloneable_grenad, merge_deladd_cbo_roaring_bitmaps, CursorClonableMmap, GrenadParameters,
+    MergeFn, MergeableReader,
+};
 use super::{helpers, TypedChunk};
 use crate::proximity::ProximityPrecision;
 use crate::vector::EmbeddingConfigs;
@@ -48,7 +52,7 @@ pub(crate) fn data_from_obkv_documents(
     primary_key_id: FieldId,
     geo_fields_ids: Option<(FieldId, FieldId)>,
     field_id_map: FieldsIdsMap,
-    stop_words: Option<fst::Set<Vec<u8>>>,
+    stop_words: Option<fst::Set<&[u8]>>,
     allowed_separators: Option<&[&str]>,
     dictionary: Option<&[&str]>,
     max_positions_per_attributes: Option<u32>,
@@ -58,154 +62,201 @@ pub(crate) fn data_from_obkv_documents(
 ) -> Result<()> {
     puffin::profile_function!();
 
-    let (original_pipeline_result, flattened_pipeline_result): (Result<_>, Result<_>) = rayon::join(
-        || {
-            original_obkv_chunks
-                .par_bridge()
-                .map(|original_documents_chunk| {
-                    send_original_documents_data(
-                        original_documents_chunk,
-                        indexer,
-                        lmdb_writer_sx.clone(),
-                        field_id_map.clone(),
-                        embedders.clone(),
-                    )
-                })
-                .collect::<Result<()>>()
-        },
-        || {
-            flattened_obkv_chunks
-                .par_bridge()
-                .map(|flattened_obkv_chunks| {
-                    send_and_extract_flattened_documents_data(
-                        flattened_obkv_chunks,
-                        indexer,
-                        lmdb_writer_sx.clone(),
-                        &searchable_fields,
-                        &faceted_fields,
-                        primary_key_id,
-                        geo_fields_ids,
-                        &stop_words,
-                        &allowed_separators,
-                        &dictionary,
-                        max_positions_per_attributes,
-                    )
-                })
-                .map(|result| {
-                    if let Ok((
-                        ref docid_word_positions_chunk,
-                        (ref fid_docid_facet_numbers_chunk, ref fid_docid_facet_strings_chunk),
-                    )) = result
-                    {
-                        run_extraction_task::<_, _, grenad::Reader<BufReader<File>>>(
-                            docid_word_positions_chunk.clone(),
-                            indexer,
-                            lmdb_writer_sx.clone(),
-                            extract_fid_word_count_docids,
-                            TypedChunk::FieldIdWordCountDocids,
-                            "field-id-wordcount-docids",
-                        );
+    original_obkv_chunks
+        .par_bridge()
+        .map(|original_documents_chunk| {
+            send_original_documents_data(
+                original_documents_chunk,
+                indexer,
+                lmdb_writer_sx.clone(),
+                field_id_map.clone(),
+                embedders.clone(),
+            )
+        })
+        .collect::<Result<()>>()?;
 
-                        let exact_attributes = exact_attributes.clone();
-                        run_extraction_task::<
-                            _,
-                            _,
-                            (
-                                grenad::Reader<BufReader<File>>,
-                                grenad::Reader<BufReader<File>>,
-                                grenad::Reader<BufReader<File>>,
-                            ),
-                        >(
-                            docid_word_positions_chunk.clone(),
-                            indexer,
-                            lmdb_writer_sx.clone(),
-                            move |doc_word_pos, indexer| {
-                                extract_word_docids(doc_word_pos, indexer, &exact_attributes)
-                            },
-                            |(
-                                word_docids_reader,
-                                exact_word_docids_reader,
-                                word_fid_docids_reader,
-                            )| {
-                                TypedChunk::WordDocids {
-                                    word_docids_reader,
-                                    exact_word_docids_reader,
-                                    word_fid_docids_reader,
-                                }
-                            },
-                            "word-docids",
-                        );
+    #[allow(clippy::type_complexity)]
+    let result: Result<(Vec<_>, (Vec<_>, (Vec<_>, (Vec<_>, (Vec<_>, Vec<_>)))))> =
+        flattened_obkv_chunks
+            .par_bridge()
+            .map(|flattened_obkv_chunks| {
+                send_and_extract_flattened_documents_data(
+                    flattened_obkv_chunks,
+                    indexer,
+                    lmdb_writer_sx.clone(),
+                    &searchable_fields,
+                    &faceted_fields,
+                    primary_key_id,
+                    geo_fields_ids,
+                    &stop_words,
+                    &allowed_separators,
+                    &dictionary,
+                    max_positions_per_attributes,
+                )
+            })
+            .collect();
 
-                        run_extraction_task::<_, _, grenad::Reader<BufReader<File>>>(
-                            docid_word_positions_chunk.clone(),
-                            indexer,
-                            lmdb_writer_sx.clone(),
-                            extract_word_position_docids,
-                            TypedChunk::WordPositionDocids,
-                            "word-position-docids",
-                        );
+    let (
+        docid_word_positions_chunks,
+        (
+            fid_docid_facet_numbers_chunks,
+            (
+                fid_docid_facet_strings_chunks,
+                (
+                    facet_is_null_docids_chunks,
+                    (facet_is_empty_docids_chunks, facet_exists_docids_chunks),
+                ),
+            ),
+        ),
+    ) = result?;
 
-                        run_extraction_task::<
-                            _,
-                            _,
-                            (grenad::Reader<BufReader<File>>, grenad::Reader<BufReader<File>>),
-                        >(
-                            fid_docid_facet_strings_chunk.clone(),
-                            indexer,
-                            lmdb_writer_sx.clone(),
-                            extract_facet_string_docids,
-                            TypedChunk::FieldIdFacetStringDocids,
-                            "field-id-facet-string-docids",
-                        );
+    // merge facet_exists_docids and send them as a typed chunk
+    {
+        let lmdb_writer_sx = lmdb_writer_sx.clone();
+        rayon::spawn(move || {
+            debug!(database = "facet-id-exists-docids", "merge");
+            match facet_exists_docids_chunks.merge(merge_deladd_cbo_roaring_bitmaps, &indexer) {
+                Ok(reader) => {
+                    let _ = lmdb_writer_sx.send(Ok(TypedChunk::FieldIdFacetExistsDocids(reader)));
+                }
+                Err(e) => {
+                    let _ = lmdb_writer_sx.send(Err(e));
+                }
+            }
+        });
+    }
 
-                        run_extraction_task::<_, _, grenad::Reader<BufReader<File>>>(
-                            fid_docid_facet_numbers_chunk.clone(),
-                            indexer,
-                            lmdb_writer_sx.clone(),
-                            extract_facet_number_docids,
-                            TypedChunk::FieldIdFacetNumberDocids,
-                            "field-id-facet-number-docids",
-                        );
+    // merge facet_is_null_docids and send them as a typed chunk
+    {
+        let lmdb_writer_sx = lmdb_writer_sx.clone();
+        rayon::spawn(move || {
+            debug!(database = "facet-id-is-null-docids", "merge");
+            match facet_is_null_docids_chunks.merge(merge_deladd_cbo_roaring_bitmaps, &indexer) {
+                Ok(reader) => {
+                    let _ = lmdb_writer_sx.send(Ok(TypedChunk::FieldIdFacetIsNullDocids(reader)));
+                }
+                Err(e) => {
+                    let _ = lmdb_writer_sx.send(Err(e));
+                }
+            }
+        });
+    }
 
-                        if proximity_precision == ProximityPrecision::ByWord {
-                            run_extraction_task::<_, _, grenad::Reader<BufReader<File>>>(
-                                docid_word_positions_chunk.clone(),
-                                indexer,
-                                lmdb_writer_sx.clone(),
-                                extract_word_pair_proximity_docids,
-                                TypedChunk::WordPairProximityDocids,
-                                "word-pair-proximity-docids",
-                            );
-                        }
-                    }
+    // merge facet_is_empty_docids and send them as a typed chunk
+    {
+        let lmdb_writer_sx = lmdb_writer_sx.clone();
+        rayon::spawn(move || {
+            debug!(database = "facet-id-is-empty-docids", "merge");
+            match facet_is_empty_docids_chunks.merge(merge_deladd_cbo_roaring_bitmaps, &indexer) {
+                Ok(reader) => {
+                    let _ = lmdb_writer_sx.send(Ok(TypedChunk::FieldIdFacetIsEmptyDocids(reader)));
+                }
+                Err(e) => {
+                    let _ = lmdb_writer_sx.send(Err(e));
+                }
+            }
+        });
+    }
 
-                    Ok(())
-                })
-                .collect::<Result<()>>()
-        },
+    if proximity_precision == ProximityPrecision::ByWord {
+        spawn_extraction_task::<_, _, Vec<grenad::Reader<BufReader<File>>>>(
+            docid_word_positions_chunks.clone(),
+            indexer,
+            lmdb_writer_sx.clone(),
+            extract_word_pair_proximity_docids,
+            merge_deladd_cbo_roaring_bitmaps,
+            TypedChunk::WordPairProximityDocids,
+            "word-pair-proximity-docids",
+        );
+    }
+
+    spawn_extraction_task::<_, _, Vec<grenad::Reader<BufReader<File>>>>(
+        docid_word_positions_chunks.clone(),
+        indexer,
+        lmdb_writer_sx.clone(),
+        extract_fid_word_count_docids,
+        merge_deladd_cbo_roaring_bitmaps,
+        TypedChunk::FieldIdWordCountDocids,
+        "field-id-wordcount-docids",
     );
 
-    original_pipeline_result.and(flattened_pipeline_result)
+    spawn_extraction_task::<
+        _,
+        _,
+        Vec<(
+            grenad::Reader<BufReader<File>>,
+            grenad::Reader<BufReader<File>>,
+            grenad::Reader<BufReader<File>>,
+        )>,
+    >(
+        docid_word_positions_chunks.clone(),
+        indexer,
+        lmdb_writer_sx.clone(),
+        move |doc_word_pos, indexer| extract_word_docids(doc_word_pos, indexer, &exact_attributes),
+        merge_deladd_cbo_roaring_bitmaps,
+        |(word_docids_reader, exact_word_docids_reader, word_fid_docids_reader)| {
+            TypedChunk::WordDocids {
+                word_docids_reader,
+                exact_word_docids_reader,
+                word_fid_docids_reader,
+            }
+        },
+        "word-docids",
+    );
+
+    spawn_extraction_task::<_, _, Vec<grenad::Reader<BufReader<File>>>>(
+        docid_word_positions_chunks.clone(),
+        indexer,
+        lmdb_writer_sx.clone(),
+        extract_word_position_docids,
+        merge_deladd_cbo_roaring_bitmaps,
+        TypedChunk::WordPositionDocids,
+        "word-position-docids",
+    );
+
+    spawn_extraction_task::<_, _, Vec<grenad::Reader<BufReader<File>>>>(
+        fid_docid_facet_strings_chunks,
+        indexer,
+        lmdb_writer_sx.clone(),
+        extract_facet_string_docids,
+        merge_deladd_cbo_roaring_bitmaps,
+        TypedChunk::FieldIdFacetStringDocids,
+        "field-id-facet-string-docids",
+    );
+
+    spawn_extraction_task::<_, _, Vec<grenad::Reader<BufReader<File>>>>(
+        fid_docid_facet_numbers_chunks,
+        indexer,
+        lmdb_writer_sx,
+        extract_facet_number_docids,
+        merge_deladd_cbo_roaring_bitmaps,
+        TypedChunk::FieldIdFacetNumberDocids,
+        "field-id-facet-number-docidsdexing::details, ",
+    );
+
+    Ok(())
 }
 
 /// Spawn a new task to extract data for a specific DB using extract_fn.
 /// Generated grenad chunks are merged using the merge_fn.
 /// The result of merged chunks is serialized as TypedChunk using the serialize_fn
 /// and sent into lmdb_writer_sx.
-fn run_extraction_task<FE, FS, M>(
-    chunk: grenad::Reader<CursorClonableMmap>,
+fn spawn_extraction_task<FE, FS, M>(
+    chunks: Vec<grenad::Reader<CursorClonableMmap>>,
     indexer: GrenadParameters,
     lmdb_writer_sx: Sender<Result<TypedChunk>>,
     extract_fn: FE,
+    merge_fn: MergeFn,
     serialize_fn: FS,
     name: &'static str,
 ) where
-    FE: Fn(grenad::Reader<CursorClonableMmap>, GrenadParameters) -> Result<M>
+    FE: Fn(grenad::Reader<CursorClonableMmap>, GrenadParameters) -> Result<M::Output>
         + Sync
         + Send
         + 'static,
-    FS: Fn(M) -> TypedChunk + Sync + Send + 'static,
-    M: Send,
+    FS: Fn(M::Output) -> TypedChunk + Sync + Send + 'static,
+    M: MergeableReader + FromParallelIterator<M::Output> + Send + 'static,
+    M::Output: Send,
 {
     let current_span = tracing::Span::current();
 
@@ -213,16 +264,25 @@ fn run_extraction_task<FE, FS, M>(
         let child_span =
             tracing::trace_span!(target: "", parent: &current_span, "extract_multiple_chunks");
         let _entered = child_span.enter();
-        puffin::profile_scope!("extract_multiple_chunks", name);
-        match extract_fn(chunk, indexer) {
-            Ok(chunk) => {
-                let _ = lmdb_writer_sx.send(Ok(serialize_fn(chunk)));
+        puffin::profile_scope!("extract_multiple_chunksdexing::details, ", name);
+        let chunks: Result<M> =
+            chunks.into_par_iter().map(|chunk| extract_fn(chunk, indexer)).collect();
+        let current_span = tracing::Span::current();
+
+        rayon::spawn(move || match chunks {
+            Ok(chunks) => {
+                let child_span = tracing::trace_span!(target: "", parent: &current_span, "merge_multiple_chunks");
+                let _entered = child_span.enter();
+                debug!(database = name, "merge");
+                puffin::profile_scope!("merge_multiple_chunks", name);
+                let reader = chunks.merge(merge_fn, &indexer);
+                let _ = lmdb_writer_sx.send(reader.map(serialize_fn));
             }
             Err(e) => {
                 let _ = lmdb_writer_sx.send(Err(e));
             }
-        }
-    })
+        })
+    });
 }
 
 /// Extract chunked data and send it into lmdb_writer_sx sender:
@@ -280,7 +340,7 @@ fn send_original_documents_data(
     });
 
     // TODO: create a custom internal error
-    let _ = lmdb_writer_sx.send(Ok(TypedChunk::Documents(original_documents_chunk)));
+    lmdb_writer_sx.send(Ok(TypedChunk::Documents(original_documents_chunk))).unwrap();
     Ok(())
 }
 
@@ -300,13 +360,22 @@ fn send_and_extract_flattened_documents_data(
     faceted_fields: &HashSet<FieldId>,
     primary_key_id: FieldId,
     geo_fields_ids: Option<(FieldId, FieldId)>,
-    stop_words: &Option<fst::Set<Vec<u8>>>,
+    stop_words: &Option<fst::Set<&[u8]>>,
     allowed_separators: &Option<&[&str]>,
     dictionary: &Option<&[&str]>,
     max_positions_per_attributes: Option<u32>,
 ) -> Result<(
     grenad::Reader<CursorClonableMmap>,
-    (grenad::Reader<CursorClonableMmap>, grenad::Reader<CursorClonableMmap>),
+    (
+        grenad::Reader<CursorClonableMmap>,
+        (
+            grenad::Reader<CursorClonableMmap>,
+            (
+                grenad::Reader<BufReader<File>>,
+                (grenad::Reader<BufReader<File>>, grenad::Reader<BufReader<File>>),
+            ),
+        ),
+    ),
 )> {
     let flattened_documents_chunk =
         flattened_documents_chunk.and_then(|c| unsafe { as_cloneable_grenad(&c) })?;
@@ -377,17 +446,16 @@ fn send_and_extract_flattened_documents_data(
                     fid_docid_facet_strings_chunk.clone(),
                 )));
 
-                let _ = lmdb_writer_sx
-                    .send(Ok(TypedChunk::FieldIdFacetIsNullDocids(fid_facet_is_null_docids_chunk)));
-
-                let _ = lmdb_writer_sx.send(Ok(TypedChunk::FieldIdFacetIsEmptyDocids(
-                    fid_facet_is_empty_docids_chunk,
-                )));
-
-                let _ = lmdb_writer_sx
-                    .send(Ok(TypedChunk::FieldIdFacetExistsDocids(fid_facet_exists_docids_chunk)));
-
-                Ok((fid_docid_facet_numbers_chunk, fid_docid_facet_strings_chunk))
+                Ok((
+                    fid_docid_facet_numbers_chunk,
+                    (
+                        fid_docid_facet_strings_chunk,
+                        (
+                            fid_facet_is_null_docids_chunk,
+                            (fid_facet_is_empty_docids_chunk, fid_facet_exists_docids_chunk),
+                        ),
+                    ),
+                ))
             },
         );
 

milli/src/update/index_documents/helpers/grenad_helpers.rs

@@ -9,10 +9,6 @@ use super::{ClonableMmap, MergeFn};
 use crate::update::index_documents::valid_lmdb_key;
 use crate::Result;
 
-/// This is something reasonable given the fact
-/// that there is one grenad sorter by thread.
-const MAX_GRENAD_SORTER_USAGE: usize = 500 * 1024 * 1024; // 500 MiB
-
 pub type CursorClonableMmap = io::Cursor<ClonableMmap>;
 
 pub fn create_writer<R: io::Write>(
@@ -28,9 +24,6 @@ pub fn create_writer<R: io::Write>(
     builder.build(BufWriter::new(file))
 }
 
-/// A helper function that creates a grenad sorter
-/// with the given parameters. The max memory is
-/// clamped to something reasonable.
 pub fn create_sorter(
     sort_algorithm: grenad::SortAlgorithm,
     merge: MergeFn,
@@ -48,7 +41,7 @@ pub fn create_sorter(
         builder.max_nb_chunks(nb_chunks);
     }
     if let Some(memory) = max_memory {
-        builder.dump_threshold(memory.min(MAX_GRENAD_SORTER_USAGE));
+        builder.dump_threshold(memory);
         builder.allow_realloc(false);
     }
     builder.sort_algorithm(sort_algorithm);
@@ -90,6 +83,90 @@ pub unsafe fn as_cloneable_grenad(
     Ok(reader)
 }
 
+pub trait MergeableReader
+where
+    Self: Sized,
+{
+    type Output;
+
+    fn merge(self, merge_fn: MergeFn, indexer: &GrenadParameters) -> Result<Self::Output>;
+}
+
+impl MergeableReader for Vec<grenad::Reader<BufReader<File>>> {
+    type Output = grenad::Reader<BufReader<File>>;
+
+    fn merge(self, merge_fn: MergeFn, params: &GrenadParameters) -> Result<Self::Output> {
+        let mut merger = MergerBuilder::new(merge_fn);
+        self.into_iter().try_for_each(|r| merger.push(r))?;
+        merger.finish(params)
+    }
+}
+
+impl MergeableReader for Vec<(grenad::Reader<BufReader<File>>, grenad::Reader<BufReader<File>>)> {
+    type Output = (grenad::Reader<BufReader<File>>, grenad::Reader<BufReader<File>>);
+
+    fn merge(self, merge_fn: MergeFn, params: &GrenadParameters) -> Result<Self::Output> {
+        let mut m1 = MergerBuilder::new(merge_fn);
+        let mut m2 = MergerBuilder::new(merge_fn);
+        for (r1, r2) in self.into_iter() {
+            m1.push(r1)?;
+            m2.push(r2)?;
+        }
+        Ok((m1.finish(params)?, m2.finish(params)?))
+    }
+}
+
+impl MergeableReader
+    for Vec<(
+        grenad::Reader<BufReader<File>>,
+        grenad::Reader<BufReader<File>>,
+        grenad::Reader<BufReader<File>>,
+    )>
+{
+    type Output = (
+        grenad::Reader<BufReader<File>>,
+        grenad::Reader<BufReader<File>>,
+        grenad::Reader<BufReader<File>>,
+    );
+
+    fn merge(self, merge_fn: MergeFn, params: &GrenadParameters) -> Result<Self::Output> {
+        let mut m1 = MergerBuilder::new(merge_fn);
+        let mut m2 = MergerBuilder::new(merge_fn);
+        let mut m3 = MergerBuilder::new(merge_fn);
+        for (r1, r2, r3) in self.into_iter() {
+            m1.push(r1)?;
+            m2.push(r2)?;
+            m3.push(r3)?;
+        }
+        Ok((m1.finish(params)?, m2.finish(params)?, m3.finish(params)?))
+    }
+}
+
+struct MergerBuilder<R>(grenad::MergerBuilder<R, MergeFn>);
+
+impl<R: io::Read + io::Seek> MergerBuilder<R> {
+    fn new(merge_fn: MergeFn) -> Self {
+        Self(grenad::MergerBuilder::new(merge_fn))
+    }
+
+    fn push(&mut self, reader: grenad::Reader<R>) -> Result<()> {
+        self.0.push(reader.into_cursor()?);
+        Ok(())
+    }
+
+    fn finish(self, params: &GrenadParameters) -> Result<grenad::Reader<BufReader<File>>> {
+        let merger = self.0.build();
+        let mut writer = create_writer(
+            params.chunk_compression_type,
+            params.chunk_compression_level,
+            tempfile::tempfile()?,
+        );
+        merger.write_into_stream_writer(&mut writer)?;
+
+        writer_into_reader(writer)
+    }
+}
+
 #[derive(Debug, Clone, Copy)]
 pub struct GrenadParameters {
     pub chunk_compression_type: CompressionType,
@@ -111,15 +188,10 @@ impl Default for GrenadParameters {
 
 impl GrenadParameters {
     /// This function use the number of threads in the current threadpool to compute the value.
-    ///
     /// This should be called inside of a rayon thread pool,
-    /// otherwise, it will take the global number of threads.
-    ///
-    /// The max memory cannot exceed a given reasonable value.
+    /// Otherwise, it will take the global number of threads.
     pub fn max_memory_by_thread(&self) -> Option<usize> {
-        self.max_memory.map(|max_memory| {
-            (max_memory / rayon::current_num_threads()).min(MAX_GRENAD_SORTER_USAGE)
-        })
+        self.max_memory.map(|max_memory| max_memory / rayon::current_num_threads())
     }
 }
 

milli/src/update/index_documents/helpers/merge_functions.rs

@@ -35,6 +35,27 @@ pub fn merge_roaring_bitmaps<'a>(_key: &[u8], values: &[Cow<'a, [u8]>]) -> Resul
     }
 }
 
+pub fn merge_btreeset_string<'a>(_key: &[u8], values: &[Cow<'a, [u8]>]) -> Result<Cow<'a, [u8]>> {
+    if values.len() == 1 {
+        Ok(values[0].clone())
+    } else {
+        // TODO improve the perf by using a `#[borrow] Cow<str>`.
+        let strings: BTreeSet<String> = values
+            .iter()
+            .map(AsRef::as_ref)
+            .map(serde_json::from_slice::<BTreeSet<String>>)
+            .map(StdResult::unwrap)
+            .reduce(|mut current, new| {
+                for x in new {
+                    current.insert(x);
+                }
+                current
+            })
+            .unwrap();
+        Ok(Cow::Owned(serde_json::to_vec(&strings).unwrap()))
+    }
+}
+
 pub fn keep_first<'a>(_key: &[u8], values: &[Cow<'a, [u8]>]) -> Result<Cow<'a, [u8]>> {
     Ok(values[0].clone())
 }
@@ -222,40 +243,3 @@ pub fn merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap<'a>(
         buffer,
     )?)
 }
-
-/// Do a union of BtreeSet on both sides of a DelAdd obkv
-/// separately and outputs a new DelAdd with both unions.
-pub fn merge_deladd_btreeset_string<'a>(
-    _key: &[u8],
-    values: &[Cow<'a, [u8]>],
-) -> Result<Cow<'a, [u8]>> {
-    if values.len() == 1 {
-        Ok(values[0].clone())
-    } else {
-        // Retrieve the bitmaps from both sides
-        let mut del_set = BTreeSet::new();
-        let mut add_set = BTreeSet::new();
-        for value in values {
-            let obkv = KvReaderDelAdd::new(value);
-            if let Some(bytes) = obkv.get(DelAdd::Deletion) {
-                let set = serde_json::from_slice::<BTreeSet<String>>(bytes).unwrap();
-                for value in set {
-                    del_set.insert(value);
-                }
-            }
-            if let Some(bytes) = obkv.get(DelAdd::Addition) {
-                let set = serde_json::from_slice::<BTreeSet<String>>(bytes).unwrap();
-                for value in set {
-                    add_set.insert(value);
-                }
-            }
-        }
-
-        let mut output_deladd_obkv = KvWriterDelAdd::memory();
-        let del = serde_json::to_vec(&del_set).unwrap();
-        output_deladd_obkv.insert(DelAdd::Deletion, &del)?;
-        let add = serde_json::to_vec(&add_set).unwrap();
-        output_deladd_obkv.insert(DelAdd::Addition, &add)?;
-        output_deladd_obkv.into_inner().map(Cow::from).map_err(Into::into)
-    }
-}

milli/src/update/index_documents/helpers/mod.rs

@@ -10,10 +10,10 @@ use fst::{IntoStreamer, Streamer};
 pub use grenad_helpers::{
     as_cloneable_grenad, create_sorter, create_writer, grenad_obkv_into_chunks,
     merge_ignore_values, sorter_into_reader, write_sorter_into_database, writer_into_reader,
-    GrenadParameters,
+    GrenadParameters, MergeableReader,
 };
 pub use merge_functions::{
-    keep_first, keep_latest_obkv, merge_cbo_roaring_bitmaps, merge_deladd_btreeset_string,
+    keep_first, keep_latest_obkv, merge_btreeset_string, merge_cbo_roaring_bitmaps,
     merge_deladd_cbo_roaring_bitmaps, merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap,
     merge_roaring_bitmaps, obkvs_keep_last_addition_merge_deletions,
     obkvs_merge_additions_and_deletions, MergeFn,

milli/src/update/index_documents/mod.rs

@@ -5,29 +5,29 @@ mod transform;
 mod typed_chunk;
 
 use std::collections::{HashMap, HashSet};
-use std::io::{Read, Seek};
+use std::io::{Cursor, Read, Seek};
 use std::iter::FromIterator;
 use std::num::NonZeroU32;
 use std::result::Result as StdResult;
 
 use crossbeam_channel::{Receiver, Sender};
-use grenad::{Merger, MergerBuilder};
 use heed::types::Str;
 use heed::Database;
 use rand::SeedableRng;
 use roaring::RoaringBitmap;
 use serde::{Deserialize, Serialize};
 use slice_group_by::GroupBy;
-use tracing::debug;
-use typed_chunk::{write_typed_chunk_into_index, ChunkAccumulator, TypedChunk};
+use tracing::{debug_span};
+use typed_chunk::{write_typed_chunk_into_index, TypedChunk};
 
 use self::enrich::enrich_documents_batch;
 pub use self::enrich::{extract_finite_float_from_value, DocumentId};
 pub use self::helpers::{
     as_cloneable_grenad, create_sorter, create_writer, fst_stream_into_hashset,
-    fst_stream_into_vec, merge_cbo_roaring_bitmaps, merge_deladd_cbo_roaring_bitmaps,
-    merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap, merge_roaring_bitmaps,
-    valid_lmdb_key, write_sorter_into_database, writer_into_reader, MergeFn,
+    fst_stream_into_vec, merge_btreeset_string, merge_cbo_roaring_bitmaps,
+    merge_deladd_cbo_roaring_bitmaps, merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap,
+    merge_roaring_bitmaps, valid_lmdb_key, write_sorter_into_database, writer_into_reader,
+    ClonableMmap, MergeFn,
 };
 use self::helpers::{grenad_obkv_into_chunks, GrenadParameters};
 pub use self::transform::{Transform, TransformOutput};
@@ -95,8 +95,8 @@ pub struct IndexDocumentsConfig {
 
 impl<'t, 'i, 'a, FP, FA> IndexDocuments<'t, 'i, 'a, FP, FA>
 where
-    FP: Fn(UpdateIndexingStep) + Sync + Send,
-    FA: Fn() -> bool + Sync + Send,
+    FP: Fn(UpdateIndexingStep) + Sync,
+    FA: Fn() -> bool + Sync,
 {
     pub fn new(
         wtxn: &'t mut heed::RwTxn<'i>,
@@ -326,6 +326,9 @@ where
             }
         };
 
+        let original_documents = grenad::Reader::new(original_documents)?;
+        let flattened_documents = grenad::Reader::new(flattened_documents)?;
+
         // create LMDB writer channel
         let (lmdb_writer_sx, lmdb_writer_rx): (
             Sender<Result<TypedChunk>>,
@@ -364,7 +367,11 @@ where
 
         let stop_words = self.index.stop_words(self.wtxn)?;
         let separators = self.index.allowed_separators(self.wtxn)?;
+        let separators: Option<Vec<_>> =
+            separators.as_ref().map(|x| x.iter().map(String::as_str).collect());
         let dictionary = self.index.dictionary(self.wtxn)?;
+        let dictionary: Option<Vec<_>> =
+            dictionary.as_ref().map(|x| x.iter().map(String::as_str).collect());
         let exact_attributes = self.index.exact_attributes_ids(self.wtxn)?;
         let proximity_precision = self.index.proximity_precision(self.wtxn)?.unwrap_or_default();
 
@@ -374,204 +381,141 @@ where
             max_memory: self.indexer_config.max_memory,
             max_nb_chunks: self.indexer_config.max_nb_chunks, // default value, may be chosen.
         };
-        let documents_chunk_size = match self.indexer_config.documents_chunk_size {
-            Some(chunk_size) => chunk_size,
-            None => {
-                let default_chunk_size = 1024 * 1024 * 4; // 4MiB
-                let min_chunk_size = 1024 * 512; // 512KiB
-
-                // compute the chunk size from the number of available threads and the inputed data size.
-                let total_size = flattened_documents.metadata().map(|m| m.len());
-                let current_num_threads = pool.current_num_threads();
-                // if we have more than 2 thread, create a number of chunk equal to 3/4 threads count
-                let chunk_count = if current_num_threads > 2 {
-                    (current_num_threads * 3 / 4).max(2)
-                } else {
-                    current_num_threads
-                };
-                total_size
-                    .map_or(default_chunk_size, |size| (size as usize) / chunk_count)
-                    .max(min_chunk_size)
-            }
-        };
-
-        let original_documents = grenad::Reader::new(original_documents)?;
-        let flattened_documents = grenad::Reader::new(flattened_documents)?;
-
+        let documents_chunk_size =
+            self.indexer_config.documents_chunk_size.unwrap_or(1024 * 1024 * 4); // 4MiB
         let max_positions_per_attributes = self.indexer_config.max_positions_per_attributes;
 
         let cloned_embedder = self.embedders.clone();
 
-        let mut final_documents_ids = RoaringBitmap::new();
-        let mut databases_seen = 0;
-        let mut word_position_docids = None;
-        let mut word_fid_docids = None;
-        let mut word_docids = None;
-        let mut exact_word_docids = None;
-        let mut chunk_accumulator = ChunkAccumulator::default();
-        let mut dimension = HashMap::new();
-        let stop_words = stop_words.map(|sw| sw.map_data(Vec::from).unwrap());
-
         let current_span = tracing::Span::current();
 
         // Run extraction pipeline in parallel.
         pool.install(|| {
-            rayon::spawn(move || {
-                let child_span = tracing::trace_span!(target: "indexing::details", parent: &current_span, "extract_and_send_grenad_chunks");
+            let child_span = tracing::trace_span!(target: "indexing::details", parent: &current_span, "extract_and_send_grenad_chunks");
             let _enter = child_span.enter();
             puffin::profile_scope!("extract_and_send_grenad_chunks");
-                // split obkv file into several chunks
-                let original_chunk_iter =
-                    grenad_obkv_into_chunks(original_documents, pool_params, documents_chunk_size);
+            // split obkv file into several chunks
+            let original_chunk_iter =
+                grenad_obkv_into_chunks(original_documents, pool_params, documents_chunk_size);
 
-                // split obkv file into several chunks
-                let flattened_chunk_iter =
-                    grenad_obkv_into_chunks(flattened_documents, pool_params, documents_chunk_size);
+            // split obkv file into several chunks
+            let flattened_chunk_iter =
+                grenad_obkv_into_chunks(flattened_documents, pool_params, documents_chunk_size);
 
-                let separators: Option<Vec<_>> =
-                    separators.as_ref().map(|x| x.iter().map(String::as_str).collect());
-                let dictionary: Option<Vec<_>> =
-                    dictionary.as_ref().map(|x| x.iter().map(String::as_str).collect());
-                let result = original_chunk_iter.and_then(|original_chunk| {
-                    let flattened_chunk = flattened_chunk_iter?;
-                    // extract all databases from the chunked obkv douments
-                    extract::data_from_obkv_documents(
-                        original_chunk,
-                        flattened_chunk,
-                        pool_params,
-                        lmdb_writer_sx.clone(),
-                        searchable_fields,
-                        faceted_fields,
-                        primary_key_id,
-                        geo_fields_ids,
-                        field_id_map,
-                        stop_words,
-                        separators.as_deref(),
-                        dictionary.as_deref(),
-                        max_positions_per_attributes,
-                        exact_attributes,
-                        proximity_precision,
-                        cloned_embedder,
-                    )
-                });
-
-                if let Err(e) = result {
-                    let _ = lmdb_writer_sx.send(Err(e));
-                }
-
-                // needs to be dropped to avoid channel waiting lock.
-                drop(lmdb_writer_sx);
+            let result = original_chunk_iter.and_then(|original_chunk| {
+                let flattened_chunk = flattened_chunk_iter?;
+                // extract all databases from the chunked obkv douments
+                extract::data_from_obkv_documents(
+                    original_chunk,
+                    flattened_chunk,
+                    pool_params,
+                    lmdb_writer_sx.clone(),
+                    searchable_fields,
+                    faceted_fields,
+                    primary_key_id,
+                    geo_fields_ids,
+                    field_id_map,
+                    stop_words,
+                    separators.as_deref(),
+                    dictionary.as_deref(),
+                    max_positions_per_attributes,
+                    exact_attributes,
+                    proximity_precision,
+                    cloned_embedder,
+                )
             });
 
-            (self.progress)(UpdateIndexingStep::MergeDataIntoFinalDatabase {
-                databases_seen,
-                total_databases: TOTAL_POSTING_DATABASE_COUNT,
-            });
-
-            loop {
-                if (self.should_abort)() {
-                    return Err(Error::InternalError(InternalError::AbortedIndexation));
-                }
-
-                match lmdb_writer_rx.clone().recv_timeout(std::time::Duration::from_millis(500)) {
-                    Err(status) => {
-                        if let Some(typed_chunks) = chunk_accumulator.pop_longest() {
-                            let (docids, is_merged_database) =
-                                write_typed_chunk_into_index(typed_chunks, self.index, self.wtxn)?;
-                            if !docids.is_empty() {
-                                final_documents_ids |= docids;
-                                let documents_seen_count = final_documents_ids.len();
-                                (self.progress)(UpdateIndexingStep::IndexDocuments {
-                                    documents_seen: documents_seen_count as usize,
-                                    total_documents: documents_count,
-                                });
-                                debug!(documents = documents_seen_count, total = documents_count, "Seen");
-                            }
-                            if is_merged_database {
-                                databases_seen += 1;
-                                (self.progress)(UpdateIndexingStep::MergeDataIntoFinalDatabase {
-                                    databases_seen,
-                                    total_databases: TOTAL_POSTING_DATABASE_COUNT,
-                                });
-                            }
-                        // If no more chunk remains in the chunk accumulator and the channel is disconected, break.
-                        } else if status == crossbeam_channel::RecvTimeoutError::Disconnected {
-                            break;
-                        } else {
-                            rayon::yield_now();
-                        }
-                    }
-                    Ok(result) => {
-                        let typed_chunk = match result? {
-                            TypedChunk::WordDocids {
-                                word_docids_reader,
-                                exact_word_docids_reader,
-                                word_fid_docids_reader,
-                            } => {
-                                let cloneable_chunk =
-                                    unsafe { as_cloneable_grenad(&word_docids_reader)? };
-                                let word_docids = word_docids.get_or_insert_with(|| {
-                                    MergerBuilder::new(merge_deladd_cbo_roaring_bitmaps as MergeFn)
-                                });
-                                word_docids.push(cloneable_chunk.into_cursor()?);
-                                let cloneable_chunk =
-                                    unsafe { as_cloneable_grenad(&exact_word_docids_reader)? };
-                                let exact_word_docids =
-                                    exact_word_docids.get_or_insert_with(|| {
-                                        MergerBuilder::new(
-                                            merge_deladd_cbo_roaring_bitmaps as MergeFn,
-                                        )
-                                    });
-                                exact_word_docids.push(cloneable_chunk.into_cursor()?);
-                                let cloneable_chunk =
-                                    unsafe { as_cloneable_grenad(&word_fid_docids_reader)? };
-                                let word_fid_docids = word_fid_docids.get_or_insert_with(|| {
-                                    MergerBuilder::new(merge_deladd_cbo_roaring_bitmaps as MergeFn)
-                                });
-                                word_fid_docids.push(cloneable_chunk.into_cursor()?);
-                                TypedChunk::WordDocids {
-                                    word_docids_reader,
-                                    exact_word_docids_reader,
-                                    word_fid_docids_reader,
-                                }
-                            }
-                            TypedChunk::WordPositionDocids(chunk) => {
-                                let cloneable_chunk = unsafe { as_cloneable_grenad(&chunk)? };
-                                let word_position_docids =
-                                    word_position_docids.get_or_insert_with(|| {
-                                        MergerBuilder::new(
-                                            merge_deladd_cbo_roaring_bitmaps as MergeFn,
-                                        )
-                                    });
-                                word_position_docids.push(cloneable_chunk.into_cursor()?);
-                                TypedChunk::WordPositionDocids(chunk)
-                            }
-                            TypedChunk::VectorPoints {
-                                expected_dimension,
-                                remove_vectors,
-                                embeddings,
-                                manual_vectors,
-                                embedder_name,
-                            } => {
-                                dimension.insert(embedder_name.clone(), expected_dimension);
-                                TypedChunk::VectorPoints {
-                                    remove_vectors,
-                                    embeddings,
-                                    expected_dimension,
-                                    manual_vectors,
-                                    embedder_name,
-                                }
-                            }
-                            otherwise => otherwise,
-                        };
-
-                        chunk_accumulator.insert(typed_chunk);
-                    }
-                }
+            if let Err(e) = result {
+                let _ = lmdb_writer_sx.send(Err(e));
             }
 
-            Ok(())
-        })?;
+            // needs to be dropped to avoid channel waiting lock.
+            drop(lmdb_writer_sx);
+        });
+
+        let index_is_empty = self.index.number_of_documents(self.wtxn)? == 0;
+        let mut final_documents_ids = RoaringBitmap::new();
+
+        let mut databases_seen = 0;
+        (self.progress)(UpdateIndexingStep::MergeDataIntoFinalDatabase {
+            databases_seen,
+            total_databases: TOTAL_POSTING_DATABASE_COUNT,
+        });
+
+        let mut word_position_docids = None;
+        let mut word_fid_docids = None;
+        let mut word_docids = None;
+        let mut exact_word_docids = None;
+
+        let mut dimension = HashMap::new();
+
+        for result in lmdb_writer_rx {
+            if (self.should_abort)() {
+                return Err(Error::InternalError(InternalError::AbortedIndexation));
+            }
+
+            let typed_chunk = match result? {
+                TypedChunk::WordDocids {
+                    word_docids_reader,
+                    exact_word_docids_reader,
+                    word_fid_docids_reader,
+                } => {
+                    let cloneable_chunk = unsafe { as_cloneable_grenad(&word_docids_reader)? };
+                    word_docids = Some(cloneable_chunk);
+                    let cloneable_chunk =
+                        unsafe { as_cloneable_grenad(&exact_word_docids_reader)? };
+                    exact_word_docids = Some(cloneable_chunk);
+                    let cloneable_chunk = unsafe { as_cloneable_grenad(&word_fid_docids_reader)? };
+                    word_fid_docids = Some(cloneable_chunk);
+                    TypedChunk::WordDocids {
+                        word_docids_reader,
+                        exact_word_docids_reader,
+                        word_fid_docids_reader,
+                    }
+                }
+                TypedChunk::WordPositionDocids(chunk) => {
+                    let cloneable_chunk = unsafe { as_cloneable_grenad(&chunk)? };
+                    word_position_docids = Some(cloneable_chunk);
+                    TypedChunk::WordPositionDocids(chunk)
+                }
+                TypedChunk::VectorPoints {
+                    expected_dimension,
+                    remove_vectors,
+                    embeddings,
+                    manual_vectors,
+                    embedder_name,
+                } => {
+                    dimension.insert(embedder_name.clone(), expected_dimension);
+                    TypedChunk::VectorPoints {
+                        remove_vectors,
+                        embeddings,
+                        expected_dimension,
+                        manual_vectors,
+                        embedder_name,
+                    }
+                }
+                otherwise => otherwise,
+            };
+
+            let (docids, is_merged_database) =
+                write_typed_chunk_into_index(typed_chunk, self.index, self.wtxn, index_is_empty)?;
+            if !docids.is_empty() {
+                final_documents_ids |= docids;
+                let documents_seen_count = final_documents_ids.len();
+                (self.progress)(UpdateIndexingStep::IndexDocuments {
+                    documents_seen: documents_seen_count as usize,
+                    total_documents: documents_count,
+                });
+                debug_span!("Seen", documents = documents_seen_count, total = documents_count);
+            }
+            if is_merged_database {
+                databases_seen += 1;
+                (self.progress)(UpdateIndexingStep::MergeDataIntoFinalDatabase {
+                    databases_seen,
+                    total_databases: TOTAL_POSTING_DATABASE_COUNT,
+                });
+            }
+        }
 
         // We write the field distribution into the main database
         self.index.put_field_distribution(self.wtxn, &field_distribution)?;
@@ -604,10 +548,10 @@ where
         }
 
         self.execute_prefix_databases(
-            word_docids.map(MergerBuilder::build),
-            exact_word_docids.map(MergerBuilder::build),
-            word_position_docids.map(MergerBuilder::build),
-            word_fid_docids.map(MergerBuilder::build),
+            word_docids,
+            exact_word_docids,
+            word_position_docids,
+            word_fid_docids,
         )?;
 
         Ok(number_of_documents)
@@ -621,10 +565,10 @@ where
     )]
     pub fn execute_prefix_databases(
         self,
-        word_docids: Option<Merger<CursorClonableMmap, MergeFn>>,
-        exact_word_docids: Option<Merger<CursorClonableMmap, MergeFn>>,
-        word_position_docids: Option<Merger<CursorClonableMmap, MergeFn>>,
-        word_fid_docids: Option<Merger<CursorClonableMmap, MergeFn>>,
+        word_docids: Option<grenad::Reader<CursorClonableMmap>>,
+        exact_word_docids: Option<grenad::Reader<CursorClonableMmap>>,
+        word_position_docids: Option<grenad::Reader<CursorClonableMmap>>,
+        word_fid_docids: Option<grenad::Reader<CursorClonableMmap>>,
     ) -> Result<()>
     where
         FP: Fn(UpdateIndexingStep) + Sync,
@@ -807,7 +751,7 @@ where
 )]
 fn execute_word_prefix_docids(
     txn: &mut heed::RwTxn,
-    merger: Merger<CursorClonableMmap, MergeFn>,
+    reader: grenad::Reader<Cursor<ClonableMmap>>,
     word_docids_db: Database<Str, CboRoaringBitmapCodec>,
     word_prefix_docids_db: Database<Str, CboRoaringBitmapCodec>,
     indexer_config: &IndexerConfig,
@@ -817,12 +761,13 @@ fn execute_word_prefix_docids(
 ) -> Result<()> {
     puffin::profile_function!();
 
+    let cursor = reader.into_cursor()?;
     let mut builder = WordPrefixDocids::new(txn, word_docids_db, word_prefix_docids_db);
     builder.chunk_compression_type = indexer_config.chunk_compression_type;
     builder.chunk_compression_level = indexer_config.chunk_compression_level;
     builder.max_nb_chunks = indexer_config.max_nb_chunks;
     builder.max_memory = indexer_config.max_memory;
-    builder.execute(merger, new_prefix_fst_words, common_prefix_fst_words, del_prefix_fst_words)?;
+    builder.execute(cursor, new_prefix_fst_words, common_prefix_fst_words, del_prefix_fst_words)?;
     Ok(())
 }
 

milli/src/update/index_documents/typed_chunk.rs

@@ -5,64 +5,27 @@ use std::io::{self, BufReader};
 
 use bytemuck::allocation::pod_collect_to_vec;
 use charabia::{Language, Script};
-use grenad::{Merger, MergerBuilder};
+use grenad::MergerBuilder;
 use heed::types::Bytes;
-use heed::RwTxn;
+use heed::{PutFlags, RwTxn};
 use obkv::{KvReader, KvWriter};
 use roaring::RoaringBitmap;
 
 use super::helpers::{
-    self, keep_first, merge_deladd_btreeset_string, merge_deladd_cbo_roaring_bitmaps,
-    merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap, merge_ignore_values, valid_lmdb_key,
-    CursorClonableMmap,
+    self, merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap, merge_ignore_values,
+    valid_lmdb_key, CursorClonableMmap,
 };
-use super::MergeFn;
+use super::{ClonableMmap, MergeFn};
 use crate::external_documents_ids::{DocumentOperation, DocumentOperationKind};
 use crate::facet::FacetType;
 use crate::index::db_name::DOCUMENTS;
 use crate::update::del_add::{deladd_serialize_add_side, DelAdd, KvReaderDelAdd};
 use crate::update::facet::FacetsUpdate;
-use crate::update::index_documents::helpers::{
-    as_cloneable_grenad, keep_latest_obkv, try_split_array_at,
-};
+use crate::update::index_documents::helpers::{as_cloneable_grenad, try_split_array_at};
 use crate::{
     lat_lng_to_xyz, DocumentId, FieldId, GeoPoint, Index, InternalError, Result, SerializationError,
 };
 
-/// This struct accumulates and group the TypedChunks
-/// and is able to give the biggest accumulated group to index them all together
-/// with a merger.
-#[derive(Default)]
-pub(crate) struct ChunkAccumulator {
-    inner: Vec<Vec<TypedChunk>>,
-}
-
-impl ChunkAccumulator {
-    pub fn pop_longest(&mut self) -> Option<Vec<TypedChunk>> {
-        match self.inner.iter().max_by_key(|v| v.len()) {
-            Some(left) => {
-                let position = self.inner.iter().position(|right| left.len() == right.len());
-                position.map(|p| self.inner.remove(p)).filter(|v| !v.is_empty())
-            }
-            None => None,
-        }
-    }
-
-    pub fn insert(&mut self, chunk: TypedChunk) {
-        match self
-            .inner
-            .iter()
-            .position(|right| right.first().map_or(false, |right| chunk.mergeable_with(right)))
-        {
-            Some(position) => {
-                let v = self.inner.get_mut(position).unwrap();
-                v.push(chunk);
-            }
-            None => self.inner.push(vec![chunk]),
-        }
-    }
-}
-
 pub(crate) enum TypedChunk {
     FieldIdDocidFacetStrings(grenad::Reader<CursorClonableMmap>),
     FieldIdDocidFacetNumbers(grenad::Reader<CursorClonableMmap>),
@@ -75,7 +38,7 @@ pub(crate) enum TypedChunk {
     },
     WordPositionDocids(grenad::Reader<BufReader<File>>),
     WordPairProximityDocids(grenad::Reader<BufReader<File>>),
-    FieldIdFacetStringDocids((grenad::Reader<BufReader<File>>, grenad::Reader<BufReader<File>>)),
+    FieldIdFacetStringDocids(grenad::Reader<BufReader<File>>),
     FieldIdFacetNumberDocids(grenad::Reader<BufReader<File>>),
     FieldIdFacetExistsDocids(grenad::Reader<BufReader<File>>),
     FieldIdFacetIsNullDocids(grenad::Reader<BufReader<File>>),
@@ -91,33 +54,6 @@ pub(crate) enum TypedChunk {
     ScriptLanguageDocids(HashMap<(Script, Language), (RoaringBitmap, RoaringBitmap)>),
 }
 
-impl TypedChunk {
-    fn mergeable_with(&self, other: &Self) -> bool {
-        use TypedChunk::*;
-        match (self, other) {
-            (FieldIdDocidFacetStrings(_), FieldIdDocidFacetStrings(_))
-            | (FieldIdDocidFacetNumbers(_), FieldIdDocidFacetNumbers(_))
-            | (Documents(_), Documents(_))
-            | (FieldIdWordCountDocids(_), FieldIdWordCountDocids(_))
-            | (WordDocids { .. }, WordDocids { .. })
-            | (WordPositionDocids(_), WordPositionDocids(_))
-            | (WordPairProximityDocids(_), WordPairProximityDocids(_))
-            | (FieldIdFacetStringDocids(_), FieldIdFacetStringDocids(_))
-            | (FieldIdFacetNumberDocids(_), FieldIdFacetNumberDocids(_))
-            | (FieldIdFacetExistsDocids(_), FieldIdFacetExistsDocids(_))
-            | (FieldIdFacetIsNullDocids(_), FieldIdFacetIsNullDocids(_))
-            | (FieldIdFacetIsEmptyDocids(_), FieldIdFacetIsEmptyDocids(_))
-            | (GeoPoints(_), GeoPoints(_))
-            | (ScriptLanguageDocids(_), ScriptLanguageDocids(_)) => true,
-            (
-                VectorPoints { embedder_name: left, expected_dimension: left_dim, .. },
-                VectorPoints { embedder_name: right, expected_dimension: right_dim, .. },
-            ) => left == right && left_dim == right_dim,
-            _ => false,
-        }
-    }
-}
-
 impl TypedChunk {
     pub fn to_debug_string(&self) -> String {
         match self {
@@ -149,7 +85,7 @@ impl TypedChunk {
             TypedChunk::WordPairProximityDocids(grenad) => {
                 format!("WordPairProximityDocids {{ number_of_entries: {} }}", grenad.len())
             }
-            TypedChunk::FieldIdFacetStringDocids((grenad, _)) => {
+            TypedChunk::FieldIdFacetStringDocids(grenad) => {
                 format!("FieldIdFacetStringDocids {{ number_of_entries: {} }}", grenad.len())
             }
             TypedChunk::FieldIdFacetNumberDocids(grenad) => {
@@ -181,32 +117,23 @@ impl TypedChunk {
 /// Return new documents seen.
 #[tracing::instrument(level = "trace", skip_all, target = "indexing::write_db")]
 pub(crate) fn write_typed_chunk_into_index(
-    typed_chunks: Vec<TypedChunk>,
+    typed_chunk: TypedChunk,
     index: &Index,
     wtxn: &mut RwTxn,
+    index_is_empty: bool,
 ) -> Result<(RoaringBitmap, bool)> {
-    puffin::profile_function!(typed_chunks[0].to_debug_string());
+    puffin::profile_function!(typed_chunk.to_debug_string());
 
     let mut is_merged_database = false;
-    match typed_chunks[0] {
-        TypedChunk::Documents(_) => {
+    match typed_chunk {
+        TypedChunk::Documents(obkv_documents_iter) => {
             let span = tracing::trace_span!(target: "indexing::write_db", "documents");
             let _entered = span.enter();
-
-            let mut builder = MergerBuilder::new(keep_latest_obkv as MergeFn);
-            for typed_chunk in typed_chunks {
-                let TypedChunk::Documents(chunk) = typed_chunk else {
-                    unreachable!();
-                };
-
-                builder.push(chunk.into_cursor()?);
-            }
-            let merger = builder.build();
             let mut operations: Vec<DocumentOperation> = Default::default();
 
             let mut docids = index.documents_ids(wtxn)?;
-            let mut iter = merger.into_stream_merger_iter()?;
-            while let Some((key, reader)) = iter.next()? {
+            let mut cursor = obkv_documents_iter.into_cursor()?;
+            while let Some((key, reader)) = cursor.move_on_next()? {
                 let mut writer: KvWriter<_, FieldId> = KvWriter::memory();
                 let reader: KvReader<FieldId> = KvReader::new(reader);
 
@@ -247,91 +174,59 @@ pub(crate) fn write_typed_chunk_into_index(
             external_documents_docids.apply(wtxn, operations)?;
             index.put_documents_ids(wtxn, &docids)?;
         }
-        TypedChunk::FieldIdWordCountDocids(_) => {
+        TypedChunk::FieldIdWordCountDocids(fid_word_count_docids_iter) => {
             let span =
                 tracing::trace_span!(target: "indexing::write_db", "field_id_word_count_docids");
             let _entered = span.enter();
-
-            let mut builder = MergerBuilder::new(merge_deladd_cbo_roaring_bitmaps as MergeFn);
-            for typed_chunk in typed_chunks {
-                let TypedChunk::FieldIdWordCountDocids(chunk) = typed_chunk else {
-                    unreachable!();
-                };
-
-                builder.push(chunk.into_cursor()?);
-            }
-            let merger = builder.build();
-
-            write_entries_into_database(
-                merger,
+            append_entries_into_database(
+                fid_word_count_docids_iter,
                 &index.field_id_word_count_docids,
                 wtxn,
+                index_is_empty,
                 deladd_serialize_add_side,
                 merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap,
             )?;
             is_merged_database = true;
         }
-        TypedChunk::WordDocids { .. } => {
+        TypedChunk::WordDocids {
+            word_docids_reader,
+            exact_word_docids_reader,
+            word_fid_docids_reader,
+        } => {
             let span = tracing::trace_span!(target: "indexing::write_db", "word_docids");
             let _entered = span.enter();
-
-            let mut word_docids_builder =
-                MergerBuilder::new(merge_deladd_cbo_roaring_bitmaps as MergeFn);
-            let mut exact_word_docids_builder =
-                MergerBuilder::new(merge_deladd_cbo_roaring_bitmaps as MergeFn);
-            let mut word_fid_docids_builder =
-                MergerBuilder::new(merge_deladd_cbo_roaring_bitmaps as MergeFn);
-            let mut fst_merger_builder = MergerBuilder::new(merge_ignore_values as MergeFn);
-            for typed_chunk in typed_chunks {
-                let TypedChunk::WordDocids {
-                    word_docids_reader,
-                    exact_word_docids_reader,
-                    word_fid_docids_reader,
-                } = typed_chunk
-                else {
-                    unreachable!();
-                };
-                let clonable_word_docids = unsafe { as_cloneable_grenad(&word_docids_reader) }?;
-                let clonable_exact_word_docids =
-                    unsafe { as_cloneable_grenad(&exact_word_docids_reader) }?;
-
-                word_docids_builder.push(word_docids_reader.into_cursor()?);
-                exact_word_docids_builder.push(exact_word_docids_reader.into_cursor()?);
-                word_fid_docids_builder.push(word_fid_docids_reader.into_cursor()?);
-                fst_merger_builder.push(clonable_word_docids.into_cursor()?);
-                fst_merger_builder.push(clonable_exact_word_docids.into_cursor()?);
-            }
-
-            let word_docids_merger = word_docids_builder.build();
-            write_entries_into_database(
-                word_docids_merger,
+            let word_docids_iter = unsafe { as_cloneable_grenad(&word_docids_reader) }?;
+            append_entries_into_database(
+                word_docids_iter.clone(),
                 &index.word_docids,
                 wtxn,
+                index_is_empty,
                 deladd_serialize_add_side,
                 merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap,
             )?;
 
-            let exact_word_docids_merger = exact_word_docids_builder.build();
-            write_entries_into_database(
-                exact_word_docids_merger,
+            let exact_word_docids_iter = unsafe { as_cloneable_grenad(&exact_word_docids_reader) }?;
+            append_entries_into_database(
+                exact_word_docids_iter.clone(),
                 &index.exact_word_docids,
                 wtxn,
+                index_is_empty,
                 deladd_serialize_add_side,
                 merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap,
             )?;
 
-            let word_fid_docids_merger = word_fid_docids_builder.build();
-            write_entries_into_database(
-                word_fid_docids_merger,
+            let word_fid_docids_iter = unsafe { as_cloneable_grenad(&word_fid_docids_reader) }?;
+            append_entries_into_database(
+                word_fid_docids_iter,
                 &index.word_fid_docids,
                 wtxn,
+                index_is_empty,
                 deladd_serialize_add_side,
                 merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap,
             )?;
 
             // create fst from word docids
-            let fst_merger = fst_merger_builder.build();
-            let fst = merge_word_docids_reader_into_fst(fst_merger)?;
+            let fst = merge_word_docids_reader_into_fst(word_docids_iter, exact_word_docids_iter)?;
             let db_fst = index.words_fst(wtxn)?;
 
             // merge new fst with database fst
@@ -342,202 +237,98 @@ pub(crate) fn write_typed_chunk_into_index(
             index.put_words_fst(wtxn, &fst)?;
             is_merged_database = true;
         }
-        TypedChunk::WordPositionDocids(_) => {
+        TypedChunk::WordPositionDocids(word_position_docids_iter) => {
             let span = tracing::trace_span!(target: "indexing::write_db", "word_position_docids");
             let _entered = span.enter();
-
-            let mut builder = MergerBuilder::new(merge_deladd_cbo_roaring_bitmaps as MergeFn);
-            for typed_chunk in typed_chunks {
-                let TypedChunk::WordPositionDocids(chunk) = typed_chunk else {
-                    unreachable!();
-                };
-
-                builder.push(chunk.into_cursor()?);
-            }
-            let merger = builder.build();
-
-            write_entries_into_database(
-                merger,
+            append_entries_into_database(
+                word_position_docids_iter,
                 &index.word_position_docids,
                 wtxn,
+                index_is_empty,
                 deladd_serialize_add_side,
                 merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap,
             )?;
             is_merged_database = true;
         }
-        TypedChunk::FieldIdFacetNumberDocids(_) => {
+        TypedChunk::FieldIdFacetNumberDocids(facet_id_number_docids_iter) => {
             let span =
                 tracing::trace_span!(target: "indexing::write_db","field_id_facet_number_docids");
             let _entered = span.enter();
-
-            let mut builder = MergerBuilder::new(merge_deladd_cbo_roaring_bitmaps as MergeFn);
-            let mut data_size = 0;
-            for typed_chunk in typed_chunks {
-                let TypedChunk::FieldIdFacetNumberDocids(facet_id_number_docids) = typed_chunk
-                else {
-                    unreachable!();
-                };
-
-                data_size += facet_id_number_docids.len();
-                builder.push(facet_id_number_docids.into_cursor()?);
-            }
-            let merger = builder.build();
-
-            let indexer = FacetsUpdate::new(index, FacetType::Number, merger, None, data_size);
+            let indexer = FacetsUpdate::new(index, FacetType::Number, facet_id_number_docids_iter);
             indexer.execute(wtxn)?;
             is_merged_database = true;
         }
-        TypedChunk::FieldIdFacetStringDocids(_) => {
+        TypedChunk::FieldIdFacetStringDocids(facet_id_string_docids_iter) => {
             let span =
                 tracing::trace_span!(target: "indexing::write_db", "field_id_facet_string_docids");
             let _entered = span.enter();
-
-            let mut facet_id_string_builder =
-                MergerBuilder::new(merge_deladd_cbo_roaring_bitmaps as MergeFn);
-            let mut normalized_facet_id_string_builder =
-                MergerBuilder::new(merge_deladd_btreeset_string as MergeFn);
-            let mut data_size = 0;
-            for typed_chunk in typed_chunks {
-                let TypedChunk::FieldIdFacetStringDocids((
-                    facet_id_string_docids,
-                    normalized_facet_id_string_docids,
-                )) = typed_chunk
-                else {
-                    unreachable!();
-                };
-
-                data_size += facet_id_string_docids.len();
-                facet_id_string_builder.push(facet_id_string_docids.into_cursor()?);
-                normalized_facet_id_string_builder
-                    .push(normalized_facet_id_string_docids.into_cursor()?);
-            }
-            let facet_id_string_merger = facet_id_string_builder.build();
-            let normalized_facet_id_string_merger = normalized_facet_id_string_builder.build();
-
-            let indexer = FacetsUpdate::new(
-                index,
-                FacetType::String,
-                facet_id_string_merger,
-                Some(normalized_facet_id_string_merger),
-                data_size,
-            );
+            let indexer = FacetsUpdate::new(index, FacetType::String, facet_id_string_docids_iter);
             indexer.execute(wtxn)?;
             is_merged_database = true;
         }
-        TypedChunk::FieldIdFacetExistsDocids(_) => {
+        TypedChunk::FieldIdFacetExistsDocids(facet_id_exists_docids) => {
             let span =
                 tracing::trace_span!(target: "indexing::write_db", "field_id_facet_exists_docids");
             let _entered = span.enter();
-
-            let mut builder = MergerBuilder::new(merge_deladd_cbo_roaring_bitmaps as MergeFn);
-            for typed_chunk in typed_chunks {
-                let TypedChunk::FieldIdFacetExistsDocids(chunk) = typed_chunk else {
-                    unreachable!();
-                };
-
-                builder.push(chunk.into_cursor()?);
-            }
-            let merger = builder.build();
-
-            write_entries_into_database(
-                merger,
+            append_entries_into_database(
+                facet_id_exists_docids,
                 &index.facet_id_exists_docids,
                 wtxn,
+                index_is_empty,
                 deladd_serialize_add_side,
                 merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap,
             )?;
             is_merged_database = true;
         }
-        TypedChunk::FieldIdFacetIsNullDocids(_) => {
+        TypedChunk::FieldIdFacetIsNullDocids(facet_id_is_null_docids) => {
             let span =
                 tracing::trace_span!(target: "indexing::write_db", "field_id_facet_is_null_docids");
             let _entered = span.enter();
-
-            let mut builder = MergerBuilder::new(merge_deladd_cbo_roaring_bitmaps as MergeFn);
-            for typed_chunk in typed_chunks {
-                let TypedChunk::FieldIdFacetIsNullDocids(chunk) = typed_chunk else {
-                    unreachable!();
-                };
-
-                builder.push(chunk.into_cursor()?);
-            }
-            let merger = builder.build();
-
-            write_entries_into_database(
-                merger,
+            append_entries_into_database(
+                facet_id_is_null_docids,
                 &index.facet_id_is_null_docids,
                 wtxn,
+                index_is_empty,
                 deladd_serialize_add_side,
                 merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap,
             )?;
             is_merged_database = true;
         }
-        TypedChunk::FieldIdFacetIsEmptyDocids(_) => {
+        TypedChunk::FieldIdFacetIsEmptyDocids(facet_id_is_empty_docids) => {
             let span = tracing::trace_span!(target: "profile::indexing::write_db", "field_id_facet_is_empty_docids");
             let _entered = span.enter();
-
-            let mut builder = MergerBuilder::new(merge_deladd_cbo_roaring_bitmaps as MergeFn);
-            for typed_chunk in typed_chunks {
-                let TypedChunk::FieldIdFacetIsEmptyDocids(chunk) = typed_chunk else {
-                    unreachable!();
-                };
-
-                builder.push(chunk.into_cursor()?);
-            }
-            let merger = builder.build();
-
-            write_entries_into_database(
-                merger,
+            append_entries_into_database(
+                facet_id_is_empty_docids,
                 &index.facet_id_is_empty_docids,
                 wtxn,
+                index_is_empty,
                 deladd_serialize_add_side,
                 merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap,
             )?;
             is_merged_database = true;
         }
-        TypedChunk::WordPairProximityDocids(_) => {
+        TypedChunk::WordPairProximityDocids(word_pair_proximity_docids_iter) => {
             let span =
                 tracing::trace_span!(target: "indexing::write_db", "word_pair_proximity_docids");
             let _entered = span.enter();
-
-            let mut builder = MergerBuilder::new(merge_deladd_cbo_roaring_bitmaps as MergeFn);
-            for typed_chunk in typed_chunks {
-                let TypedChunk::WordPairProximityDocids(chunk) = typed_chunk else {
-                    unreachable!();
-                };
-
-                builder.push(chunk.into_cursor()?);
-            }
-            let merger = builder.build();
-
-            write_entries_into_database(
-                merger,
+            append_entries_into_database(
+                word_pair_proximity_docids_iter,
                 &index.word_pair_proximity_docids,
                 wtxn,
+                index_is_empty,
                 deladd_serialize_add_side,
                 merge_deladd_cbo_roaring_bitmaps_into_cbo_roaring_bitmap,
             )?;
             is_merged_database = true;
         }
-        TypedChunk::FieldIdDocidFacetNumbers(_) => {
+        TypedChunk::FieldIdDocidFacetNumbers(fid_docid_facet_number) => {
             let span =
                 tracing::trace_span!(target: "indexing::write_db", "field_id_docid_facet_numbers");
             let _entered = span.enter();
-
-            let mut builder = MergerBuilder::new(keep_first as MergeFn);
-            for typed_chunk in typed_chunks {
-                let TypedChunk::FieldIdDocidFacetNumbers(chunk) = typed_chunk else {
-                    unreachable!();
-                };
-
-                builder.push(chunk.into_cursor()?);
-            }
-            let merger = builder.build();
-
             let index_fid_docid_facet_numbers =
                 index.field_id_docid_facet_f64s.remap_types::<Bytes, Bytes>();
-            let mut iter = merger.into_stream_merger_iter()?;
-            while let Some((key, value)) = iter.next()? {
+            let mut cursor = fid_docid_facet_number.into_cursor()?;
+            while let Some((key, value)) = cursor.move_on_next()? {
                 let reader = KvReaderDelAdd::new(value);
                 if valid_lmdb_key(key) {
                     match (reader.get(DelAdd::Deletion), reader.get(DelAdd::Addition)) {
@@ -553,25 +344,14 @@ pub(crate) fn write_typed_chunk_into_index(
                 }
             }
         }
-        TypedChunk::FieldIdDocidFacetStrings(_) => {
+        TypedChunk::FieldIdDocidFacetStrings(fid_docid_facet_string) => {
             let span =
                 tracing::trace_span!(target: "indexing::write_db", "field_id_docid_facet_strings");
             let _entered = span.enter();
-
-            let mut builder = MergerBuilder::new(keep_first as MergeFn);
-            for typed_chunk in typed_chunks {
-                let TypedChunk::FieldIdDocidFacetStrings(chunk) = typed_chunk else {
-                    unreachable!();
-                };
-
-                builder.push(chunk.into_cursor()?);
-            }
-            let merger = builder.build();
-
             let index_fid_docid_facet_strings =
                 index.field_id_docid_facet_strings.remap_types::<Bytes, Bytes>();
-            let mut iter = merger.into_stream_merger_iter()?;
-            while let Some((key, value)) = iter.next()? {
+            let mut cursor = fid_docid_facet_string.into_cursor()?;
+            while let Some((key, value)) = cursor.move_on_next()? {
                 let reader = KvReaderDelAdd::new(value);
                 if valid_lmdb_key(key) {
                     match (reader.get(DelAdd::Deletion), reader.get(DelAdd::Addition)) {
@@ -587,25 +367,14 @@ pub(crate) fn write_typed_chunk_into_index(
                 }
             }
         }
-        TypedChunk::GeoPoints(_) => {
+        TypedChunk::GeoPoints(geo_points) => {
             let span = tracing::trace_span!(target: "indexing::write_db", "geo_points");
             let _entered = span.enter();
-
-            let mut builder = MergerBuilder::new(keep_first as MergeFn);
-            for typed_chunk in typed_chunks {
-                let TypedChunk::GeoPoints(chunk) = typed_chunk else {
-                    unreachable!();
-                };
-
-                builder.push(chunk.into_cursor()?);
-            }
-            let merger = builder.build();
-
             let mut rtree = index.geo_rtree(wtxn)?.unwrap_or_default();
             let mut geo_faceted_docids = index.geo_faceted_documents_ids(wtxn)?;
 
-            let mut iter = merger.into_stream_merger_iter()?;
-            while let Some((key, value)) = iter.next()? {
+            let mut cursor = geo_points.into_cursor()?;
+            while let Some((key, value)) = cursor.move_on_next()? {
                 // convert the key back to a u32 (4 bytes)
                 let docid = key.try_into().map(DocumentId::from_be_bytes).unwrap();
 
@@ -624,38 +393,15 @@ pub(crate) fn write_typed_chunk_into_index(
             index.put_geo_rtree(wtxn, &rtree)?;
             index.put_geo_faceted_documents_ids(wtxn, &geo_faceted_docids)?;
         }
-        TypedChunk::VectorPoints { .. } => {
+        TypedChunk::VectorPoints {
+            remove_vectors,
+            manual_vectors,
+            embeddings,
+            expected_dimension,
+            embedder_name,
+        } => {
             let span = tracing::trace_span!(target: "indexing::write_db", "vector_points");
             let _entered = span.enter();
-
-            let mut remove_vectors_builder = MergerBuilder::new(keep_first as MergeFn);
-            let mut manual_vectors_builder = MergerBuilder::new(keep_first as MergeFn);
-            let mut embeddings_builder = MergerBuilder::new(keep_first as MergeFn);
-            let mut params = None;
-            for typed_chunk in typed_chunks {
-                let TypedChunk::VectorPoints {
-                    remove_vectors,
-                    manual_vectors,
-                    embeddings,
-                    expected_dimension,
-                    embedder_name,
-                } = typed_chunk
-                else {
-                    unreachable!();
-                };
-
-                params = Some((expected_dimension, embedder_name));
-
-                remove_vectors_builder.push(remove_vectors.into_cursor()?);
-                manual_vectors_builder.push(manual_vectors.into_cursor()?);
-                if let Some(embeddings) = embeddings {
-                    embeddings_builder.push(embeddings.into_cursor()?);
-                }
-            }
-
-            // typed chunks has always at least 1 chunk.
-            let Some((expected_dimension, embedder_name)) = params else { unreachable!() };
-
             let embedder_index = index.embedder_category_id.get(wtxn, &embedder_name)?.ok_or(
                 InternalError::DatabaseMissingEntry { db_name: "embedder_category_id", key: None },
             )?;
@@ -673,9 +419,8 @@ pub(crate) fn write_typed_chunk_into_index(
             let writers = writers?;
 
             // remove vectors for docids we want them removed
-            let merger = remove_vectors_builder.build();
-            let mut iter = merger.into_stream_merger_iter()?;
-            while let Some((key, _)) = iter.next()? {
+            let mut cursor = remove_vectors.into_cursor()?;
+            while let Some((key, _)) = cursor.move_on_next()? {
                 let docid = key.try_into().map(DocumentId::from_be_bytes).unwrap();
 
                 for writer in &writers {
@@ -687,39 +432,40 @@ pub(crate) fn write_typed_chunk_into_index(
             }
 
             // add generated embeddings
-            let merger = embeddings_builder.build();
-            let mut iter = merger.into_stream_merger_iter()?;
-            while let Some((key, value)) = iter.next()? {
-                let docid = key.try_into().map(DocumentId::from_be_bytes).unwrap();
-                let data = pod_collect_to_vec(value);
-                // it is a code error to have embeddings and not expected_dimension
-                let embeddings = crate::vector::Embeddings::from_inner(data, expected_dimension)
-                    // code error if we somehow got the wrong dimension
-                    .unwrap();
+            if let Some(embeddings) = embeddings {
+                let mut cursor = embeddings.into_cursor()?;
+                while let Some((key, value)) = cursor.move_on_next()? {
+                    let docid = key.try_into().map(DocumentId::from_be_bytes).unwrap();
+                    let data = pod_collect_to_vec(value);
+                    // it is a code error to have embeddings and not expected_dimension
+                    let embeddings =
+                        crate::vector::Embeddings::from_inner(data, expected_dimension)
+                            // code error if we somehow got the wrong dimension
+                            .unwrap();
 
-                if embeddings.embedding_count() > usize::from(u8::MAX) {
-                    let external_docid = if let Ok(Some(Ok(index))) = index
-                        .external_id_of(wtxn, std::iter::once(docid))
-                        .map(|it| it.into_iter().next())
-                    {
-                        index
-                    } else {
-                        format!("internal docid={docid}")
-                    };
-                    return Err(crate::Error::UserError(crate::UserError::TooManyVectors(
-                        external_docid,
-                        embeddings.embedding_count(),
-                    )));
-                }
-                for (embedding, writer) in embeddings.iter().zip(&writers) {
-                    writer.add_item(wtxn, docid, embedding)?;
+                    if embeddings.embedding_count() > usize::from(u8::MAX) {
+                        let external_docid = if let Ok(Some(Ok(index))) = index
+                            .external_id_of(wtxn, std::iter::once(docid))
+                            .map(|it| it.into_iter().next())
+                        {
+                            index
+                        } else {
+                            format!("internal docid={docid}")
+                        };
+                        return Err(crate::Error::UserError(crate::UserError::TooManyVectors(
+                            external_docid,
+                            embeddings.embedding_count(),
+                        )));
+                    }
+                    for (embedding, writer) in embeddings.iter().zip(&writers) {
+                        writer.add_item(wtxn, docid, embedding)?;
+                    }
                 }
             }
 
             // perform the manual diff
-            let merger = manual_vectors_builder.build();
-            let mut iter = merger.into_stream_merger_iter()?;
-            while let Some((key, value)) = iter.next()? {
+            let mut cursor = manual_vectors.into_cursor()?;
+            while let Some((key, value)) = cursor.move_on_next()? {
                 // convert the key back to a u32 (4 bytes)
                 let (left, _index) = try_split_array_at(key).unwrap();
                 let docid = DocumentId::from_be_bytes(left);
@@ -773,30 +519,26 @@ pub(crate) fn write_typed_chunk_into_index(
 
             tracing::debug!("Finished vector chunk for {}", embedder_name);
         }
-        TypedChunk::ScriptLanguageDocids(_) => {
+        TypedChunk::ScriptLanguageDocids(sl_map) => {
             let span = tracing::trace_span!(target: "indexing::write_db", "script_language_docids");
             let _entered = span.enter();
-
-            for typed_chunk in typed_chunks {
-                let TypedChunk::ScriptLanguageDocids(sl_map) = typed_chunk else { unreachable!() };
-                for (key, (deletion, addition)) in sl_map {
-                    let mut db_key_exists = false;
-                    let final_value = match index.script_language_docids.get(wtxn, &key)? {
-                        Some(db_values) => {
-                            db_key_exists = true;
-                            (db_values - deletion) | addition
-                        }
-                        None => addition,
-                    };
-
-                    if final_value.is_empty() {
-                        // If the database entry exists, delete it.
-                        if db_key_exists {
-                            index.script_language_docids.delete(wtxn, &key)?;
-                        }
-                    } else {
-                        index.script_language_docids.put(wtxn, &key, &final_value)?;
+            for (key, (deletion, addition)) in sl_map {
+                let mut db_key_exists = false;
+                let final_value = match index.script_language_docids.get(wtxn, &key)? {
+                    Some(db_values) => {
+                        db_key_exists = true;
+                        (db_values - deletion) | addition
                     }
+                    None => addition,
+                };
+
+                if final_value.is_empty() {
+                    // If the database entry exists, delete it.
+                    if db_key_exists {
+                        index.script_language_docids.delete(wtxn, &key)?;
+                    }
+                } else {
+                    index.script_language_docids.put(wtxn, &key, &final_value)?;
                 }
             }
         }
@@ -815,9 +557,13 @@ fn extract_geo_point(value: &[u8], docid: DocumentId) -> GeoPoint {
 }
 
 fn merge_word_docids_reader_into_fst(
-    merger: Merger<CursorClonableMmap, MergeFn>,
+    word_docids_iter: grenad::Reader<io::Cursor<ClonableMmap>>,
+    exact_word_docids_iter: grenad::Reader<io::Cursor<ClonableMmap>>,
 ) -> Result<fst::Set<Vec<u8>>> {
-    let mut iter = merger.into_stream_merger_iter()?;
+    let mut merger_builder = MergerBuilder::new(merge_ignore_values as MergeFn);
+    merger_builder.push(word_docids_iter.into_cursor()?);
+    merger_builder.push(exact_word_docids_iter.into_cursor()?);
+    let mut iter = merger_builder.build().into_stream_merger_iter()?;
     let mut builder = fst::SetBuilder::memory();
 
     while let Some((k, _)) = iter.next()? {
@@ -831,9 +577,10 @@ fn merge_word_docids_reader_into_fst(
 /// merge_values function is used if an entry already exist in the database.
 #[tracing::instrument(level = "trace", skip_all, target = "indexing::write_db")]
 fn write_entries_into_database<R, K, V, FS, FM>(
-    merger: Merger<R, MergeFn>,
+    data: grenad::Reader<R>,
     database: &heed::Database<K, V>,
     wtxn: &mut RwTxn,
+    index_is_empty: bool,
     serialize_value: FS,
     merge_values: FM,
 ) -> Result<()>
@@ -842,17 +589,22 @@ where
     FS: for<'a> Fn(&'a [u8], &'a mut Vec<u8>) -> Result<&'a [u8]>,
     FM: for<'a> Fn(&[u8], &[u8], &'a mut Vec<u8>) -> Result<Option<&'a [u8]>>,
 {
-    puffin::profile_function!();
+    puffin::profile_function!(format!("number of entries: {}", data.len()));
+
     let mut buffer = Vec::new();
     let database = database.remap_types::<Bytes, Bytes>();
 
-    let mut iter = merger.into_stream_merger_iter()?;
-    while let Some((key, value)) = iter.next()? {
+    let mut cursor = data.into_cursor()?;
+    while let Some((key, value)) = cursor.move_on_next()? {
         if valid_lmdb_key(key) {
             buffer.clear();
-            let value = match database.get(wtxn, key)? {
-                Some(prev_value) => merge_values(value, prev_value, &mut buffer)?,
-                None => Some(serialize_value(value, &mut buffer)?),
+            let value = if index_is_empty {
+                Some(serialize_value(value, &mut buffer)?)
+            } else {
+                match database.get(wtxn, key)? {
+                    Some(prev_value) => merge_values(value, prev_value, &mut buffer)?,
+                    None => Some(serialize_value(value, &mut buffer)?),
+                }
             };
             match value {
                 Some(value) => database.put(wtxn, key, value)?,
@@ -862,5 +614,62 @@ where
             }
         }
     }
+
+    Ok(())
+}
+
+/// Write provided entries in database using serialize_value function.
+/// merge_values function is used if an entry already exist in the database.
+/// All provided entries must be ordered.
+/// If the index is not empty, write_entries_into_database is called instead.
+#[tracing::instrument(level = "trace", skip_all, target = "indexing::write_db")]
+fn append_entries_into_database<R, K, V, FS, FM>(
+    data: grenad::Reader<R>,
+    database: &heed::Database<K, V>,
+    wtxn: &mut RwTxn,
+    index_is_empty: bool,
+    serialize_value: FS,
+    merge_values: FM,
+) -> Result<()>
+where
+    R: io::Read + io::Seek,
+    FS: for<'a> Fn(&'a [u8], &'a mut Vec<u8>) -> Result<&'a [u8]>,
+    FM: for<'a> Fn(&[u8], &[u8], &'a mut Vec<u8>) -> Result<Option<&'a [u8]>>,
+    K: for<'a> heed::BytesDecode<'a>,
+{
+    puffin::profile_function!(format!("number of entries: {}", data.len()));
+
+    if !index_is_empty {
+        return write_entries_into_database(
+            data,
+            database,
+            wtxn,
+            false,
+            serialize_value,
+            merge_values,
+        );
+    }
+
+    let mut buffer = Vec::new();
+    let mut database = database.iter_mut(wtxn)?.remap_types::<Bytes, Bytes>();
+
+    let mut cursor = data.into_cursor()?;
+    while let Some((key, value)) = cursor.move_on_next()? {
+        if valid_lmdb_key(key) {
+            debug_assert!(
+                K::bytes_decode(key).is_ok(),
+                "Couldn't decode key with the database decoder, key length: {} - key bytes: {:x?}",
+                key.len(),
+                &key
+            );
+            buffer.clear();
+            let value = serialize_value(value, &mut buffer)?;
+            unsafe {
+                // safety: We do not keep a reference to anything that lives inside the database
+                database.put_current_with_options::<Bytes>(PutFlags::APPEND, key, value)?
+            };
+        }
+    }
+
     Ok(())
 }

milli/src/update/mod.rs

@@ -3,8 +3,9 @@ pub use self::clear_documents::ClearDocuments;
 pub use self::facet::bulk::FacetsUpdateBulk;
 pub use self::facet::incremental::FacetsUpdateIncrementalInner;
 pub use self::index_documents::{
-    merge_cbo_roaring_bitmaps, merge_roaring_bitmaps, DocumentAdditionResult, DocumentId,
-    IndexDocuments, IndexDocumentsConfig, IndexDocumentsMethod, MergeFn,
+    merge_btreeset_string, merge_cbo_roaring_bitmaps, merge_roaring_bitmaps,
+    DocumentAdditionResult, DocumentId, IndexDocuments, IndexDocumentsConfig, IndexDocumentsMethod,
+    MergeFn,
 };
 pub use self::indexer_config::IndexerConfig;
 pub use self::settings::{validate_embedding_settings, Setting, Settings};

milli/src/update/settings.rs

@@ -979,9 +979,6 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> {
                             crate::vector::settings::EmbeddingSettings::apply_default_source(
                                 &mut setting,
                             );
-                            crate::vector::settings::EmbeddingSettings::apply_default_openai_model(
-                                &mut setting,
-                            );
                             let setting = validate_embedding_settings(setting, &name)?;
                             changed = true;
                             new_configs.insert(name, setting);
@@ -1127,14 +1124,6 @@ pub fn validate_embedding_settings(
     let Setting::Set(settings) = settings else { return Ok(settings) };
     let EmbeddingSettings { source, model, revision, api_key, dimensions, document_template } =
         settings;
-
-    if let Some(0) = dimensions.set() {
-        return Err(crate::error::UserError::InvalidSettingsDimensions {
-            embedder_name: name.to_owned(),
-        }
-        .into());
-    }
-
     let Some(inferred_source) = source.set() else {
         return Ok(Setting::Set(EmbeddingSettings {
             source,
@@ -1148,34 +1137,14 @@ pub fn validate_embedding_settings(
     match inferred_source {
         EmbedderSource::OpenAi => {
             check_unset(&revision, "revision", inferred_source, name)?;
+            check_unset(&dimensions, "dimensions", inferred_source, name)?;
             if let Setting::Set(model) = &model {
-                let model = crate::vector::openai::EmbeddingModel::from_name(model.as_str())
-                    .ok_or(crate::error::UserError::InvalidOpenAiModel {
+                crate::vector::openai::EmbeddingModel::from_name(model.as_str()).ok_or(
+                    crate::error::UserError::InvalidOpenAiModel {
                         embedder_name: name.to_owned(),
                         model: model.clone(),
-                    })?;
-                if let Setting::Set(dimensions) = dimensions {
-                    if !model.supports_overriding_dimensions()
-                        && dimensions != model.default_dimensions()
-                    {
-                        return Err(crate::error::UserError::InvalidOpenAiModelDimensions {
-                            embedder_name: name.to_owned(),
-                            model: model.name(),
-                            dimensions,
-                            expected_dimensions: model.default_dimensions(),
-                        }
-                        .into());
-                    }
-                    if dimensions > model.default_dimensions() {
-                        return Err(crate::error::UserError::InvalidOpenAiModelDimensionsMax {
-                            embedder_name: name.to_owned(),
-                            model: model.name(),
-                            dimensions,
-                            max_dimensions: model.default_dimensions(),
-                        }
-                        .into());
-                    }
-                }
+                    },
+                )?;
             }
         }
         EmbedderSource::HuggingFace => {

milli/src/update/word_prefix_docids.rs

@@ -47,7 +47,7 @@ impl<'t, 'i> WordPrefixDocids<'t, 'i> {
     )]
     pub fn execute(
         self,
-        new_word_docids: grenad::Merger<CursorClonableMmap, MergeFn>,
+        mut new_word_docids_iter: grenad::ReaderCursor<CursorClonableMmap>,
         new_prefix_fst_words: &[String],
         common_prefix_fst_words: &[&[String]],
         del_prefix_fst_words: &HashSet<Vec<u8>>,
@@ -68,8 +68,7 @@ impl<'t, 'i> WordPrefixDocids<'t, 'i> {
         if !common_prefix_fst_words.is_empty() {
             let mut current_prefixes: Option<&&[String]> = None;
             let mut prefixes_cache = HashMap::new();
-            let mut new_word_docids_iter = new_word_docids.into_stream_merger_iter()?;
-            while let Some((word, data)) = new_word_docids_iter.next()? {
+            while let Some((word, data)) = new_word_docids_iter.move_on_next()? {
                 current_prefixes = match current_prefixes.take() {
                     Some(prefixes) if word.starts_with(prefixes[0].as_bytes()) => Some(prefixes),
                     _otherwise => {

milli/src/update/words_prefix_integer_docids.rs

@@ -52,7 +52,7 @@ impl<'t, 'i> WordPrefixIntegerDocids<'t, 'i> {
     )]
     pub fn execute(
         self,
-        new_word_integer_docids: grenad::Merger<CursorClonableMmap, MergeFn>,
+        new_word_integer_docids: grenad::Reader<CursorClonableMmap>,
         new_prefix_fst_words: &[String],
         common_prefix_fst_words: &[&[String]],
         del_prefix_fst_words: &HashSet<Vec<u8>>,
@@ -69,14 +69,14 @@ impl<'t, 'i> WordPrefixIntegerDocids<'t, 'i> {
             self.max_memory,
         );
 
+        let mut new_word_integer_docids_iter = new_word_integer_docids.into_cursor()?;
+
         if !common_prefix_fst_words.is_empty() {
             // We fetch all the new common prefixes between the previous and new prefix fst.
             let mut buffer = Vec::new();
             let mut current_prefixes: Option<&&[String]> = None;
             let mut prefixes_cache = HashMap::new();
-            let mut new_word_integer_docids_iter =
-                new_word_integer_docids.into_stream_merger_iter()?;
-            while let Some((key, data)) = new_word_integer_docids_iter.next()? {
+            while let Some((key, data)) = new_word_integer_docids_iter.move_on_next()? {
                 let (word, pos) =
                     StrBEU16Codec::bytes_decode(key).map_err(heed::Error::Decoding)?;
 

milli/src/vector/error.rs

@@ -67,10 +67,6 @@ pub enum EmbedErrorKind {
     OpenAiUnhandledStatusCode(u16),
     #[error("attempt to embed the following text in a configuration where embeddings must be user provided: {0:?}")]
     ManualEmbed(String),
-    #[error("could not initialize asynchronous runtime: {0}")]
-    OpenAiRuntimeInit(std::io::Error),
-    #[error("initializing web client for sending embedding requests failed: {0}")]
-    InitWebClient(reqwest::Error),
 }
 
 impl EmbedError {
@@ -121,14 +117,6 @@ impl EmbedError {
     pub(crate) fn embed_on_manual_embedder(texts: String) -> EmbedError {
         Self { kind: EmbedErrorKind::ManualEmbed(texts), fault: FaultSource::User }
     }
-
-    pub(crate) fn openai_runtime_init(inner: std::io::Error) -> EmbedError {
-        Self { kind: EmbedErrorKind::OpenAiRuntimeInit(inner), fault: FaultSource::Runtime }
-    }
-
-    pub fn openai_initialize_web_client(inner: reqwest::Error) -> Self {
-        Self { kind: EmbedErrorKind::InitWebClient(inner), fault: FaultSource::Runtime }
-    }
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -195,6 +183,10 @@ impl NewEmbedderError {
         }
     }
 
+    pub fn openai_initialize_web_client(inner: reqwest::Error) -> Self {
+        Self { kind: NewEmbedderErrorKind::InitWebClient(inner), fault: FaultSource::Runtime }
+    }
+
     pub fn openai_invalid_api_key_format(inner: reqwest::header::InvalidHeaderValue) -> Self {
         Self { kind: NewEmbedderErrorKind::InvalidApiKeyFormat(inner), fault: FaultSource::User }
     }
@@ -245,6 +237,8 @@ pub enum NewEmbedderErrorKind {
     #[error("loading model failed: {0}")]
     LoadModel(candle_core::Error),
     // openai
+    #[error("initializing web client for sending embedding requests failed: {0}")]
+    InitWebClient(reqwest::Error),
     #[error("The API key passed to Authorization error was in an invalid format: {0}")]
     InvalidApiKeyFormat(reqwest::header::InvalidHeaderValue),
 }

milli/src/vector/hf.rs

@@ -151,8 +151,7 @@ impl Embedder {
         let token_ids = tokens
             .iter()
             .map(|tokens| {
-                let mut tokens = tokens.get_ids().to_vec();
-                tokens.truncate(512);
+                let tokens = tokens.get_ids().to_vec();
                 Tensor::new(tokens.as_slice(), &self.model.device).map_err(EmbedError::tensor_shape)
             })
             .collect::<Result<Vec<_>, EmbedError>>()?;

milli/src/vector/mod.rs

@@ -163,24 +163,18 @@ impl Embedder {
     ) -> std::result::Result<Vec<Embeddings<f32>>, EmbedError> {
         match self {
             Embedder::HuggingFace(embedder) => embedder.embed(texts),
-            Embedder::OpenAi(embedder) => {
-                let client = embedder.new_client()?;
-                embedder.embed(texts, &client).await
-            }
+            Embedder::OpenAi(embedder) => embedder.embed(texts).await,
             Embedder::UserProvided(embedder) => embedder.embed(texts),
         }
     }
 
-    /// # Panics
-    ///
-    /// - if called from an asynchronous context
-    pub fn embed_chunks(
+    pub async fn embed_chunks(
         &self,
         text_chunks: Vec<Vec<String>>,
     ) -> std::result::Result<Vec<Vec<Embeddings<f32>>>, EmbedError> {
         match self {
             Embedder::HuggingFace(embedder) => embedder.embed_chunks(text_chunks),
-            Embedder::OpenAi(embedder) => embedder.embed_chunks(text_chunks),
+            Embedder::OpenAi(embedder) => embedder.embed_chunks(text_chunks).await,
             Embedder::UserProvided(embedder) => embedder.embed_chunks(text_chunks),
         }
     }

milli/src/vector/openai.rs

@@ -8,7 +8,7 @@ use super::{DistributionShift, Embedding, Embeddings};
 
 #[derive(Debug)]
 pub struct Embedder {
-    headers: reqwest::header::HeaderMap,
+    client: reqwest::Client,
     tokenizer: tiktoken_rs::CoreBPE,
     options: EmbedderOptions,
 }
@@ -17,7 +17,6 @@ pub struct Embedder {
 pub struct EmbedderOptions {
     pub api_key: Option<String>,
     pub embedding_model: EmbeddingModel,
-    pub dimensions: Option<usize>,
 }
 
 #[derive(
@@ -42,50 +41,34 @@ pub enum EmbeddingModel {
     #[serde(rename = "text-embedding-ada-002")]
     #[deserr(rename = "text-embedding-ada-002")]
     TextEmbeddingAda002,
-
-    #[serde(rename = "text-embedding-3-small")]
-    #[deserr(rename = "text-embedding-3-small")]
-    TextEmbedding3Small,
-
-    #[serde(rename = "text-embedding-3-large")]
-    #[deserr(rename = "text-embedding-3-large")]
-    TextEmbedding3Large,
 }
 
 impl EmbeddingModel {
     pub fn supported_models() -> &'static [&'static str] {
-        &["text-embedding-ada-002", "text-embedding-3-small", "text-embedding-3-large"]
+        &["text-embedding-ada-002"]
     }
 
     pub fn max_token(&self) -> usize {
         match self {
             EmbeddingModel::TextEmbeddingAda002 => 8191,
-            EmbeddingModel::TextEmbedding3Large => 8191,
-            EmbeddingModel::TextEmbedding3Small => 8191,
         }
     }
 
-    pub fn default_dimensions(&self) -> usize {
+    pub fn dimensions(&self) -> usize {
         match self {
             EmbeddingModel::TextEmbeddingAda002 => 1536,
-            EmbeddingModel::TextEmbedding3Large => 3072,
-            EmbeddingModel::TextEmbedding3Small => 1536,
         }
     }
 
     pub fn name(&self) -> &'static str {
         match self {
             EmbeddingModel::TextEmbeddingAda002 => "text-embedding-ada-002",
-            EmbeddingModel::TextEmbedding3Large => "text-embedding-3-large",
-            EmbeddingModel::TextEmbedding3Small => "text-embedding-3-small",
         }
     }
 
     pub fn from_name(name: &str) -> Option<Self> {
         match name {
             "text-embedding-ada-002" => Some(EmbeddingModel::TextEmbeddingAda002),
-            "text-embedding-3-large" => Some(EmbeddingModel::TextEmbedding3Large),
-            "text-embedding-3-small" => Some(EmbeddingModel::TextEmbedding3Small),
             _ => None,
         }
     }
@@ -95,20 +78,6 @@ impl EmbeddingModel {
             EmbeddingModel::TextEmbeddingAda002 => {
                 Some(DistributionShift { current_mean: 0.90, current_sigma: 0.08 })
             }
-            EmbeddingModel::TextEmbedding3Large => {
-                Some(DistributionShift { current_mean: 0.70, current_sigma: 0.1 })
-            }
-            EmbeddingModel::TextEmbedding3Small => {
-                Some(DistributionShift { current_mean: 0.75, current_sigma: 0.1 })
-            }
-        }
-    }
-
-    pub fn supports_overriding_dimensions(&self) -> bool {
-        match self {
-            EmbeddingModel::TextEmbeddingAda002 => false,
-            EmbeddingModel::TextEmbedding3Large => true,
-            EmbeddingModel::TextEmbedding3Small => true,
         }
     }
 }
@@ -117,22 +86,15 @@ pub const OPENAI_EMBEDDINGS_URL: &str = "https://api.openai.com/v1/embeddings";
 
 impl EmbedderOptions {
     pub fn with_default_model(api_key: Option<String>) -> Self {
-        Self { api_key, embedding_model: Default::default(), dimensions: None }
+        Self { api_key, embedding_model: Default::default() }
     }
 
     pub fn with_embedding_model(api_key: Option<String>, embedding_model: EmbeddingModel) -> Self {
-        Self { api_key, embedding_model, dimensions: None }
+        Self { api_key, embedding_model }
     }
 }
 
 impl Embedder {
-    pub fn new_client(&self) -> Result<reqwest::Client, EmbedError> {
-        reqwest::ClientBuilder::new()
-            .default_headers(self.headers.clone())
-            .build()
-            .map_err(EmbedError::openai_initialize_web_client)
-    }
-
     pub fn new(options: EmbedderOptions) -> Result<Self, NewEmbedderError> {
         let mut headers = reqwest::header::HeaderMap::new();
         let mut inferred_api_key = Default::default();
@@ -149,25 +111,25 @@ impl Embedder {
             reqwest::header::CONTENT_TYPE,
             reqwest::header::HeaderValue::from_static("application/json"),
         );
+        let client = reqwest::ClientBuilder::new()
+            .default_headers(headers)
+            .build()
+            .map_err(NewEmbedderError::openai_initialize_web_client)?;
 
         // looking at the code it is very unclear that this can actually fail.
         let tokenizer = tiktoken_rs::cl100k_base().unwrap();
 
-        Ok(Self { options, headers, tokenizer })
+        Ok(Self { options, client, tokenizer })
     }
 
-    pub async fn embed(
-        &self,
-        texts: Vec<String>,
-        client: &reqwest::Client,
-    ) -> Result<Vec<Embeddings<f32>>, EmbedError> {
+    pub async fn embed(&self, texts: Vec<String>) -> Result<Vec<Embeddings<f32>>, EmbedError> {
         let mut tokenized = false;
 
         for attempt in 0..7 {
             let result = if tokenized {
-                self.try_embed_tokenized(&texts, client).await
+                self.try_embed_tokenized(&texts).await
             } else {
-                self.try_embed(&texts, client).await
+                self.try_embed(&texts).await
             };
 
             let retry_duration = match result {
@@ -187,9 +149,9 @@ impl Embedder {
         }
 
         let result = if tokenized {
-            self.try_embed_tokenized(&texts, client).await
+            self.try_embed_tokenized(&texts).await
         } else {
-            self.try_embed(&texts, client).await
+            self.try_embed(&texts).await
         };
 
         result.map_err(Retry::into_error)
@@ -267,17 +229,13 @@ impl Embedder {
     async fn try_embed<S: AsRef<str> + serde::Serialize>(
         &self,
         texts: &[S],
-        client: &reqwest::Client,
     ) -> Result<Vec<Embeddings<f32>>, Retry> {
         for text in texts {
             tracing::trace!("Received prompt: {}", text.as_ref())
         }
-        let request = OpenAiRequest {
-            model: self.options.embedding_model.name(),
-            input: texts,
-            dimensions: self.overriden_dimensions(),
-        };
-        let response = client
+        let request = OpenAiRequest { model: self.options.embedding_model.name(), input: texts };
+        let response = self
+            .client
             .post(OPENAI_EMBEDDINGS_URL)
             .json(&request)
             .send()
@@ -302,11 +260,7 @@ impl Embedder {
             .collect())
     }
 
-    async fn try_embed_tokenized(
-        &self,
-        text: &[String],
-        client: &reqwest::Client,
-    ) -> Result<Vec<Embeddings<f32>>, Retry> {
+    async fn try_embed_tokenized(&self, text: &[String]) -> Result<Vec<Embeddings<f32>>, Retry> {
         pub const OVERLAP_SIZE: usize = 200;
         let mut all_embeddings = Vec::with_capacity(text.len());
         for text in text {
@@ -314,34 +268,31 @@ impl Embedder {
             let encoded = self.tokenizer.encode_ordinary(text.as_str());
             let len = encoded.len();
             if len < max_token_count {
-                all_embeddings.append(&mut self.try_embed(&[text], client).await?);
+                all_embeddings.append(&mut self.try_embed(&[text]).await?);
                 continue;
             }
 
             let mut tokens = encoded.as_slice();
-            let mut embeddings_for_prompt = Embeddings::new(self.dimensions());
+            let mut embeddings_for_prompt =
+                Embeddings::new(self.options.embedding_model.dimensions());
             while tokens.len() > max_token_count {
                 let window = &tokens[..max_token_count];
-                embeddings_for_prompt.push(self.embed_tokens(window, client).await?).unwrap();
+                embeddings_for_prompt.push(self.embed_tokens(window).await?).unwrap();
 
                 tokens = &tokens[max_token_count - OVERLAP_SIZE..];
             }
 
             // end of text
-            embeddings_for_prompt.push(self.embed_tokens(tokens, client).await?).unwrap();
+            embeddings_for_prompt.push(self.embed_tokens(tokens).await?).unwrap();
 
             all_embeddings.push(embeddings_for_prompt);
         }
         Ok(all_embeddings)
     }
 
-    async fn embed_tokens(
-        &self,
-        tokens: &[usize],
-        client: &reqwest::Client,
-    ) -> Result<Embedding, Retry> {
+    async fn embed_tokens(&self, tokens: &[usize]) -> Result<Embedding, Retry> {
         for attempt in 0..9 {
-            let duration = match self.try_embed_tokens(tokens, client).await {
+            let duration = match self.try_embed_tokens(tokens).await {
                 Ok(embedding) => return Ok(embedding),
                 Err(retry) => retry.into_duration(attempt),
             }
@@ -350,22 +301,14 @@ impl Embedder {
             tokio::time::sleep(duration).await;
         }
 
-        self.try_embed_tokens(tokens, client)
-            .await
-            .map_err(|retry| Retry::give_up(retry.into_error()))
+        self.try_embed_tokens(tokens).await.map_err(|retry| Retry::give_up(retry.into_error()))
     }
 
-    async fn try_embed_tokens(
-        &self,
-        tokens: &[usize],
-        client: &reqwest::Client,
-    ) -> Result<Embedding, Retry> {
-        let request = OpenAiTokensRequest {
-            model: self.options.embedding_model.name(),
-            input: tokens,
-            dimensions: self.overriden_dimensions(),
-        };
-        let response = client
+    async fn try_embed_tokens(&self, tokens: &[usize]) -> Result<Embedding, Retry> {
+        let request =
+            OpenAiTokensRequest { model: self.options.embedding_model.name(), input: tokens };
+        let response = self
+            .client
             .post(OPENAI_EMBEDDINGS_URL)
             .json(&request)
             .send()
@@ -383,19 +326,12 @@ impl Embedder {
         Ok(response.data.pop().map(|data| data.embedding).unwrap_or_default())
     }
 
-    pub fn embed_chunks(
+    pub async fn embed_chunks(
         &self,
         text_chunks: Vec<Vec<String>>,
     ) -> Result<Vec<Vec<Embeddings<f32>>>, EmbedError> {
-        let rt = tokio::runtime::Builder::new_current_thread()
-            .enable_io()
-            .enable_time()
-            .build()
-            .map_err(EmbedError::openai_runtime_init)?;
-        let client = self.new_client()?;
-        rt.block_on(futures::future::try_join_all(
-            text_chunks.into_iter().map(|prompts| self.embed(prompts, &client)),
-        ))
+        futures::future::try_join_all(text_chunks.into_iter().map(|prompts| self.embed(prompts)))
+            .await
     }
 
     pub fn chunk_count_hint(&self) -> usize {
@@ -407,24 +343,12 @@ impl Embedder {
     }
 
     pub fn dimensions(&self) -> usize {
-        if self.options.embedding_model.supports_overriding_dimensions() {
-            self.options.dimensions.unwrap_or(self.options.embedding_model.default_dimensions())
-        } else {
-            self.options.embedding_model.default_dimensions()
-        }
+        self.options.embedding_model.dimensions()
     }
 
     pub fn distribution(&self) -> Option<DistributionShift> {
         self.options.embedding_model.distribution()
     }
-
-    fn overriden_dimensions(&self) -> Option<usize> {
-        if self.options.embedding_model.supports_overriding_dimensions() {
-            self.options.dimensions
-        } else {
-            None
-        }
-    }
 }
 
 // retrying in case of failure
@@ -484,16 +408,12 @@ impl Retry {
 struct OpenAiRequest<'a, S: AsRef<str> + serde::Serialize> {
     model: &'a str,
     input: &'a [S],
-    #[serde(skip_serializing_if = "Option::is_none")]
-    dimensions: Option<usize>,
 }
 
 #[derive(Debug, Serialize)]
 struct OpenAiTokensRequest<'a> {
     model: &'a str,
     input: &'a [usize],
-    #[serde(skip_serializing_if = "Option::is_none")]
-    dimensions: Option<usize>,
 }
 
 #[derive(Debug, Deserialize)]

milli/src/vector/settings.rs

@@ -1,7 +1,6 @@
 use deserr::Deserr;
 use serde::{Deserialize, Serialize};
 
-use super::openai;
 use crate::prompt::PromptData;
 use crate::update::Setting;
 use crate::vector::EmbeddingConfig;
@@ -83,7 +82,7 @@ impl EmbeddingSettings {
             Self::MODEL => &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi],
             Self::REVISION => &[EmbedderSource::HuggingFace],
             Self::API_KEY => &[EmbedderSource::OpenAi],
-            Self::DIMENSIONS => &[EmbedderSource::OpenAi, EmbedderSource::UserProvided],
+            Self::DIMENSIONS => &[EmbedderSource::UserProvided],
             Self::DOCUMENT_TEMPLATE => &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi],
             _other => unreachable!("unknown field"),
         }
@@ -91,13 +90,9 @@ impl EmbeddingSettings {
 
     pub fn allowed_fields_for_source(source: EmbedderSource) -> &'static [&'static str] {
         match source {
-            EmbedderSource::OpenAi => &[
-                Self::SOURCE,
-                Self::MODEL,
-                Self::API_KEY,
-                Self::DOCUMENT_TEMPLATE,
-                Self::DIMENSIONS,
-            ],
+            EmbedderSource::OpenAi => {
+                &[Self::SOURCE, Self::MODEL, Self::API_KEY, Self::DOCUMENT_TEMPLATE]
+            }
             EmbedderSource::HuggingFace => {
                 &[Self::SOURCE, Self::MODEL, Self::REVISION, Self::DOCUMENT_TEMPLATE]
             }
@@ -114,17 +109,6 @@ impl EmbeddingSettings {
             *source = Setting::Set(EmbedderSource::default())
         }
     }
-
-    pub(crate) fn apply_default_openai_model(setting: &mut Setting<EmbeddingSettings>) {
-        if let Setting::Set(EmbeddingSettings {
-            source: Setting::Set(EmbedderSource::OpenAi),
-            model: model @ (Setting::NotSet | Setting::Reset),
-            ..
-        }) = setting
-        {
-            *model = Setting::Set(openai::EmbeddingModel::default().name().to_owned())
-        }
-    }
 }
 
 #[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq, Deserr)]
@@ -192,7 +176,7 @@ impl From<EmbeddingConfig> for EmbeddingSettings {
                 model: Setting::Set(options.embedding_model.name().to_owned()),
                 revision: Setting::NotSet,
                 api_key: options.api_key.map(Setting::Set).unwrap_or_default(),
-                dimensions: options.dimensions.map(Setting::Set).unwrap_or_default(),
+                dimensions: Setting::NotSet,
                 document_template: Setting::Set(prompt.template),
             },
             super::EmbedderOptions::UserProvided(options) => Self {
@@ -224,9 +208,6 @@ impl From<EmbeddingSettings> for EmbeddingConfig {
                     if let Some(api_key) = api_key.set() {
                         options.api_key = Some(api_key);
                     }
-                    if let Some(dimensions) = dimensions.set() {
-                        options.dimensions = Some(dimensions);
-                    }
                     this.embedder_options = super::EmbedderOptions::OpenAi(options);
                 }
                 EmbedderSource::HuggingFace => {

tracing-trace/Cargo.toml

@@ -13,11 +13,10 @@ serde_json = "1.0.111"
 tracing = "0.1.40"
 tracing-error = "0.2.0"
 tracing-subscriber = "0.3.18"
-byte-unit = { version = "4.0.19", default-features = false, features = [
-    "std",
-    "serde",
-] }
 tokio = { version = "1.35.1", features = ["sync"] }
+clap = { version = "4.4.18", features = ["derive"] }
+anyhow = "1.0.79"
+byte-unit = { version = "5.1.4", features = ["byte"] }
 
 [target.'cfg(any(target_os = "linux", target_os = "macos"))'.dependencies]
 libproc = "0.14.2"

tracing-trace/src/bin/filter-trace.rs

@@ -0,0 +1,103 @@
+use std::collections::vec_deque::Drain;
+use std::collections::VecDeque;
+use std::io::{self, BufReader, BufWriter, Stdout, Write};
+use std::mem;
+
+use anyhow::Context;
+use byte_unit::Byte;
+use clap::Parser;
+use tracing_trace::entry::{Entry, NewSpan};
+
+/// A program that filters trace logs to only keeps
+/// the logs related to memory usage above the given threshold.
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+struct Args {
+    /// The threshold that a log must have to be returned by this program.
+    #[arg(short, long)]
+    memory_threshold: Byte,
+
+    /// Number of context lines to keep around high memory log lines.
+    #[arg(long, default_value_t = 10)]
+    context: usize,
+}
+
+fn main() -> anyhow::Result<()> {
+    let Args { memory_threshold, context } = Args::parse();
+
+    let mut context = EntryContext::new(context);
+    let mut currently_in_threshold = false;
+
+    let input = BufReader::new(io::stdin());
+    let mut output = io::BufWriter::new(io::stdout());
+    for result in tracing_trace::TraceReader::new(input) {
+        let entry = result?;
+
+        match entry {
+            Entry::NewCallsite(_) | Entry::NewThread(_) => {
+                write_to_output(&mut output, &entry)?;
+            }
+            Entry::NewSpan(NewSpan { id, call_id, parent_id, thread_id }) => todo!(),
+            Entry::SpanEnter(_) => todo!(),
+            Entry::SpanExit(_) => todo!(),
+            Entry::SpanClose(_) => todo!(),
+            Entry::Event(_) => todo!(),
+        }
+
+        // if matches!(entry, Entry::NewCallsite(_) | Entry::NewThread(_)) {
+        //     write_to_output(&mut output, &entry)?;
+        // } else if entry.memory().map_or(true, |m| m.resident < memory_threshold.as_u64()) {
+        //     if mem::replace(&mut currently_in_threshold, false) {
+        //         for entry in context.drain() {
+        //             write_to_output(&mut output, &entry)?;
+        //         }
+        //     }
+        //     context.push(entry);
+        // } else {
+        //     currently_in_threshold = true;
+        //     for entry in context.drain() {
+        //         write_to_output(&mut output, &entry)?;
+        //     }
+        //     write_to_output(&mut output, &entry)?;
+        // }
+    }
+
+    for entry in context.drain() {
+        write_to_output(&mut output, &entry)?;
+    }
+
+    output.flush().context("flushing stdout")?;
+    Ok(())
+}
+
+fn write_to_output(writer: &mut BufWriter<Stdout>, entry: &Entry) -> anyhow::Result<()> {
+    serde_json::to_writer(writer, &entry).context("while serializing and writing to stdout")
+}
+
+/// Keeps only the last `size` element in memory.
+/// It's basically a sliding window.
+pub struct EntryContext {
+    size: usize,
+    queue: VecDeque<Entry>,
+}
+
+impl EntryContext {
+    pub fn new(size: usize) -> EntryContext {
+        EntryContext { size, queue: VecDeque::with_capacity(size) }
+    }
+
+    pub fn is_full(&self) -> bool {
+        self.size >= self.queue.len()
+    }
+
+    pub fn push(&mut self, entry: Entry) {
+        if self.queue.len() == self.size {
+            self.queue.pop_front();
+        }
+        self.queue.push_back(entry);
+    }
+
+    pub fn drain(&mut self) -> Drain<Entry> {
+        self.queue.drain(..)
+    }
+}

tracing-trace/src/entry.rs

@@ -38,6 +38,20 @@ pub enum Entry {
     Event(Event),
 }
 
+impl Entry {
+    pub fn memory(&self) -> Option<MemoryStats> {
+        match self {
+            Entry::NewCallsite(_)
+            | Entry::NewThread(_)
+            | Entry::NewSpan(_)
+            | Entry::SpanClose(_) => None,
+            Entry::SpanEnter(event) => event.memory,
+            Entry::SpanExit(event) => event.memory,
+            Entry::Event(event) => event.memory,
+        }
+    }
+}
+
 #[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq, Hash)]
 pub struct SpanId(u64);
 

tracing-trace/src/processor/fmt.rs

@@ -189,7 +189,7 @@ fn print_duration(duration: std::time::Duration) -> String {
 
 /// Format only the allocated bytes, deallocated bytes and reallocated bytes in GiB, MiB, KiB, Bytes.
 fn print_memory(MemoryStats { resident }: MemoryStats) -> String {
-    use byte_unit::Byte;
-    let rss_bytes = Byte::from_bytes(resident).get_appropriate_unit(true);
+    use byte_unit::{Byte, UnitType};
+    let rss_bytes = Byte::from_u64(resident).get_appropriate_unit(UnitType::Binary);
     format!("RSS {rss_bytes:.2}")
 }