From 34e98d5217a47d0b9cd4c87ae9812e79e0a07d50 Mon Sep 17 00:00:00 2001 From: Lu Qiu Date: Thu, 23 Jul 2026 11:08:19 -0700 Subject: [PATCH 1/3] perf(index): search fts partitions in pipelined chunks One tiny cpu-pool task per partition (~350 per query) interleaves so heavily under concurrent load that scoring loses cache locality and the dispatch rate floods the blocking pool. Partitions are now searched in chunks of LANCE_FTS_SEARCH_CHUNK (default 16): each chunk loads its postings and scoring DocSets concurrently, then a single cpu task scores the chunk's partitions back-to-back on one thread; chunks pipeline against each other through buffer_unordered. On a 320-core node against a 100M-doc 42-language corpus (5-term OR, k=100, warm), per-query scoring CPU drops ~5x and c16 throughput doubles (227 -> 428 qps at half the latency); combined with a contention-free cache read path the chunked layout reaches 1340 qps at c128 versus 345 for per-partition tasks. --- rust/lance-index/src/scalar/inverted/index.rs | 196 ++++++++++++------ 1 file changed, 128 insertions(+), 68 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index d00d270c1a2..d9d56f6ac5f 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -98,6 +98,21 @@ pub const INVERT_LIST_FILE: &str = "invert.lance"; pub const DOCS_FILE: &str = "docs.lance"; pub const METADATA_FILE: &str = "metadata.lance"; +/// Partitions searched per cpu-pool task in `bm25_search`. Chunking bounds +/// the task rate and lets the shared top-k floor propagate between the +/// partitions a thread scores back-to-back. `LANCE_FTS_SEARCH_CHUNK=1` +/// restores one-task-per-partition. +fn fts_search_chunk() -> usize { + static CHUNK: LazyLock = LazyLock::new(|| { + std::env::var("LANCE_FTS_SEARCH_CHUNK") + .ok() + .and_then(|v| v.parse().ok()) + .filter(|&n| n >= 1) + .unwrap_or(16) + }); + *CHUNK +} + pub const TOKEN_COL: &str = "_token"; pub const TOKEN_ID_COL: &str = "_token_id"; pub const TOKEN_FST_BYTES_COL: &str = "_token_fst_bytes"; @@ -273,16 +288,6 @@ struct PartitionCandidates { candidates: Vec, } -impl PartitionCandidates { - fn empty() -> Self { - Self { - tokens_by_position: Vec::new(), - grouped_expansions: Vec::new(), - candidates: Vec::new(), - } - } -} - #[derive(Debug)] struct LoadedPostings { postings: Vec, @@ -910,11 +915,18 @@ impl InvertedIndex { // global k-th - see `Wand::shared_threshold`). let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); let legacy_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + // Partitions are processed in chunks: each chunk loads its postings + // and scoring DocSets asynchronously, then ONE cpu-pool task searches + // the whole chunk sequentially. Chunks pipeline against each other + // through buffer_unordered (chunk i searches while chunk i+1 loads). + // Chunking bounds the cpu-pool task rate (vs one tiny task per + // partition) and lets the shared top-k floor propagate immediately + // between partitions searched on the same thread. let parts = self .partitions - .iter() - .map(|part| { - let part = part.clone(); + .chunks(fts_search_chunk()) + .map(|chunk| { + let chunk = chunk.to_vec(); let tokens = tokens.clone(); let params = params.clone(); let mask = mask.clone(); @@ -923,73 +935,121 @@ impl InvertedIndex { let impact_shared_threshold = impact_shared_threshold.clone(); let legacy_shared_threshold = legacy_shared_threshold.clone(); async move { - let loaded_postings = part - .load_posting_lists( - tokens.as_ref(), - params.as_ref(), - operator, - impact_scorer.as_ref(), - metrics.as_ref(), - ) - .await?; - let LoadedPostings { - postings, - grouped_expansions, - impact_safe, - exact_scoring_required, - } = loaded_postings; - if postings.is_empty() { - // No hits in this partition; its DocSet stays - // unloaded, so we never pay the per-doc - // row_id/num_tokens download for it. - return Result::Ok((part, PartitionCandidates::empty())); - } - let docs_for_wand = part.docs.docs_for_wand(operator, mask.as_ref()).await?; - let max_position = postings - .iter() - .map(|posting| posting.term_index() as usize) - .max() - .unwrap_or_default(); - let mut tokens_by_position = vec![String::new(); max_position + 1]; - for posting in &postings { - let idx = posting.term_index() as usize; - tokens_by_position[idx] = posting.token().to_owned(); + // Load the chunk's partitions concurrently; only the + // search below is batched onto one cpu task. + let loads = chunk.into_iter().map(|part| { + let tokens = tokens.clone(); + let params = params.clone(); + let mask = mask.clone(); + let metrics = metrics.clone(); + let impact_scorer = impact_scorer.clone(); + let impact_shared_threshold = impact_shared_threshold.clone(); + let legacy_shared_threshold = legacy_shared_threshold.clone(); + async move { + let loaded_postings = part + .load_posting_lists( + tokens.as_ref(), + params.as_ref(), + operator, + impact_scorer.as_ref(), + metrics.as_ref(), + ) + .await?; + let LoadedPostings { + postings, + grouped_expansions, + impact_safe, + exact_scoring_required, + } = loaded_postings; + if postings.is_empty() { + // No hits in this partition; its DocSet stays + // unloaded, so we never pay the per-doc + // row_id/num_tokens download for it. + return Result::Ok(None); + } + let docs_for_wand = + part.docs.docs_for_wand(operator, mask.as_ref()).await?; + let max_position = postings + .iter() + .map(|posting| posting.term_index() as usize) + .max() + .unwrap_or_default(); + let mut tokens_by_position = vec![String::new(); max_position + 1]; + for posting in &postings { + let idx = posting.term_index() as usize; + tokens_by_position[idx] = posting.token().to_owned(); + } + let use_global_scorer = impact_safe || exact_scoring_required; + let partition_threshold = if use_global_scorer { + impact_shared_threshold + } else { + legacy_shared_threshold + }; + let wand_scorer = use_global_scorer.then(|| impact_scorer.clone()); + Result::Ok(Some(( + part, + docs_for_wand, + postings, + wand_scorer, + partition_threshold, + tokens_by_position, + grouped_expansions, + ))) + } + }); + let loaded: Vec<_> = futures::future::try_join_all(loads) + .await? + .into_iter() + .flatten() + .collect(); + if loaded.is_empty() { + return Result::Ok(Vec::new()); } let params = params.clone(); let mask = mask.clone(); let metrics = metrics.clone(); - let part_for_wand = part.clone(); - let use_global_scorer = impact_safe || exact_scoring_required; - let partition_threshold = if use_global_scorer { - impact_shared_threshold - } else { - legacy_shared_threshold - }; - let wand_scorer = use_global_scorer.then(|| impact_scorer.clone()); - let candidates = spawn_cpu(move || { - let candidates = part_for_wand.bm25_search( - docs_for_wand.as_ref(), - params.as_ref(), - operator, - mask, + let results = spawn_cpu(move || { + let mut results = Vec::with_capacity(loaded.len()); + for ( + part, + docs_for_wand, postings, wand_scorer, - metrics.as_ref(), partition_threshold, - )?; - std::result::Result::<_, Error>::Ok(candidates) + tokens_by_position, + grouped_expansions, + ) in loaded + { + let candidates = part.bm25_search( + docs_for_wand.as_ref(), + params.as_ref(), + operator, + mask.clone(), + postings, + wand_scorer, + metrics.as_ref(), + partition_threshold, + )?; + results.push(( + part, + PartitionCandidates { + tokens_by_position, + grouped_expansions, + candidates, + }, + )); + } + std::result::Result::<_, Error>::Ok(results) }) .await?; - let partition_result = PartitionCandidates { - tokens_by_position, - grouped_expansions, - candidates, - }; - Result::Ok((part, partition_result)) + Result::Ok(results) } }) .collect::>(); - let mut parts = stream::iter(parts).buffer_unordered(get_num_compute_intensive_cpus()); + let mut parts = stream::iter(parts) + .buffer_unordered(get_num_compute_intensive_cpus().min(32)) + .map_ok(|results| stream::iter(results.into_iter().map(Result::Ok))) + .try_flatten(); let mut idf_cache: HashMap = HashMap::new(); // Partitions that produced candidates, indexed by the `slot` carried // in deferred heap entries. From 84d8fe518da4c06bcb2676e9099220ddbf936a31 Mon Sep 17 00:00:00 2001 From: Lu Qiu Date: Thu, 23 Jul 2026 16:39:26 -0700 Subject: [PATCH 2/3] perf(index): bound in-chunk posting loads by store io parallelism --- rust/lance-index/src/scalar/inverted/index.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index d9d56f6ac5f..a47b5d8bb32 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -997,7 +997,9 @@ impl InvertedIndex { ))) } }); - let loaded: Vec<_> = futures::future::try_join_all(loads) + let loaded: Vec<_> = stream::iter(loads) + .buffer_unordered(self.store.io_parallelism()) + .try_collect::>() .await? .into_iter() .flatten() From b537edfb551a66bd4137a316ed102bf388f72deb Mon Sep 17 00:00:00 2001 From: Lu Qiu Date: Fri, 24 Jul 2026 10:32:25 -0700 Subject: [PATCH 3/3] fix(fts): give legacy partitions private top-k thresholds Legacy (no-impact) partitions score with partition-local BM25 statistics, so their scores are not comparable across partitions. Publishing them into a threshold shared by the whole query lets one partition's local k-th score prune another partition's true global winner before global rescoring. The sequential chunk loop made this pre-existing schedule-sensitive bug deterministic. Keep the shared floor only for the globally-comparable impact scorer and seed each legacy partition with a private floor. Regression: two legacy partitions in one chunk, asserting the global winner's row id and score. --- rust/lance-index/src/scalar/inverted/index.rs | 66 ++++++++++++++++--- 1 file changed, 58 insertions(+), 8 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index a47b5d8bb32..368660133ca 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -912,9 +912,11 @@ impl InvertedIndex { // Shared top-k floor across this query's partitions. Seeded to -inf so // the first real score wins; each partition publishes its local k-th // and prunes against the running global k-th (a lower bound on the true - // global k-th - see `Wand::shared_threshold`). + // global k-th - see `Wand::shared_threshold`). Only sound for the + // impact (global) scorer, whose scores are comparable across + // partitions; legacy BM25 scores are partition-local scale, so those + // partitions each get a private floor below. let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); - let legacy_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); // Partitions are processed in chunks: each chunk loads its postings // and scoring DocSets asynchronously, then ONE cpu-pool task searches // the whole chunk sequentially. Chunks pipeline against each other @@ -933,7 +935,6 @@ impl InvertedIndex { let metrics = metrics.clone(); let impact_scorer = impact_scorer.clone(); let impact_shared_threshold = impact_shared_threshold.clone(); - let legacy_shared_threshold = legacy_shared_threshold.clone(); async move { // Load the chunk's partitions concurrently; only the // search below is batched onto one cpu task. @@ -944,7 +945,6 @@ impl InvertedIndex { let metrics = metrics.clone(); let impact_scorer = impact_scorer.clone(); let impact_shared_threshold = impact_shared_threshold.clone(); - let legacy_shared_threshold = legacy_shared_threshold.clone(); async move { let loaded_postings = part .load_posting_lists( @@ -983,7 +983,7 @@ impl InvertedIndex { let partition_threshold = if use_global_scorer { impact_shared_threshold } else { - legacy_shared_threshold + Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())) }; let wand_scorer = use_global_scorer.then(|| impact_scorer.clone()); Result::Ok(Some(( @@ -10858,6 +10858,7 @@ mod tests { } async fn load_global_scoring_test_index( + first_partition_has_impacts: bool, second_partition_has_impacts: bool, ) -> (TempObjDir, Arc) { let tmpdir = TempObjDir::default(); @@ -10867,7 +10868,7 @@ mod tests { Arc::new(LanceCache::no_cache()), )); let partition_specs = [ - (0, 100, 5_000, 101..111, 5_000, true), + (0, 100, 5_000, 101..111, 5_000, first_partition_has_impacts), (1, 200, 1_000, 201..301, 1, second_partition_has_impacts), ]; for ( @@ -10968,7 +10969,7 @@ mod tests { // Partition 0 wins under its local corpus statistics but loses under // the global statistics. If its local score escapes into the shared // floor, partition 1 will incorrectly prune the real global winner. - let (_tmpdir, index) = load_global_scoring_test_index(true).await; + let (_tmpdir, index) = load_global_scoring_test_index(true, true).await; let first_partition = index .partitions .iter() @@ -11071,7 +11072,7 @@ mod tests { #[tokio::test] async fn test_mixed_impact_and_legacy_partitions_use_global_final_scores() { - let (_tmpdir, index) = load_global_scoring_test_index(false).await; + let (_tmpdir, index) = load_global_scoring_test_index(true, false).await; let impact_partition = index .partitions @@ -11128,6 +11129,55 @@ mod tests { ); } + #[tokio::test] + async fn test_two_legacy_partitions_keep_private_thresholds() { + // Both partitions are legacy (no impacts), so their BM25 statistics + // are partition-local scale. Partition 0's matching doc scores high + // under its own statistics while partition 1's matching doc (the true + // global winner, row 200) scores low under partition 1's local + // statistics. The chunked search runs both partitions sequentially in + // one chunk: if partition 0's local k-th score leaked into a shared + // threshold, partition 1 would prune row 200 before global rescoring + // and return the wrong winner. + let (_tmpdir, index) = load_global_scoring_test_index(false, false).await; + for partition in index.partitions.iter() { + let posting = partition + .inverted_list + .posting_list(0, false, &NoOpMetricsCollector) + .await + .unwrap(); + assert!(!posting.has_impacts()); + } + + let tokens = Arc::new(Tokens::new(vec!["alpha".to_string()], DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(1))); + let (row_ids, scores) = index + .bm25_search( + tokens.clone(), + params.clone(), + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + + assert_eq!(row_ids, vec![200]); + assert_eq!(scores.len(), 1); + let scorer = index + .bm25_base_scorer(tokens.as_ref(), params.as_ref()) + .await + .unwrap(); + let expected_score = scorer.query_weight("alpha") * scorer.doc_weight(1, 1_000); + assert!( + (scores[0] - expected_score).abs() < 1e-6, + "score: {}, expected global score: {}", + scores[0], + expected_score + ); + } + #[tokio::test] async fn test_and_query_returns_empty_when_exact_term_missing() { let tmpdir = TempObjDir::default();