Merge pull request #5527 from nnethercott/all-cpus-in-import-dump

Use all CPUs during an import dump

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

@@ -33,7 +33,6 @@ use crate::documents::{obkv_to_object, DocumentsBatchReader};
 use crate::error::{Error, InternalError};
 use crate::index::{PrefixSearch, PrefixSettings};
 use crate::progress::Progress;
-use crate::thread_pool_no_abort::ThreadPoolNoAbortBuilder;
 pub use crate::update::index_documents::helpers::CursorClonableMmap;
 use crate::update::{
     IndexerConfig, UpdateIndexingStep, WordPrefixDocids, WordPrefixIntegerDocids, WordsPrefixesFst,
@@ -228,24 +227,7 @@ where
         let possible_embedding_mistakes =
             crate::vector::error::PossibleEmbeddingMistakes::new(&field_distribution);
 
-        let backup_pool;
-        let pool = match self.indexer_config.thread_pool {
-            Some(ref pool) => pool,
-            None => {
-                // We initialize a backup pool with the default
-                // settings if none have already been set.
-                #[allow(unused_mut)]
-                let mut pool_builder = ThreadPoolNoAbortBuilder::new();
-
-                #[cfg(test)]
-                {
-                    pool_builder = pool_builder.num_threads(1);
-                }
-
-                backup_pool = pool_builder.build()?;
-                &backup_pool
-            }
-        };
+        let pool = &self.indexer_config.thread_pool;
 
         // create LMDB writer channel
         let (lmdb_writer_sx, lmdb_writer_rx): (

crates/milli/src/update/indexer_config.rs

@@ -1,7 +1,7 @@
 use grenad::CompressionType;
 
 use super::GrenadParameters;
-use crate::thread_pool_no_abort::ThreadPoolNoAbort;
+use crate::{thread_pool_no_abort::ThreadPoolNoAbort, ThreadPoolNoAbortBuilder};
 
 #[derive(Debug)]
 pub struct IndexerConfig {
@@ -9,9 +9,10 @@ pub struct IndexerConfig {
     pub max_nb_chunks: Option<usize>,
     pub documents_chunk_size: Option<usize>,
     pub max_memory: Option<usize>,
+    pub max_threads: Option<usize>,
     pub chunk_compression_type: CompressionType,
     pub chunk_compression_level: Option<u32>,
-    pub thread_pool: Option<ThreadPoolNoAbort>,
+    pub thread_pool: ThreadPoolNoAbort,
     pub max_positions_per_attributes: Option<u32>,
     pub skip_index_budget: bool,
 }
@@ -27,16 +28,39 @@ impl IndexerConfig {
     }
 }
 
+/// By default use only 1 thread for indexing in tests
+#[cfg(test)]
+pub fn default_thread_pool_and_threads() -> (ThreadPoolNoAbort, Option<usize>) {
+    let pool = ThreadPoolNoAbortBuilder::new_for_indexing()
+        .num_threads(1)
+        .build()
+        .expect("failed to build default rayon thread pool");
+
+    (pool, Some(1))
+}
+
+#[cfg(not(test))]
+pub fn default_thread_pool_and_threads() -> (ThreadPoolNoAbort, Option<usize>) {
+    let pool = ThreadPoolNoAbortBuilder::new_for_indexing()
+        .build()
+        .expect("failed to build default rayon thread pool");
+
+    (pool, None)
+}
+
 impl Default for IndexerConfig {
     fn default() -> Self {
+        let (thread_pool, max_threads) = default_thread_pool_and_threads();
+
         Self {
+            max_threads,
+            thread_pool,
             log_every_n: None,
             max_nb_chunks: None,
             documents_chunk_size: None,
             max_memory: None,
             chunk_compression_type: CompressionType::None,
             chunk_compression_level: None,
-            thread_pool: None,
             max_positions_per_attributes: None,
             skip_index_budget: false,
         }

crates/milli/src/update/mod.rs

@@ -4,7 +4,7 @@ pub use self::concurrent_available_ids::ConcurrentAvailableIds;
 pub use self::facet::bulk::FacetsUpdateBulk;
 pub use self::facet::incremental::FacetsUpdateIncrementalInner;
 pub use self::index_documents::*;
-pub use self::indexer_config::IndexerConfig;
+pub use self::indexer_config::{default_thread_pool_and_threads, IndexerConfig};
 pub use self::new::ChannelCongestion;
 pub use self::settings::{validate_embedding_settings, Setting, Settings};
 pub use self::update_step::UpdateIndexingStep;