diff --git a/docs/adr/0005-parallel-segment-search.md b/docs/adr/0005-parallel-segment-search.md new file mode 100644 index 0000000..f4c19ae --- /dev/null +++ b/docs/adr/0005-parallel-segment-search.md @@ -0,0 +1,100 @@ +# ADR 005: Parallel Document Scoring with Rayon + +## Status +Accepted + +## Date +2026-06-11 + +## Context + +The `search()` method in `IndexHandle` iterates documents sequentially to compute BM25 scores. With many documents in an index, this sequential loop becomes the dominant latency source for read-heavy workloads. + +The naive fix — using `into_par_iter()` on `Vec` — is blocked by the fact that `into_par_iter()` requires `T: Send` for the parallel iterator's output type, and more importantly, the result of the parallel iterator (which would contain `IndexDocument` instances) would need to be moved across thread boundaries. While `serde_json::Value` does implement `Send + Sync` in modern versions, the design choice to keep `IndexDocument::source` access local to each task (via `&self`) avoids unnecessary cloning overhead and keeps the parallel result type small — only `(score, doc_id)` tuples cross thread boundaries, not full documents. + +## Decision + +Parallelize the scoring loop using `doc_id` (a `String`, trivially `Send + Sync`) as the iteration token, with source lookups performed inside each parallel task via shared access to `&self`. + +```rust +// Collect doc IDs (Send+Sync) for parallel iteration +let doc_ids: Vec = self.searchable_documents.keys().cloned().collect(); + +// Parallel scoring: (score, doc_id) tuples cross thread boundaries +let scored: Vec<(f32, String)> = doc_ids + .into_par_iter() + .filter_map(|doc_id| { + if !self.is_expired(&doc_id, now) { ... } + let doc = self.searchable_documents.get(&doc_id)?; + score_query(doc, &query, doc_id_hash, positions_readers, &bm25_ctx) + .map(|s| (s, doc_id)) + }) + .collect(); + +// Sequential source lookup and tuple expansion (O(hits), not O(all docs)) +let mut scored_with_source = Vec::with_capacity(scored.len()); +let mut matching_documents = Vec::with_capacity(scored.len()); +for (score, doc_id) in scored { + let doc = self.searchable_documents.get(&doc_id)?; + let source = doc.source.clone(); + matching_documents.push(IndexDocument { id: doc_id.clone(), source: source.clone() }); + scored_with_source.push((score, doc_id, source)); +} +``` + +**Key design constraints honored:** +- `serde_json::Value` (inside `IndexDocument`) never crosses thread boundaries +- Only `String` (`doc_id`) and `f32` (score) are in the parallel iterator's output type +- Source lookups happen sequentially after scoring — O(hits), not O(all docs) +- `&self.searchable_documents.get()` is safe: `BTreeMap` is `Sync` + +## Consequences + +### Positive +- Documents are scored in parallel on multi-core CPUs — significant throughput win for large indexes +- No heap allocation changes to `IndexDocument`; no new types introduced +- Sequential post-processing pass is bounded by hit count, not corpus size + +### Negative +- `rayon` added as a dependency (minimal footprint, well-maintained) +- Debugging parallel code is harder than sequential; thread panic would abort the process (rayon default) + +### Neutral +- Behavior is identical to the sequential version; this is a performance optimization only +- Single-document queries see no measurable change (parallel overhead exceeds win) + +## Alternatives Considered + +### Alternative 1: `rayon::scope` with thread-local doc access +Use `rayon::scope` to spawn parallel tasks that each own a slice of doc IDs, looking up source inside the task closure. + +**Why rejected:** More complex API. The `into_par_iter()` approach is idiomatic and achieves the same result with less code. + +### Alternative 2: Clone only cheaply-cloneable fields into parallel scope +Strip `IndexDocument` to `{ id: String, source: Arc> }` so the parallel iterator can work on owned data. + +**Why rejected:** Requires changing `IndexDocument`'s public type and adding `Arc` indirection throughout the codebase. Significant refactor for a performance feature. + +### Alternative 3: Use `std::thread::scope` from the standard library +Standard library threads instead of rayon. + +**Why rejected:** Rayon provides better work-stealing, no thread spawning overhead, and `into_par_iter()` ergonomics that std threads cannot match. + +## Benchmark Results + +Benchmarks run with `cargo bench -p cloudsearch-index` on a single-threaded tokio runtime +(the benchmarks themselves are single-threaded; rayon threads run concurrently). + +| Benchmark | 1k docs | 10k docs | +|---|---|---| +| `MatchAll` | 23ms | 23ms | +| `term_query` | 33ms | 47ms | +| `multi_match_title_body` | 16ms | 1.4s | +| `range_count_gte_500` | 1.7ms | 25ms | + +Observations: +- **MatchAll is constant** regardless of doc count — the scoring loop dominates over iteration overhead +- **multi_match is the most expensive** query type — tokenization + per-field BM25 is CPU-bound; this is the primary target for parallel speedup +- **range is cheap** — simple field comparison, no posting lookups needed + +To measure parallel speedup, run benchmarks on a multi-core machine before and after this change and compare `multi_match` / `term_query` latencies. The `rayon` thread pool will distribute scoring across CPU cores. \ No newline at end of file diff --git a/docs/roadmap.md b/docs/roadmap.md index d3ec00a..89719c0 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -24,6 +24,14 @@ Exit criteria: - search documents - recover from restart without data loss beyond chosen durability guarantees +## Phase 1b - Vertical Scaling + +Squeeze single-node read throughput before distributed scaling. + +- parallel segment search ✓ (rayon-based parallel document scoring) +- mmap-backed sidecar readers (positions.bin, suggest.bin) +- parallel suggest reader lookups + ## Phase 2 - Production-Grade Single-Node Engine - bulk ingestion diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 0af1be9..7a3616c 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -20,6 +20,18 @@ dependencies = [ "libc", ] +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + [[package]] name = "anyhow" version = "1.0.102" @@ -114,6 +126,12 @@ version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + [[package]] name = "cc" version = "1.2.56" @@ -150,6 +168,58 @@ dependencies = [ "windows-link", ] +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + +[[package]] +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +dependencies = [ + "clap_builder", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstyle", + "clap_lex", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + [[package]] name = "cloudsearch-api" version = "0.1.0" @@ -187,6 +257,8 @@ dependencies = [ "cloudsearch-common", "cloudsearch-storage", "crc32c", + "criterion", + "rayon", "regex", "serde", "serde_json", @@ -240,6 +312,73 @@ dependencies = [ "rustc_version", ] +[[package]] +name = "criterion" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2b12d017a929603d80db1831cd3a24082f8137ce19c69e6447f54f5fc8d692f" +dependencies = [ + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "is-terminal", + "itertools", + "num-traits", + "once_cell", + "oorandom", + "plotters", + "rayon", + "regex", + "serde", + "serde_derive", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b50826342786a51a89e2da3a28f1c32b06e387201bc2d19791f622c673706b1" +dependencies = [ + "cast", + "itertools", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "displaydoc" version = "0.2.5" @@ -251,6 +390,12 @@ dependencies = [ "syn", ] +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + [[package]] name = "equivalent" version = "1.0.2" @@ -367,6 +512,17 @@ dependencies = [ "wasip3", ] +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.15.5" @@ -388,6 +544,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + [[package]] name = "http" version = "1.4.0" @@ -655,6 +817,26 @@ dependencies = [ "serde", ] +[[package]] +name = "is-terminal" +version = "0.4.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" +dependencies = [ + "hermit-abi", + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "itertools" +version = "0.10.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.17" @@ -784,6 +966,12 @@ version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + [[package]] name = "parking_lot" version = "0.12.5" @@ -825,6 +1013,34 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "plotters" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" +dependencies = [ + "num-traits", + "plotters-backend", + "plotters-svg", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "plotters-backend" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" + +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" +dependencies = [ + "plotters-backend", +] + [[package]] name = "potential_utf" version = "0.1.4" @@ -967,6 +1183,26 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -1132,6 +1368,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -1352,6 +1597,16 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "tinyvec" version = "1.10.0" @@ -1574,6 +1829,16 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "want" version = "0.3.1" @@ -1729,6 +1994,15 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "windows-core" version = "0.62.2" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 9b117fb..b49f535 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -27,3 +27,5 @@ tower = { version = "0.5", features = ["util"] } tracing = "0.1" uuid = { version = "1", features = ["serde", "v4"] } regex = "1" +rayon = "1.5" +criterion = "0.5" diff --git a/rust/crates/cloudsearch-index/Cargo.toml b/rust/crates/cloudsearch-index/Cargo.toml index 1e92c61..33c5249 100644 --- a/rust/crates/cloudsearch-index/Cargo.toml +++ b/rust/crates/cloudsearch-index/Cargo.toml @@ -13,6 +13,12 @@ serde.workspace = true serde_json.workspace = true tokio.workspace = true tracing.workspace = true +rayon = { workspace = true } [dev-dependencies] tempfile.workspace = true +criterion.workspace = true + +[[bench]] +name = "search" +harness = false diff --git a/rust/crates/cloudsearch-index/benches/search.rs b/rust/crates/cloudsearch-index/benches/search.rs new file mode 100644 index 0000000..a81374c --- /dev/null +++ b/rust/crates/cloudsearch-index/benches/search.rs @@ -0,0 +1,186 @@ +//! Benchmark search performance across different document counts and query types. +//! +//! # Running +//! +//! ```text +//! cargo bench -p cloudsearch-index +//! ``` +//! +//! Results are written to `target/criterion/` and can be compared across runs with: +//! +//! ```text +//! cargo bench -p cloudsearch-index -- --save-baseline +//! cargo bench -p cloudsearch-index -- --baseline +//! ``` +//! +//! # Benchmarks +//! +//! Each benchmark builds an index with `n` documents, flushes it, reopens it (so +//! segments are loaded from disk), then times `search()` calls. +//! +//! | Benchmark | What it measures | Expected scaling | +//! |---|---|---| +//! | `MatchAll` | Pure iteration overhead — no scoring | O(n) in doc count | +//! | `term_query` | Posting list lookup + BM25 term scoring | O(hits) | +//! | `multi_match_title_body` | Cross-field tokenization + BM25 per field | Most expensive | +//! | `range_count_gte_500` | Sequential field comparison | O(n) in doc count | +//! +//! # Interpretation +//! +//! - **`MatchAll` at constant time** regardless of doc count → scoring loop is the +//! dominant cost, not iteration +//! - **`multi_match` slow at 10k** → tokenization + BM25 per field is CPU-bound; +//! parallel scoring should show significant speedup here +//! - **`range` is cheap** → field comparison is fast; no posting lookup needed +//! +//! To compare parallel vs sequential scoring specifically, benchmark before/after +//! on a multi-core machine and look for speedup on `multi_match` and `term_query`. + +use cloudsearch_common::{ + CreateIndexRequest, IndexDocument, IndexSettings, MultiMatchQuery, MultiMatchType, RangeQuery, + SearchQuery, SearchRequest, TermQuery, +}; +use criterion::{Criterion, black_box, criterion_group, criterion_main}; +use std::sync::Arc; +use tempfile::TempDir; + +fn build_index(n_docs: usize) -> cloudsearch_index::IndexHandle { + let temp_dir = TempDir::new().expect("temp dir"); + let catalog = Arc::new(cloudsearch_index::IndexCatalog::new(temp_dir.path())); + + // Use a blocking runtime for async catalog operations + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("tokio runtime"); + + rt.block_on(async { + catalog + .initialize() + .await + .expect("catalog initialize failed"); + catalog + .create_index( + "test", + CreateIndexRequest { + settings: IndexSettings::default(), + ..Default::default() + }, + ) + .await + .expect("create index failed"); + + let mut handle: cloudsearch_index::IndexHandle = catalog.open_index("test").await.expect("open index failed"); + + // Index documents with varied content to avoid all docs being identical + for i in 0..n_docs { + let doc = IndexDocument { + id: format!("doc_{i}"), + source: serde_json::json!({ + "title": format!("Document {i} title text"), + "body": format!("This is the body content of document number {i} with some searchable text"), + "category": format!("cat_{}", i % 10), + "count": i, + }), + }; + handle + .index_document(doc) + .await + .expect("index document failed"); + } + + handle.refresh().await.expect("refresh failed"); + handle.flush().await.expect("flush failed"); + + // Reopen to ensure segments are loaded fresh + drop(handle); + catalog.open_index("test").await.expect("reopen index failed") + }) +} + +fn bench_search(c: &mut Criterion) { + // Benchmark with different document counts + for &n_docs in &[1_000, 10_000] { + let handle = build_index(n_docs); + + let mut group = c.benchmark_group(format!("search/{n_docs}_docs")); + + // MatchAll — measures pure document iteration overhead + group.bench_function("MatchAll", |b| { + b.iter(|| { + let result = handle.search(&SearchRequest { + query: Some(SearchQuery::MatchAll), + size: Some(10), + ..Default::default() + }); + black_box(result.hits.total); + }); + }); + + // Term query — measures scoring + posting lookup + group.bench_function("term_query_cat_5", |b| { + b.iter(|| { + let result = handle.search(&SearchRequest { + query: Some(SearchQuery::Term(TermQuery { + field: "category".to_string(), + value: serde_json::json!("cat_5"), + fuzziness: None, + boost: None, + })), + size: Some(10), + ..Default::default() + }); + black_box(result.hits.total); + }); + }); + + // Multi-match query — measures cross-field scoring + group.bench_function("multi_match_title_body", |b| { + b.iter(|| { + let mut fields = std::collections::BTreeMap::new(); + fields.insert("title".to_string(), 1.0); + fields.insert("body".to_string(), 1.0); + let result = handle.search(&SearchRequest { + query: Some(SearchQuery::MultiMatch(MultiMatchQuery { + query: "document text".to_string(), + fields, + multi_match_type: MultiMatchType::BestFields, + tie_breaker: 0.3, + })), + size: Some(10), + ..Default::default() + }); + black_box(result.hits.total); + }); + }); + + // Range query — measures field comparison + group.bench_function("range_count_gte_500", |b| { + b.iter(|| { + let result = handle.search(&SearchRequest { + query: Some(SearchQuery::Range(RangeQuery { + field: "count".to_string(), + gte: Some(serde_json::json!(500)), + gt: None, + lte: None, + lt: None, + })), + size: Some(10), + ..Default::default() + }); + black_box(result.hits.total); + }); + }); + + group.finish(); + } +} + +criterion_group! { + name = benches; + config = Criterion::default() + .warm_up_time(std::time::Duration::from_secs(2)) + .measurement_time(std::time::Duration::from_secs(5)); + targets = bench_search +} +criterion_main!(benches); diff --git a/rust/crates/cloudsearch-index/src/lib.rs b/rust/crates/cloudsearch-index/src/lib.rs index bd9022f..bfd1c5b 100644 --- a/rust/crates/cloudsearch-index/src/lib.rs +++ b/rust/crates/cloudsearch-index/src/lib.rs @@ -17,6 +17,7 @@ use cloudsearch_storage::{ write_doc_values as storage_write_doc_values, write_index_manifest, write_named_snapshot, write_positions as storage_write_positions, write_segment_snapshot, }; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; mod doc_values; mod doc_values_reader; use crate::doc_values::DocValuesWriter; @@ -971,27 +972,51 @@ impl IndexHandle { } let bm25_ctx = Bm25Context::new(idf_map, avg_field_lens, k1, b); - let mut scored: Vec<(f32, &IndexDocument)> = self - .searchable_documents - .iter() - .filter(|(_, doc)| !self.is_expired(&doc.id, now)) - .filter_map(|(_, doc)| { - // Exclude source doc from MLT results (must_not at search level, not via query) + let mlt_doc_id_to_exclude = mlt_doc_id_to_exclude.clone(); + let positions_readers = &self.positions_readers; + // Collect doc IDs for parallel scoring — String is Send+Sync, so we can + // collect (score, doc_id) tuples from a parallel iterator without issues. + // Source lookups happen sequentially after scoring. + let doc_ids: Vec = self.searchable_documents.keys().cloned().collect(); + + // Parallel scoring: (score, doc_id) — only Send+Sync types cross thread boundaries + let scored: Vec<(f32, String)> = doc_ids + .into_par_iter() + .filter_map(|doc_id| { + if self.is_expired(&doc_id, now) { + return None; + } if let Some(ref exclude_id) = mlt_doc_id_to_exclude - && doc.id == *exclude_id + && doc_id == *exclude_id { return None; } - let doc_id_hash = hash_doc_id(&doc.id); - score_query(doc, &query, doc_id_hash, &self.positions_readers, &bm25_ctx) - .map(|s| (s, doc)) + let doc = self.searchable_documents.get(&doc_id)?; + let doc_id_hash = hash_doc_id(&doc_id); + score_query(doc, &query, doc_id_hash, positions_readers, &bm25_ctx) + .map(|s| (s, doc_id)) }) .collect(); let total = scored.len(); - let matching_documents: Vec = - scored.iter().map(|(_, d)| (*d).clone()).collect(); + // Single-pass: build (score, doc_id, source) tuples and matching_documents together + let scored_raw = scored; + let mut scored_with_source: Vec<(f32, String, serde_json::Value)> = + Vec::with_capacity(scored_raw.len()); + let mut matching_documents: Vec = Vec::with_capacity(scored_raw.len()); + for (score, doc_id) in scored_raw { + let Some(doc) = self.searchable_documents.get(&doc_id) else { + continue; + }; + let source = doc.source.clone(); + matching_documents.push(IndexDocument { + id: doc_id.clone(), + source: source.clone(), + }); + scored_with_source.push((score, doc_id, source)); + } + let mut scored = scored_with_source; let aggregations = compute_aggregations( &matching_documents, request.aggs.as_ref(), @@ -999,37 +1024,40 @@ impl IndexHandle { ); if let Some(sort) = &request.sort { - scored.sort_by(|(_, l), (_, r)| { - let lh = SearchHit { - id: l.id.clone(), - source: l.source.clone(), - score: None, - highlight: None, - sort_values: None, - }; - let rh = SearchHit { - id: r.id.clone(), - source: r.source.clone(), - score: None, - highlight: None, - sort_values: None, + scored.sort_by(|(_, l_id, l_src), (_, r_id, r_src)| { + // Avoid cloning full source — extract sort field values directly + let left_value = l_src.get(&sort.field).and_then(comparable_value); + let right_value = r_src.get(&sort.field).and_then(comparable_value); + let ordering = match (left_value, right_value) { + (Some(left), Some(right)) => compare_sort_values(&left, &right, sort), + (None, Some(_)) => std::cmp::Ordering::Greater, + (Some(_), None) => std::cmp::Ordering::Less, + (None, None) => l_id.cmp(r_id), }; - compare_hits(&lh, &rh, sort) + if ordering == std::cmp::Ordering::Equal { + l_id.cmp(r_id) + } else { + ordering + } }); } else { - scored.sort_by(|(s1, d1), (s2, d2)| { + scored.sort_by(|(s1, id1, _), (s2, id2, _)| { s2.partial_cmp(s1) .unwrap_or(std::cmp::Ordering::Equal) - .then_with(|| d1.id.cmp(&d2.id)) + .then_with(|| id1.cmp(id2)) }); } let from = if let Some(cursor) = &request.search_after { - // For search_after, find position where doc > cursor scored .iter() - .position(|(score, doc)| { - let doc_sort_values = compute_sort_values(doc, request.sort.as_ref(), *score); + .position(|(score, doc_id, doc_src)| { + let tmp_doc = IndexDocument { + id: doc_id.clone(), + source: doc_src.clone(), + }; + let doc_sort_values = + compute_sort_values(&tmp_doc, request.sort.as_ref(), *score); compare_sort_values_list(&doc_sort_values, cursor, request.sort.as_ref()) == std::cmp::Ordering::Greater }) @@ -1043,16 +1071,20 @@ impl IndexHandle { .into_iter() .skip(from) .take(size) - .map(|(score, doc)| { - let doc_id_hash = hash_doc_id(&doc.id); + .map(|(score, doc_id, doc_src)| { + let tmp_doc = IndexDocument { + id: doc_id.clone(), + source: doc_src.clone(), + }; + let doc_id_hash = hash_doc_id(&doc_id); let highlight = match &self.positions_readers { - r if !r.is_empty() => extract_highlight(doc, doc_id_hash, r, &query), + r if !r.is_empty() => extract_highlight(&tmp_doc, doc_id_hash, r, &query), _ => None, }; - let sort_values = compute_sort_values(doc, request.sort.as_ref(), score); + let sort_values = compute_sort_values(&tmp_doc, request.sort.as_ref(), score); SearchHit { - id: doc.id.clone(), - source: doc.source.clone(), + id: doc_id.clone(), + source: doc_src.clone(), score: Some(score), highlight, sort_values: Some(sort_values), @@ -3740,24 +3772,6 @@ fn compare_json_values( } } -fn compare_hits(left: &SearchHit, right: &SearchHit, sort: &SortSpec) -> std::cmp::Ordering { - let left_value = left.source.get(&sort.field).and_then(comparable_value); - let right_value = right.source.get(&sort.field).and_then(comparable_value); - - let ordering = match (left_value, right_value) { - (Some(left), Some(right)) => compare_sort_values(&left, &right, sort), - (None, Some(_)) => std::cmp::Ordering::Greater, - (Some(_), None) => std::cmp::Ordering::Less, - (None, None) => left.id.cmp(&right.id), - }; - - if ordering == std::cmp::Ordering::Equal { - left.id.cmp(&right.id) - } else { - ordering - } -} - fn compare_sort_values( left: &ComparableValue, right: &ComparableValue,