diff --git a/Cargo.lock b/Cargo.lock index bfd6219..1087441 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6263,6 +6263,7 @@ dependencies = [ "gix-pack", "hex", "indicatif 0.18.6", + "lance", "lancedb", "model2vec-rs", "ndarray 0.17.2", diff --git a/Cargo.toml b/Cargo.toml index 32ed797..5346340 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -66,6 +66,11 @@ chrono = "0.4.44" chrono-english = "0.1" # Natural language date parsing base64 = "0.22" # Base64 decoding for MIME-encoded email bodies +[dev-dependencies] +# Tests reach lance's WriteParams, which lancedb does not re-export. Pin +# the exact version lancedb resolves to, so the two cannot drift apart. +lance = { version = "=4.0.0", default-features = false } + [[bin]] name = "semcode-index" path = "src/bin/index.rs" diff --git a/src/bin/index.rs b/src/bin/index.rs index aaa1fa3..efc33bb 100644 --- a/src/bin/index.rs +++ b/src/bin/index.rs @@ -1202,7 +1202,8 @@ async fn main() -> Result<()> { // combined working set triggers the OOM killer. println!("\nCompacting lore tables..."); match db_manager.compact_lore_tables().await { - Ok(_) => println!("Lore table compaction completed successfully"), + Ok(true) => println!("Lore table compaction completed successfully"), + Ok(false) => {} Err(e) => error!("Failed to compact lore tables: {}", e), } @@ -1216,7 +1217,8 @@ async fn main() -> Result<()> { Err(e) => eprintln!("Warning: Failed to ensure FTS indices: {}", e), } match db_manager.optimize_lore_fts_indices().await { - Ok(_) => println!("FTS indices updated successfully"), + Ok(true) => println!("FTS indices updated successfully"), + Ok(false) => {} Err(e) => eprintln!("Warning: Failed to optimize FTS indices: {}", e), } } @@ -1376,7 +1378,8 @@ async fn main() -> Result<()> { if total_new_emails > 0 { println!("\nCompacting lore tables..."); match db_manager.compact_lore_tables().await { - Ok(_) => println!("Lore table compaction completed successfully"), + Ok(true) => println!("Lore table compaction completed successfully"), + Ok(false) => {} Err(e) => error!("Failed to compact lore tables: {}", e), } @@ -1386,7 +1389,8 @@ async fn main() -> Result<()> { Err(e) => eprintln!("Warning: Failed to ensure FTS indices: {}", e), } match db_manager.optimize_lore_fts_indices().await { - Ok(_) => println!("FTS indices updated successfully"), + Ok(true) => println!("FTS indices updated successfully"), + Ok(false) => {} Err(e) => eprintln!("Warning: Failed to optimize FTS indices: {}", e), } } diff --git a/src/database/connection.rs b/src/database/connection.rs index 22d8d48..47e51c7 100644 --- a/src/database/connection.rs +++ b/src/database/connection.rs @@ -889,7 +889,8 @@ impl DatabaseManager { self.schema_manager.compact_and_cleanup().await } - pub async fn compact_lore_tables(&self) -> Result<()> { + /// Returns false, after printing the problem, if any table failed. + pub async fn compact_lore_tables(&self) -> Result { self.schema_manager.compact_lore_tables().await } @@ -909,7 +910,8 @@ impl DatabaseManager { } /// Merge newly-inserted rows into existing lore FTS indices. - pub async fn optimize_lore_fts_indices(&self) -> Result<()> { + /// Returns false, after printing why, if the merge was skipped. + pub async fn optimize_lore_fts_indices(&self) -> Result { self.schema_manager.optimize_lore_fts_indices().await } diff --git a/src/database/schema.rs b/src/database/schema.rs index 56f0a58..60a0479 100644 --- a/src/database/schema.rs +++ b/src/database/schema.rs @@ -7,7 +7,7 @@ use futures::{stream, StreamExt, TryStreamExt}; use lancedb::connection::Connection; use lancedb::index::{scalar::BTreeIndexBuilder, scalar::FtsIndexBuilder, Index as LanceIndex}; use lancedb::query::{ExecutableQuery, QueryBase}; -use lancedb::table::{OptimizeAction, OptimizeOptions}; +use lancedb::table::{CompactionOptions, OptimizeAction, OptimizeOptions}; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; @@ -1721,18 +1721,7 @@ impl SchemaManager { // Prune orphaned index data left behind by drop_index(). if dropped { tracing::info!("Pruning orphaned index data from lore table..."); - if let Err(e) = table - .optimize(OptimizeAction::Prune { - older_than: Some( - lancedb::table::Duration::try_seconds(0).expect("valid duration"), - ), - delete_unverified: Some(true), - error_if_tagged_old_versions: Some(false), - }) - .await - { - tracing::warn!("Failed to prune lore table after FTS rebuild: {}", e); - } + Self::prune_old_versions(&table, "lore").await; } Ok(()) @@ -1787,7 +1776,9 @@ impl SchemaManager { /// brute-force fallback at query time, so queries remain correct /// even before this call. Running optimize merges those rows /// into the inverted index structure, eliminating the scan cost. - pub async fn optimize_lore_fts_indices(&self) -> Result<()> { + /// + /// Returns false, after printing why, if it skipped the merge. + pub async fn optimize_lore_fts_indices(&self) -> Result { // Guard against running on a table with a large _indices/ // backlog. lance/index/append.rs opens every delta fragment // for a column before merging any, so peak memory scales @@ -1796,23 +1787,37 @@ impl SchemaManager { // drives semcode-index into swap and gets it OOM-killed. // Query correctness is preserved regardless: unindexed rows // still fall back to a brute-force scan. + // + // Count only non-empty entries: pruning old versions deletes a + // superseded index's files but leaves its directory behind. + // Those empty directories accumulate, one per FTS and BTree + // index each lore run; they cost a directory entry each and + // nothing else. What stays non-empty is the current set of + // indices plus the BTree indices that compact_lore_tables() + // superseded earlier in this run, which the next run prunes. const MAX_LORE_INDEX_FRAGMENTS: usize = 100; let uri = self.connection.uri(); let indices_dir = std::path::Path::new(uri) .join("lore.lance") .join("_indices"); if let Ok(rd) = std::fs::read_dir(&indices_dir) { - let count = rd.count(); + let count = rd + .filter_map(|entry| entry.ok()) + .filter(|entry| { + std::fs::read_dir(entry.path()) + .map(|mut files| files.next().is_some()) + .unwrap_or(true) + }) + .count(); if count > MAX_LORE_INDEX_FRAGMENTS { - tracing::warn!( - "Skipping lore FTS index optimization: \ + eprintln!( + "Warning: Skipping lore FTS index optimization: \ {} _indices/ fragments exceeds {} threshold. \ Queries remain correct via brute-force fallback. \ Rebuild the lore table on a larger host to recover.", - count, - MAX_LORE_INDEX_FRAGMENTS + count, MAX_LORE_INDEX_FRAGMENTS ); - return Ok(()); + return Ok(false); } } @@ -1828,8 +1833,8 @@ impl SchemaManager { .collect(); if fts_index_names.is_empty() { - tracing::warn!("Skipping lore FTS optimization: no FTS indices found"); - return Ok(()); + eprintln!("Warning: Skipping lore FTS index optimization: no FTS indices found"); + return Ok(false); } let start_time = std::time::Instant::now(); @@ -1849,7 +1854,7 @@ impl SchemaManager { "Lore FTS index optimization completed in {:.1}s", elapsed.as_secs_f64() ); - Ok(()) + Ok(true) } /// Shared helper: create all 5 FTS indices on an already-opened table. @@ -2050,19 +2055,31 @@ impl SchemaManager { /// the database, including code-index tables and content shards /// that a lore run never touches. On memory-constrained systems /// the combined working set of those compactions triggers the OOM - /// killer. This method limits work to the two lore tables and + /// killer. This method limits work to the lore tables and /// processes them sequentially to keep peak memory low. - pub async fn compact_lore_tables(&self) -> Result<()> { + /// + /// Returns false, after printing the problem, if any table failed. + pub async fn compact_lore_tables(&self) -> Result { tracing::info!("Running compaction for lore tables..."); let table_names = self.connection.table_names().execute().await?; - let lore_tables = ["lore", "lore_indexed_commits"]; + let lore_tables = ["lore", "lore_indexed_commits", "lore_vectors"]; + // Print problems to stderr: the tracing warnings carrying the + // details are filtered out below SEMCODE_DEBUG=warn, and a prune + // that fails silently lets _indices/ grow until lore FTS updates + // stop, as they did when lore was never pruned. + let mut all_ok = true; for name in &lore_tables { if !table_names.iter().any(|n| n == name) { continue; } - match Self::optimize_single_table(&self.connection, name).await { + let outcome = if *name == "lore" { + Self::maintain_lore_table(&self.connection).await + } else { + Self::optimize_single_table(&self.connection, name).await + }; + match outcome { Ok(OptimizeOutcome::Optimized) => { tracing::info!("Compacted table {}", name); } @@ -2070,15 +2087,234 @@ impl SchemaManager { tracing::info!("Skipped table {} (too small)", name); } Ok(OptimizeOutcome::PartialFailure) => { - tracing::warn!("Partial failure compacting table {}", name); + eprintln!( + "Warning: Partial failure compacting table {} \ + (set SEMCODE_DEBUG=warn for details)", + name + ); + all_ok = false; } Err(e) => { - tracing::warn!("Failed to compact table {}: {}", name, e); + eprintln!("Warning: Failed to compact table {}: {}", name, e); + all_ok = false; } } } - Ok(()) + Ok(all_ok) + } + + /// Delete every version of `table` but the latest, along with the + /// data and index files only those versions referenced. Returns + /// false if the prune failed. + async fn prune_old_versions(table: &lancedb::table::Table, table_name: &str) -> bool { + match table + .optimize(OptimizeAction::Prune { + older_than: Some(lancedb::table::Duration::try_seconds(0).expect("valid duration")), + delete_unverified: Some(true), + error_if_tagged_old_versions: Some(false), + }) + .await + { + Ok(_) => true, + Err(e) => { + tracing::warn!( + "Failed to prune old versions from table {}: {}", + table_name, + e + ); + false + } + } + } + + /// Prune the lore table's old versions, merge its small fragments and + /// rebuild the BTree indices on its non-nullable columns: + /// optimize_single_table() skips lore entirely. + /// + /// Every lore run rewrites the FTS indices, and unless old versions + /// are pruned the superseded copies stay under _indices/ until the + /// count trips the guard in optimize_lore_fts_indices() and FTS + /// updates stop for good. Prune reads only each manifest's index + /// metadata, not the index files. + /// + /// The prune runs first, so the fragments a compaction replaces and the + /// BTree files the rebuild supersedes stay until the next run prunes + /// them, and a query still reading them is not cut off. + /// + /// Only the lore run calls this, not optimize_database(): with + /// delete_unverified set, a prune can delete files that another + /// process writing to lore has not committed yet. That also relies + /// on lore runs not overlapping, e.g. as a systemd oneshot unit ensures. + async fn maintain_lore_table(connection: &Connection) -> Result { + let table = connection.open_table("lore").execute().await?; + let pruned = Self::prune_old_versions(&table, "lore").await; + let compacted = Self::compact_lore_fragments(&table).await; + let btrees_current = Self::rebuild_lore_btree_indices(&table).await; + Ok(if pruned && compacted && btrees_current { + OptimizeOutcome::Optimized + } else { + OptimizeOutcome::PartialFailure + }) + } + + /// Rewrite the lore table's small fragments into fragments of up to + /// LORE_TARGET_ROWS_PER_FRAGMENT rows once it has more than + /// LORE_COMPACT_FRAGMENTS. Each lore run appends a few fragments, and + /// queries that scan lore pay per fragment. Returns false, after + /// logging why, on failure. + /// + /// Fragments at or above the target are not candidates, and lance + /// 4.0.0 merges only fragments covered by the same set of indices, so + /// a compaction rewrites the fragments appended since the last one and + /// leaves an earlier compaction's output alone even when it is under + /// the target: one leftover fragment per compaction. + /// + /// Index remapping is deferred: lance records the old-to-new fragment + /// mapping in its fragment-reuse index and applies it when an index is + /// read, instead of rewriting every index now, which is what costs + /// memory on a table with large FTS indices. The FTS update applies + /// the mapping too, but lance 4.0.0's BTree update does not: it drops + /// every row the index held in a rewritten fragment, while the index + /// statistics still report full coverage. The BTrees are rebuilt + /// from the table as it is instead (rebuild_lore_btree_indices()), so + /// the mapping does not concern them. lance 4.0.0 never trims the + /// fragment-reuse index, which gains a version of a few KB per + /// compaction; the fragment threshold keeps that to one compaction + /// every few hundred appended fragments. + /// + /// A table of nothing but small fragments, such as one built from + /// scratch, is rewritten whole on its first compaction, one task per + /// LORE_TARGET_ROWS_PER_FRAGMENT rows. The tasks run one at a time: + /// each one in flight adds to peak memory, and a compaction this + /// large happens once per database. + async fn compact_lore_fragments(table: &lancedb::table::Table) -> bool { + const LORE_COMPACT_FRAGMENTS: usize = 256; + const LORE_TARGET_ROWS_PER_FRAGMENT: usize = 150_000; + const LORE_COMPACT_THREADS: usize = 1; + + let fragments = match table.stats().await { + Ok(stats) => stats.fragment_stats.num_fragments, + Err(e) => { + tracing::warn!("Failed to read stats for table lore: {}", e); + return false; + } + }; + if fragments <= LORE_COMPACT_FRAGMENTS { + return true; + } + + tracing::info!("Compacting lore table ({} fragments)...", fragments); + let table = table.clone(); + Self::run_lore_task("compaction", async move { + table + .optimize(OptimizeAction::Compact { + options: CompactionOptions { + target_rows_per_fragment: LORE_TARGET_ROWS_PER_FRAGMENT, + defer_index_remap: true, + num_threads: Some(LORE_COMPACT_THREADS), + ..Default::default() + }, + remap_options: None, + }) + .await + }) + .await + } + + /// Run one maintenance step on the lore table and report whether it + /// succeeded, logging why not. The step runs as its own task: a lance + /// panic would otherwise end semcode-index before the FTS update, on + /// every run until the table is repaired by hand, and as a task it + /// becomes a JoinError. The next run's prune deletes the files an + /// aborted step wrote. + async fn run_lore_task(what: &str, step: F) -> bool + where + T: Send + 'static, + F: std::future::Future> + Send + 'static, + { + let start_time = std::time::Instant::now(); + match tokio::spawn(step).await { + Ok(Ok(_)) => { + tracing::info!( + "Lore {} completed in {:.1}s", + what, + start_time.elapsed().as_secs_f64() + ); + true + } + Ok(Err(e)) => { + tracing::warn!("Lore {} failed: {}", what, e); + false + } + Err(e) => { + tracing::warn!("Lore {} did not finish: {}", what, e); + false + } + } + } + + /// Build the lore table's BTree indices on non-nullable columns again + /// from the whole table, under their existing names. The BTrees are + /// created at schema setup, usually on an empty table, and a lookup by + /// message_id scans every row they do not cover. Returns false, after + /// logging why, on failure; a table without such indices has nothing + /// to rebuild. + /// + /// A rebuild rather than lance's in-place update, which merges only + /// new rows: a rebuild depends on nothing but the table as it is now, + /// so no earlier failed run or rewrite of the table's fragments can + /// leave rows out of the result. + /// + /// BTrees on nullable columns (in_reply_to, references) are left as + /// they are: lance 4.0.0 writes a BTree over a nullable string column + /// with enough rows into pages it cannot decode. Reading the index + /// back panics in lance-encoding, and so does any lookup that lands on + /// a bad page. Rebuilding them would replace a stale but readable + /// index with an unreadable one. Lift this once semcode moves to a + /// lance whose writer reads back its own output. + async fn rebuild_lore_btree_indices(table: &lancedb::table::Table) -> bool { + use lancedb::index::IndexType; + + let (indices, schema) = match tokio::try_join!(table.list_indices(), table.schema()) { + Ok(found) => found, + Err(e) => { + tracing::warn!("Failed to read indices of table lore: {}", e); + return false; + } + }; + let non_nullable = |column: &String| { + schema + .field_with_name(column) + .is_ok_and(|field| !field.is_nullable()) + }; + let btrees: Vec = indices + .into_iter() + .filter(|index| { + index.index_type == IndexType::BTree && index.columns.iter().all(non_nullable) + }) + .collect(); + if btrees.is_empty() { + return true; + } + + tracing::info!("Rebuilding {} lore BTree indices...", btrees.len()); + let table = table.clone(); + Self::run_lore_task("BTree index rebuild", async move { + for index in btrees { + table + .create_index( + &index.columns, + LanceIndex::BTree(BTreeIndexBuilder::default()), + ) + .name(index.name) + .replace(true) + .execute() + .await?; + } + Ok(()) + }) + .await } /// Optimize a single table - runs compact, prune, and index operations @@ -2098,11 +2334,14 @@ impl SchemaManager { // the operation completes. On memory-constrained systems the // resident set grows into swap and the OOM killer terminates // semcode-index before compaction finishes, leaving fresh - // delta fragments behind each time. Skip the table entirely. + // delta fragments behind each time. Skip the table entirely; + // compact_lore_tables() prunes its old versions, compacts it + // without remapping its indices and rebuilds the BTree indices on + // its non-nullable columns. if table_name == "lore" { tracing::info!( "Skipping generic optimization for lore table \ - (handled by optimize_lore_fts_indices)" + (maintained by compact_lore_tables)" ); return Ok(OptimizeOutcome::Skipped); } @@ -2129,11 +2368,22 @@ impl SchemaManager { let mut success = true; - // 1. Compact files + // 1. Compact files. Tables above MAX_COMPACT_FRAGMENTS are left + // alone, as a guard against lance compaction spinning on thousands + // of small fragments. lore_indexed_commits is exempt: it holds one + // 40-char SHA per row and no index, and a table of that shape + // compacts in seconds at 3,660 fragments on lance 4.0.0. A cold + // lore build leaves it with one fragment per inserted batch, over a + // thousand for a full archive set, so the guard would keep it + // uncompacted. const MAX_COMPACT_FRAGMENTS: usize = 500; + let fragment_limit_applies = table_name != "lore_indexed_commits"; let should_compact = match table.stats().await { - Ok(stats) if stats.fragment_stats.num_fragments > MAX_COMPACT_FRAGMENTS => { + Ok(stats) + if fragment_limit_applies + && stats.fragment_stats.num_fragments > MAX_COMPACT_FRAGMENTS => + { tracing::warn!( "Skipping compaction for table {} \ ({} fragments exceeds {} limit -- \ @@ -2142,6 +2392,9 @@ impl SchemaManager { stats.fragment_stats.num_fragments, MAX_COMPACT_FRAGMENTS ); + // Not compacting is not success: the fragment count only + // grows from here, and callers report the outcome. + success = false; false } Ok(_) => true, @@ -2165,19 +2418,7 @@ impl SchemaManager { } // 2. Prune ALL old versions - if let Err(e) = table - .optimize(OptimizeAction::Prune { - older_than: Some(lancedb::table::Duration::try_seconds(0).expect("valid duration")), - delete_unverified: Some(true), - error_if_tagged_old_versions: Some(false), - }) - .await - { - tracing::warn!( - "Failed to prune old versions from table {}: {}", - table_name, - e - ); + if !Self::prune_old_versions(&table, table_name).await { success = false; } @@ -2481,3 +2722,437 @@ impl SchemaManager { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::database::DatabaseManager; + use crate::types::LoreEmailInfo; + use arrow::array::StringArray; + use lancedb::connection::LanceFileVersion; + use lancedb::database::listing::ListingDatabaseOptions; + use lancedb::index::scalar::FullTextSearchQuery; + use lancedb::index::IndexType; + use lancedb::query::Select; + + async fn fresh_schema_manager() -> (tempfile::TempDir, SchemaManager) { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("schema.lancedb"); + let options = ListingDatabaseOptions::builder() + .data_storage_version(LanceFileVersion::V2_2) + .build(); + let connection = lancedb::connect(db_path.to_str().unwrap()) + .database_options(&options) + .execute() + .await + .unwrap(); + (dir, SchemaManager::new(connection)) + } + + /// A cold lore build appends one lore_indexed_commits fragment per + /// batch. 600 two-row fragments put the table above + /// MAX_COMPACT_FRAGMENTS and above MIN_ROWS_FOR_OPTIMIZATION, so the + /// lore compaction has to compact it rather than skip it. One write + /// capped at two rows per file builds them all in a single commit. + #[tokio::test] + async fn lore_indexed_commits_compacts_above_the_fragment_limit() { + const FRAGMENTS: usize = 600; + const ROWS_PER_FRAGMENT: usize = 2; + + let (_dir, schema_manager) = fresh_schema_manager().await; + schema_manager + .create_lore_indexed_commits_table() + .await + .unwrap(); + let table = schema_manager + .connection + .open_table("lore_indexed_commits") + .execute() + .await + .unwrap(); + let shas: Vec = (0..FRAGMENTS * ROWS_PER_FRAGMENT) + .map(|row| format!("{row:040x}")) + .collect(); + let batch = RecordBatch::try_new( + table.schema().await.unwrap(), + vec![Arc::new(StringArray::from(shas))], + ) + .unwrap(); + table + .add(vec![batch]) + .write_options(lancedb::table::WriteOptions { + lance_write_params: Some(lance::dataset::WriteParams { + max_rows_per_file: ROWS_PER_FRAGMENT, + max_rows_per_group: ROWS_PER_FRAGMENT, + mode: lance::dataset::WriteMode::Append, + ..Default::default() + }), + }) + .execute() + .await + .unwrap(); + let before = table.stats().await.unwrap().fragment_stats.num_fragments; + assert_eq!(before, FRAGMENTS); + + assert!(schema_manager.compact_lore_tables().await.unwrap()); + + table.checkout_latest().await.unwrap(); + let after = table.stats().await.unwrap().fragment_stats.num_fragments; + assert!(after < 10, "{after} fragments left after compaction"); + assert_eq!( + table.count_rows(None).await.unwrap(), + FRAGMENTS * ROWS_PER_FRAGMENT + ); + } + + /// BTrees on non-nullable columns that every lore database carries: + /// each lore run brings them up to date. + const UPDATED_BTREES: [&str; 3] = ["message_id_idx", "git_commit_sha_idx", "date_idx"]; + + /// BTrees on nullable columns, which lance 4.0.0 cannot rewrite + /// readably: lore runs leave them as they were created, on the empty + /// table. + const NULLABLE_BTREES: [&str; 2] = ["in_reply_to_idx", "references_idx"]; + + async fn lore_database() -> (tempfile::TempDir, DatabaseManager) { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("lore-btree.db"); + let db = DatabaseManager::new(db_path.to_str().unwrap(), String::new()) + .await + .unwrap(); + db.create_tables().await.unwrap(); + (dir, db) + } + + /// Emails `first..first + count`, every second one a reply. + fn emails(first: usize, count: usize) -> Vec { + (first..first + count) + .map(|i| LoreEmailInfo { + git_commit_sha: format!("{i:040x}"), + from: format!("Sender {i} "), + date: "Tue, 10 Nov 2026 12:00:00 +0000".to_string(), + date_timestamp: 1_794_312_000 + i as i64, + message_id: format!(""), + in_reply_to: (i % 2 == 1).then(|| format!("", i - 1)), + subject: format!("[PATCH] change {i}"), + references: (i % 2 == 1).then(|| format!("", i - 1)), + recipients: "list@example.org".to_string(), + body: format!("body {i} {}", if i % 3 == 0 { "needle" } else { "hay" }), + symbols: Vec::new(), + }) + .collect() + } + + /// Every BTree index on lore as (name, indexed rows, unindexed rows). + async fn btree_coverage(db: &DatabaseManager) -> Vec<(String, usize, usize)> { + let table = db.connection().open_table("lore").execute().await.unwrap(); + let mut coverage = Vec::new(); + for index in table.list_indices().await.unwrap() { + if index.index_type != IndexType::BTree { + continue; + } + let stats = table.index_stats(&index.name).await.unwrap().unwrap(); + coverage.push((index.name, stats.num_indexed_rows, stats.num_unindexed_rows)); + } + coverage.sort(); + coverage + } + + /// Asserts that every BTree on a non-nullable column covers all `rows` + /// and that the nullable-column BTrees still cover none. + fn assert_covers(coverage: &[(String, usize, usize)], rows: usize) { + for name in UPDATED_BTREES.iter().chain(&NULLABLE_BTREES) { + assert!( + coverage.iter().any(|(n, _, _)| n == name), + "{name} missing from {coverage:?}" + ); + } + for (name, indexed, unindexed) in coverage { + let expected = if NULLABLE_BTREES.contains(&name.as_str()) { + (0, rows) + } else { + (rows, 0) + }; + assert_eq!( + (*indexed, *unindexed), + expected, + "{name}: indexed and unindexed rows" + ); + } + } + + #[tokio::test] + async fn lore_btree_indices_cover_new_rows() { + let (_dir, db) = lore_database().await; + assert!(db + .insert_lore_emails(&emails(0, 50)) + .await + .unwrap() + .is_empty()); + + assert!(db.compact_lore_tables().await.unwrap()); + + assert_covers(&btree_coverage(&db).await, 50); + } + + /// A later run, with the FTS indices in place as every run after the + /// first has them: the BTrees take in the new rows and the FTS indices + /// are left as they were. + #[tokio::test] + async fn lore_btree_indices_stay_current_across_runs() { + let (_dir, db) = lore_database().await; + assert!(db + .insert_lore_emails(&emails(0, 50)) + .await + .unwrap() + .is_empty()); + assert!(db.compact_lore_tables().await.unwrap()); + db.ensure_lore_fts_indices().await.unwrap(); + assert!(db.optimize_lore_fts_indices().await.unwrap()); + + assert!(db + .insert_lore_emails(&emails(50, 30)) + .await + .unwrap() + .is_empty()); + assert!(db.compact_lore_tables().await.unwrap()); + + assert_covers(&btree_coverage(&db).await, 80); + let table = db.connection().open_table("lore").execute().await.unwrap(); + let mut fts: Vec<(String, usize)> = Vec::new(); + for index in table.list_indices().await.unwrap() { + if index.index_type == IndexType::FTS { + let stats = table.index_stats(&index.name).await.unwrap().unwrap(); + fts.push((index.name, stats.num_indexed_rows)); + } + } + fts.sort(); + assert_eq!( + fts, + [ + "body_idx", + "from_idx", + "recipients_idx", + "subject_idx", + "symbols_idx" + ] + .map(|name| (name.to_string(), 50)) + .to_vec() + ); + } + + async fn lore_fragments(db: &DatabaseManager) -> usize { + let table = db.connection().open_table("lore").execute().await.unwrap(); + table.stats().await.unwrap().fragment_stats.num_fragments + } + + /// Adds emails `first..first + 2 * fragments` to lore as `fragments` + /// fragments of two rows, the layout many small appends leave, in one + /// commit, since an append per fragment makes the test slow. The rows + /// are built by insert_lore_emails() in a scratch database and copied + /// over. + async fn append_lore_fragments(db: &DatabaseManager, first: usize, fragments: usize) { + let (_scratch_dir, scratch) = lore_database().await; + let failed = scratch + .insert_lore_emails(&emails(first, 2 * fragments)) + .await + .unwrap(); + assert!(failed.is_empty()); + let rows: Vec = scratch + .connection() + .open_table("lore") + .execute() + .await + .unwrap() + .query() + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + + db.connection() + .open_table("lore") + .execute() + .await + .unwrap() + .add(rows) + .write_options(lancedb::table::WriteOptions { + lance_write_params: Some(lance::dataset::WriteParams { + max_rows_per_file: 2, + max_rows_per_group: 2, + mode: lance::dataset::WriteMode::Append, + ..Default::default() + }), + }) + .execute() + .await + .unwrap(); + } + + /// Sorted message_ids of the emails whose body FTS matches `needle`. + async fn needle_message_ids(db: &DatabaseManager) -> Vec { + let table = db.connection().open_table("lore").execute().await.unwrap(); + let query = FullTextSearchQuery::new("needle".to_string()) + .with_column("body".to_string()) + .unwrap(); + let batches: Vec = table + .query() + .full_text_search(query) + .select(Select::columns(&["message_id"])) + .limit(100_000) + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let mut ids: Vec = batches + .iter() + .flat_map(|batch| { + let column = batch + .column_by_name("message_id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..column.len()).map(|row| column.value(row).to_string()) + }) + .collect(); + ids.sort(); + ids + } + + /// The message_ids `needle_message_ids()` must return once emails + /// `0..count` are stored. + fn expected_needles(count: usize) -> Vec { + let mut ids: Vec = (0..count) + .filter(|i| i % 3 == 0) + .map(|i| format!("")) + .collect(); + ids.sort(); + ids + } + + /// Above the fragment threshold the lore table is compacted, queries + /// return the same rows through the indices built before it, and both + /// kinds of index keep taking in new rows afterwards. + #[tokio::test] + async fn compact_lore_merges_small_fragments() { + const FRAGMENTS: usize = 257; + let (_dir, db) = lore_database().await; + append_lore_fragments(&db, 0, FRAGMENTS).await; + db.ensure_lore_fts_indices().await.unwrap(); + assert_eq!(lore_fragments(&db).await, FRAGMENTS); + assert_eq!( + needle_message_ids(&db).await, + expected_needles(2 * FRAGMENTS) + ); + + assert!(db.compact_lore_tables().await.unwrap()); + + assert!(lore_fragments(&db).await < 10); + let table = db.connection().open_table("lore").execute().await.unwrap(); + assert_eq!(table.count_rows(None).await.unwrap(), 2 * FRAGMENTS); + assert_eq!( + needle_message_ids(&db).await, + expected_needles(2 * FRAGMENTS) + ); + assert_covers(&btree_coverage(&db).await, 2 * FRAGMENTS); + let found = db + .get_lore_email_by_message_id("") + .await + .unwrap() + .unwrap(); + assert_eq!(found.subject, "[PATCH] change 7"); + + let failed = db + .insert_lore_emails(&emails(2 * FRAGMENTS, 10)) + .await + .unwrap(); + assert!(failed.is_empty()); + assert!(db.optimize_lore_fts_indices().await.unwrap()); + assert!(db.compact_lore_tables().await.unwrap()); + assert_eq!( + needle_message_ids(&db).await, + expected_needles(2 * FRAGMENTS + 10) + ); + assert_covers(&btree_coverage(&db).await, 2 * FRAGMENTS + 10); + } + + /// At or below the fragment threshold the lore table is not compacted, + /// so a run adds no version to lance's fragment-reuse index. + #[tokio::test] + async fn compact_lore_leaves_few_fragments_alone() { + let (_dir, db) = lore_database().await; + append_lore_fragments(&db, 0, 256).await; + assert_eq!(lore_fragments(&db).await, 256); + + assert!(db.compact_lore_tables().await.unwrap()); + + assert_eq!(lore_fragments(&db).await, 256); + } + + /// How many rows `column IN (values)` finds. The column's BTree + /// serves the lookup, so a row the index lost is not found. + async fn rows_found(db: &DatabaseManager, column: &str, values: &[String]) -> usize { + let list = values + .iter() + .map(|value| format!("'{value}'")) + .collect::>() + .join(", "); + let table = db.connection().open_table("lore").execute().await.unwrap(); + let batches: Vec = table + .query() + .only_if(format!("{column} IN ({list})")) + .select(Select::columns(&["message_id"])) + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + batches.iter().map(|batch| batch.num_rows()).sum() + } + + /// Asserts that `message_id` and `git_commit_sha` lookups find each of + /// emails `0..count` exactly once. + async fn assert_lookups_find(db: &DatabaseManager, count: usize) { + let ids: Vec = (0..count) + .map(|i| format!("")) + .collect(); + let shas: Vec = (0..count).map(|i| format!("{i:040x}")).collect(); + assert_eq!(rows_found(db, "message_id", &ids).await, count); + assert_eq!(rows_found(db, "git_commit_sha", &shas).await, count); + } + + /// A compaction rewrites fragments the BTrees already index, as every + /// compaction after a database's first does. Lookups through the + /// BTrees must still find those rows, in that run and the next. + #[tokio::test] + async fn compact_lore_keeps_indexed_rows_findable() { + let (_dir, db) = lore_database().await; + append_lore_fragments(&db, 0, 200).await; + assert!(db.compact_lore_tables().await.unwrap()); + db.ensure_lore_fts_indices().await.unwrap(); + assert_covers(&btree_coverage(&db).await, 400); + append_lore_fragments(&db, 400, 57).await; + assert_eq!(lore_fragments(&db).await, 257); + + assert!(db.compact_lore_tables().await.unwrap()); + + assert!(lore_fragments(&db).await < 10); + assert_covers(&btree_coverage(&db).await, 514); + assert_lookups_find(&db, 514).await; + + let failed = db.insert_lore_emails(&emails(514, 4)).await.unwrap(); + assert!(failed.is_empty()); + assert!(db.optimize_lore_fts_indices().await.unwrap()); + assert!(db.compact_lore_tables().await.unwrap()); + + assert_covers(&btree_coverage(&db).await, 518); + assert_lookups_find(&db, 518).await; + assert_eq!(needle_message_ids(&db).await, expected_needles(518)); + } +} diff --git a/src/database/search.rs b/src/database/search.rs index 56458cf..b958ee8 100644 --- a/src/database/search.rs +++ b/src/database/search.rs @@ -2950,20 +2950,45 @@ impl VectorSearchManager { .execute() .await?; + // Workers send one Vec per lore RecordBatch that still has + // emails without a vector. On an incremental run most of + // those hold only a few, and every merge_insert joins + // against the whole table, so write once at least + // LORE_VECTOR_FLUSH_ROWS have accumulated. + const LORE_VECTOR_FLUSH_ROWS: usize = 8192; + let mut pending: Vec = Vec::new(); let mut total_inserted = 0; - while let Some(batch_vectors) = result_rx.recv().await { + let mut open = true; + while open { + match result_rx.recv().await { + Some(batch_vectors) => { + let received = batch_vectors.len(); + let count = + processed_clone.fetch_add(received, Ordering::Relaxed) + received; + pb_clone.set_position(count as u64); + pb_clone.set_message(format!("Vectorized {} emails", count)); + pending.extend(batch_vectors); + } + None => open = false, + } + if pending.is_empty() || (open && pending.len() < LORE_VECTOR_FLUSH_ROWS) { + continue; + } + + // merge_insert inserts every unmatched source row, so two + // copies of a message_id in one buffer would become two + // rows. Keep one. + let batch = crate::database::one_row_per_key(std::mem::take(&mut pending), |e| { + e.content_hash.clone() + }); if let Err(e) = - insert_lore_vectors_batch_with_table(&lore_vectors_table, &batch_vectors).await + insert_lore_vectors_batch_with_table(&lore_vectors_table, &batch).await { tracing::error!("Failed to insert lore vectors batch: {}", e); return Err(e); } - total_inserted += batch_vectors.len(); - let count = processed_clone.fetch_add(batch_vectors.len(), Ordering::Relaxed) - + batch_vectors.len(); - pb_clone.set_position(count as u64); - pb_clone.set_message(format!("Inserted {} vectors", count)); + total_inserted += batch.len(); } Ok::<_, anyhow::Error>(total_inserted) }); diff --git a/src/indexer.rs b/src/indexer.rs index 2752ea1..1573023 100644 --- a/src/indexer.rs +++ b/src/indexer.rs @@ -957,10 +957,24 @@ pub async fn process_lore_commits_pipeline( let handle = tokio::spawn(async move { loop { - // Get next batch from shared receiver + // Get next batch from shared receiver. A parser only + // sends a short batch when its input has run out, so + // keep receiving until the batch holds at least + // batch_size emails or every parser has finished. + // Otherwise each parser's last partial batch becomes + // its own small append, dozens of tiny fragments per + // archive on a refresh. let batch = { match result_rx_clone.lock() { - Ok(rx) => rx.recv(), + Ok(rx) => rx.recv().map(|mut emails| { + while emails.len() < batch_size { + match rx.recv() { + Ok(more) => emails.extend(more), + Err(_) => break, + } + } + emails + }), Err(e) => { error!("Inserter {} failed to lock receiver: {}", inserter_id, e); break;