Parallel indexing

crates/pile-dataset/Cargo.toml

@@ -22,3 +22,4 @@ jsonpath-rust = { workspace = true }
 chrono = { workspace = true }
 toml = { workspace = true }
 thiserror = { workspace = true }
+rayon = { workspace = true }

crates/pile-dataset/src/dataset.rs

@@ -1,13 +1,23 @@
 use chrono::{DateTime, Utc};
 use pile_config::{ConfigToml, Source};
 use pile_toolbox::cancelabletask::{CancelFlag, CancelableTaskError};
-use std::{io::ErrorKind, path::PathBuf, sync::Arc};
+use rayon::{
+	ThreadPoolBuilder,
+	iter::{IntoParallelIterator, ParallelIterator},
+};
+use std::{
+	io::ErrorKind,
+	path::PathBuf,
+	sync::{Arc, mpsc::Receiver},
+	thread::JoinHandle,
+	time::Instant,
+};
 use tantivy::{Executor, Index, IndexWriter, TantivyError, collector::TopDocs};
 use thiserror::Error;
-use tracing::{info, trace, warn};
+use tracing::{debug, info, trace, warn};
 
 use crate::{
-	DataSource,
+	DataSource, Item,
 	index::{DbFtsIndex, FtsLookupResult},
 	path_ts_earliest,
 	source::DirDataSource,
@@ -85,6 +95,7 @@ impl Dataset {
 	/// Refresh this dataset's fts index
 	pub fn fts_refresh(
 		&self,
+		threads: usize,
 		flag: Option<CancelFlag>,
 	) -> Result<(), CancelableTaskError<DatasetError>> {
 		let fts_tmp_dir = self.path_workdir.join(".tmp-fts");
@@ -94,23 +105,9 @@ impl Dataset {
 			warn!("Removing temporary index in {}", fts_dir.display());
 			std::fs::remove_dir_all(&fts_tmp_dir).map_err(DatasetError::from)?;
 		}
-		if fts_dir.is_dir() {
-			warn!("Removing existing index in {}", fts_dir.display());
-			std::fs::remove_dir_all(&fts_dir).map_err(DatasetError::from)?;
-		}
 
 		std::fs::create_dir_all(&fts_tmp_dir).map_err(DatasetError::from)?;
 
-		let mut sources = Vec::new();
-		for (name, source) in &self.config.dataset.source {
-			match source {
-				Source::Flac { path: dir } => {
-					let source = DirDataSource::new(name, dir.clone().to_vec());
-					sources.push(source);
-				}
-			}
-		}
-
 		let db_index = DbFtsIndex::new(&fts_tmp_dir, &self.config);
 		let mut index = Index::create_in_dir(&fts_tmp_dir, db_index.schema.clone())
 			.map_err(DatasetError::from)?;
@@ -118,35 +115,72 @@ impl Dataset {
 		let mut index_writer: IndexWriter =
 			index.writer(50 * 1024 * 1024).map_err(DatasetError::from)?;
 
-		for s in sources {
-			info!("Processing source {:?}", s.name);
+		let batch_size = 1000;
+		let (_read_task, read_rx) = start_read_task(&self.config, batch_size);
 
-			for i in s.iter() {
-				let (k, v) = i.map_err(DatasetError::from)?;
+		#[expect(clippy::unwrap_used)]
+		let write_pool = ThreadPoolBuilder::new()
+			.num_threads(threads.max(1))
+			.thread_name(|x| format!("fts_refresh_thread_{x}"))
+			.build()
+			.unwrap();
 
-				let doc = match db_index.entry_to_document(&*v) {
-					Ok(Some(x)) => x,
-					Ok(None) => {
-						warn!("Skipping {k:?}, document is empty");
-						continue;
-					}
-					Err(err) => {
-						warn!("Could not read {k:?}, skipping. {err}");
-						continue;
-					}
-				};
-				index_writer.add_document(doc).map_err(DatasetError::from)?;
+		let mut total = 0u64;
+		while let Ok(batch) = read_rx.recv() {
+			let batch = batch.map_err(DatasetError::from)?;
+			let len = batch.len() as u64;
 
-				if let Some(flag) = flag.as_ref()
-					&& flag.is_cancelled()
-				{
-					return Err(CancelableTaskError::Cancelled);
-				}
+			if let Some(flag) = &flag
+				&& flag.is_cancelled()
+			{
+				return Err(CancelableTaskError::Cancelled);
 			}
+
+			let start = Instant::now();
+			write_pool
+				.install(|| {
+					batch
+						.into_par_iter()
+						.filter_map(|(key, item)| match db_index.entry_to_document(&*item) {
+							Ok(Some(doc)) => Some((key, doc)),
+							Ok(None) => {
+								warn!("Skipping {key:?}, document is empty");
+								None
+							}
+							Err(err) => {
+								warn!("Could not read {key:?}, skipping. {err}");
+								None
+							}
+						})
+						.map(|(key, doc)| {
+							index_writer
+								.add_document(doc)
+								.map_err(|err| (key, err))
+								.map(|_| ())
+						})
+						.find_first(|x| x.is_err())
+						.unwrap_or(Ok(()))
+				})
+				.map_err(|(_key, err)| DatasetError::from(err))?;
+
+			total += len;
+			let time_ms = start.elapsed().as_millis();
+			debug!("Added a batch of {len} in {time_ms} ms ({total} total)");
+		}
+
+		if let Some(flag) = flag.as_ref()
+			&& flag.is_cancelled()
+		{
+			return Err(CancelableTaskError::Cancelled);
 		}
 
 		info!("Committing index");
 		index_writer.commit().map_err(DatasetError::from)?;
+
+		if fts_dir.is_dir() {
+			warn!("Removing existing index in {}", fts_dir.display());
+			std::fs::remove_dir_all(&fts_dir).map_err(DatasetError::from)?;
+		}
 		std::fs::rename(&fts_tmp_dir, &fts_dir).map_err(DatasetError::from)?;
 
 		return Ok(());
@@ -229,3 +263,52 @@ impl Dataset {
 		}
 	}
 }
+
+//
+// MARK: read_task
+//
+
+fn start_read_task(
+	config: &ConfigToml,
+	batch_size: usize,
+) -> (
+	JoinHandle<()>,
+	Receiver<Result<Vec<(PathBuf, Box<dyn Item<Key = PathBuf>>)>, DatasetError>>,
+) {
+	let config = config.clone();
+	let (read_tx, read_rx) = std::sync::mpsc::sync_channel(2);
+
+	let read_task = std::thread::spawn(move || {
+		let mut batch = Vec::with_capacity(batch_size);
+		for (name, source) in &config.dataset.source {
+			info!("Loading source {name}");
+
+			match source {
+				Source::Flac { path: dir } => {
+					let source = DirDataSource::new(name, dir.clone().to_vec());
+					for i in source.iter() {
+						match i {
+							Ok(x) => batch.push(x),
+							Err(err) => {
+								let err = Err(DatasetError::from(err));
+								let _ = read_tx.send(err);
+								return;
+							}
+						}
+
+						if batch.len() >= batch_size {
+							let b = std::mem::replace(&mut batch, Vec::with_capacity(batch_size));
+
+							match read_tx.send(Ok(b)) {
+								Ok(()) => {}
+								Err(_) => return,
+							};
+						}
+					}
+				}
+			}
+		}
+	});
+
+	return (read_task, read_rx);
+}

crates/pile-dataset/src/source/dir.rs

@@ -44,20 +44,28 @@ impl DataSource for DirDataSource {
 			.iter()
 			.flat_map(|x| WalkDir::new(x).into_iter().map_ok(move |d| (x, d)))
 			.filter_ok(|(_, entry)| !entry.file_type().is_dir())
-			.map(|x| match x {
+			.filter_map(|x| match x {
 				Err(err) => {
 					let msg = format!("other walkdir error: {err:?}");
-					Err(err.into_io_error().unwrap_or(std::io::Error::other(msg)))
+					Some(Err(err
+						.into_io_error()
+						.unwrap_or(std::io::Error::other(msg))))
 				}
+
 				Ok((_, entry)) => {
 					let path = entry.into_path();
-					let item = FlacItem {
-						source_name: self.name.clone(),
-						path: path.clone(),
-					};
 
-					let item: Box<dyn Item<Key = Self::Key>> = Box::new(item);
-					Ok((path, item))
+					let item: Box<dyn Item<Key = Self::Key>> =
+						match path.extension().map(|x| x.to_str()).flatten() {
+							None => return None,
+							Some("flac") => Box::new(FlacItem {
+								source_name: self.name.clone(),
+								path: path.clone(),
+							}),
+							Some(_) => return None,
+						};
+
+					Some(Ok((path, item)))
 				}
 			});
 	}

crates/pile-dataset/src/traits.rs

@@ -7,7 +7,7 @@ pub trait DataSource {
 	/// (e.g, a PathBuf or a primary key)
 	type Key: Key;
 
-	type Error: Error;
+	type Error: Error + Sync + Send;
 
 	/// Get an item from this datasource
 	fn get(&self, key: &Self::Key) -> Result<Option<Box<dyn Item<Key = Self::Key>>>, Self::Error>;