use crate::directories::RamDirectory;
use crate::dsl::{Document, SchemaBuilder};
use crate::index::{Index, IndexConfig, IndexWriter};
use crate::query::SparseVectorQuery;
use crate::structures::{SparseFormat, SparseVectorConfig, WeightQuantization};
fn bmp_config() -> SparseVectorConfig {
SparseVectorConfig {
format: SparseFormat::Bmp,
weight_quantization: WeightQuantization::UInt8,
bmp_block_size: 64,
..SparseVectorConfig::default()
}
}
fn maxscore_config() -> SparseVectorConfig {
SparseVectorConfig {
format: SparseFormat::MaxScore,
weight_quantization: WeightQuantization::UInt8,
..SparseVectorConfig::default()
}
}
fn bmp_schema() -> (crate::dsl::Schema, crate::dsl::Field, crate::dsl::Field) {
let mut sb = SchemaBuilder::default();
let title = sb.add_text_field("title", true, true);
let sparse = sb.add_sparse_vector_field_with_config("sparse", true, true, bmp_config());
sb.set_reorder(sparse, true);
(sb.build(), title, sparse)
}
fn maxscore_schema() -> (crate::dsl::Schema, crate::dsl::Field, crate::dsl::Field) {
let mut sb = SchemaBuilder::default();
let title = sb.add_text_field("title", true, true);
let sparse = sb.add_sparse_vector_field_with_config("sparse", true, true, maxscore_config());
(sb.build(), title, sparse)
}
#[tokio::test]
async fn test_bmp_needle_in_haystack() {
let (schema, title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
.await
.unwrap();
for i in 0..100 {
let mut doc = Document::new();
doc.add_text(title, format!("Hay doc {}", i));
let entries: Vec<(u32, f32)> = (0..10)
.map(|d| (d, 0.1 + (i as f32 * 0.001) + (d as f32 * 0.01)))
.collect();
doc.add_sparse_vector(sparse, entries);
writer.add_document(doc).unwrap();
}
let mut needle = Document::new();
needle.add_text(title, "Needle BMP document");
needle.add_sparse_vector(sparse, vec![(1000, 0.9), (1001, 0.8), (1002, 0.7)]);
writer.add_document(needle).unwrap();
writer.commit().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
assert_eq!(index.num_docs().await.unwrap(), 101);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
let query = SparseVectorQuery::new(sparse, vec![(1000, 1.0), (1001, 1.0), (1002, 1.0)]);
let results = searcher.search(&query, 10).await.unwrap();
assert_eq!(results.len(), 1, "Only the needle has dims 1000-1002");
assert!(results[0].score > 0.0);
let doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
assert_eq!(
doc.get_first(title).unwrap().as_text().unwrap(),
"Needle BMP document"
);
let query_shared = SparseVectorQuery::new(sparse, vec![(5, 1.0)]);
let results = searcher.search(&query_shared, 200).await.unwrap();
assert!(
results.len() >= 50,
"Shared dim 5 should match many docs, got {}",
results.len()
);
let query_missing = SparseVectorQuery::new(sparse, vec![(99999, 1.0)]);
let results = searcher.search(&query_missing, 10).await.unwrap();
assert_eq!(results.len(), 0);
}
#[tokio::test]
async fn test_bmp_merge() {
let (schema, title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
.await
.unwrap();
for i in 0..30 {
let mut doc = Document::new();
doc.add_text(title, format!("seg1 hay {}", i));
doc.add_sparse_vector(sparse, vec![(0, 0.5), (1, 0.3), (2, 0.2)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
let mut needle = Document::new();
needle.add_text(title, "seg2 needle");
needle.add_sparse_vector(sparse, vec![(500, 0.95), (501, 0.85)]);
writer.add_document(needle).unwrap();
for i in 0..29 {
let mut doc = Document::new();
doc.add_text(title, format!("seg2 hay {}", i));
doc.add_sparse_vector(sparse, vec![(0, 0.4), (3, 0.6)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
assert_eq!(index.num_docs().await.unwrap(), 60);
let segments = index.segment_readers().await.unwrap();
assert!(segments.len() >= 2, "Should have at least 2 segments");
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
let query = SparseVectorQuery::new(sparse, vec![(500, 1.0), (501, 1.0)]);
let results = searcher.search(&query, 10).await.unwrap();
assert_eq!(results.len(), 1, "Pre-merge: needle should be found");
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.force_merge().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
assert_eq!(index.segment_readers().await.unwrap().len(), 1);
assert_eq!(index.num_docs().await.unwrap(), 60);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
let results = searcher.search(&query, 10).await.unwrap();
assert_eq!(results.len(), 1, "Post-merge: needle should still be found");
let doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
assert_eq!(
doc.get_first(title).unwrap().as_text().unwrap(),
"seg2 needle"
);
let query_hay = SparseVectorQuery::new(sparse, vec![(0, 1.0)]);
let results = searcher.search(&query_hay, 100).await.unwrap();
assert!(
results.len() >= 50,
"Post-merge: dim 0 should match >=50 docs, got {}",
results.len()
);
}
#[tokio::test]
async fn test_bmp_score_ranking() {
let (schema, title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
.await
.unwrap();
for i in 0..50 {
let mut doc = Document::new();
doc.add_text(title, format!("Doc weight {}", i));
let weight = (i + 1) as f32 / 50.0;
doc.add_sparse_vector(sparse, vec![(0, weight)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
let query = SparseVectorQuery::new(sparse, vec![(0, 1.0)]);
let results = searcher.search(&query, 10).await.unwrap();
assert_eq!(results.len(), 10);
for i in 1..results.len() {
assert!(
results[i - 1].score >= results[i].score,
"Results should be sorted descending: score[{}]={} < score[{}]={}",
i - 1,
results[i - 1].score,
i,
results[i].score
);
}
let top_doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
assert_eq!(
top_doc.get_first(title).unwrap().as_text().unwrap(),
"Doc weight 49"
);
}
#[tokio::test]
async fn test_bmp_vs_maxscore_equivalence() {
let (schema_bmp, _title_bmp, sparse_bmp) = bmp_schema();
let dir_bmp = RamDirectory::new();
let config = IndexConfig::default();
let mut writer_bmp = IndexWriter::create(dir_bmp.clone(), schema_bmp.clone(), config.clone())
.await
.unwrap();
let (schema_ms, _title_ms, sparse_ms) = maxscore_schema();
let dir_ms = RamDirectory::new();
let mut writer_ms = IndexWriter::create(dir_ms.clone(), schema_ms.clone(), config.clone())
.await
.unwrap();
let mut rng_state: u32 = 42;
for i in 0..200 {
let mut entries = Vec::new();
let num_dims = 5 + (rng_state % 11);
for _ in 0..num_dims {
rng_state = rng_state.wrapping_mul(1103515245).wrapping_add(12345);
let dim = rng_state % 1000;
rng_state = rng_state.wrapping_mul(1103515245).wrapping_add(12345);
let weight = (rng_state % 100) as f32 / 100.0;
if weight > 0.01 {
entries.push((dim, weight));
}
}
let mut doc_bmp = Document::new();
doc_bmp.add_text(_title_bmp, format!("doc {}", i));
doc_bmp.add_sparse_vector(sparse_bmp, entries.clone());
writer_bmp.add_document(doc_bmp).unwrap();
let mut doc_ms = Document::new();
doc_ms.add_text(_title_ms, format!("doc {}", i));
doc_ms.add_sparse_vector(sparse_ms, entries);
writer_ms.add_document(doc_ms).unwrap();
}
writer_bmp.commit().await.unwrap();
writer_ms.commit().await.unwrap();
let index_bmp = Index::open(dir_bmp, config.clone()).await.unwrap();
let index_ms = Index::open(dir_ms, config.clone()).await.unwrap();
let query_dims = vec![(42, 0.8), (100, 0.6), (200, 0.4), (500, 0.9)];
let reader_bmp = index_bmp.reader().await.unwrap();
let searcher_bmp = reader_bmp.searcher().await.unwrap();
let query_bmp = SparseVectorQuery::new(sparse_bmp, query_dims.clone());
let results_bmp = searcher_bmp.search(&query_bmp, 20).await.unwrap();
let reader_ms = index_ms.reader().await.unwrap();
let searcher_ms = reader_ms.searcher().await.unwrap();
let query_ms = SparseVectorQuery::new(sparse_ms, query_dims);
let results_ms = searcher_ms.search(&query_ms, 20).await.unwrap();
assert_eq!(
results_bmp.len(),
results_ms.len(),
"BMP and MaxScore should return same number of results: BMP={}, MS={}",
results_bmp.len(),
results_ms.len()
);
if !results_bmp.is_empty() {
let bmp_top = results_bmp[0].score;
let ms_top = results_ms[0].score;
let diff = (bmp_top - ms_top).abs();
assert!(
diff < 0.2 * ms_top.max(0.01),
"Top scores should be close: BMP={:.4}, MS={:.4}, diff={:.4}",
bmp_top,
ms_top,
diff
);
}
}
#[tokio::test]
async fn test_bmp_vs_maxscore_multi_ordinal() {
let (schema_bmp, _title_bmp, sparse_bmp) = bmp_schema();
let dir_bmp = RamDirectory::new();
let config = IndexConfig::default();
let mut writer_bmp = IndexWriter::create(dir_bmp.clone(), schema_bmp.clone(), config.clone())
.await
.unwrap();
let (schema_ms, _title_ms, sparse_ms) = maxscore_schema();
let dir_ms = RamDirectory::new();
let mut writer_ms = IndexWriter::create(dir_ms.clone(), schema_ms.clone(), config.clone())
.await
.unwrap();
let vectors_per_doc: Vec<Vec<Vec<(u32, f32)>>> = vec![
vec![
vec![(10, 0.9), (20, 0.3)],
vec![(10, 0.7), (30, 0.5)],
vec![(20, 0.8), (40, 0.2)],
],
vec![
vec![(10, 0.4), (30, 0.9)],
vec![(30, 0.8), (50, 0.3)],
vec![(10, 0.2), (60, 0.5)],
],
vec![
vec![(20, 0.95), (10, 0.1)],
vec![(20, 0.85)],
vec![(20, 0.6), (70, 0.4)],
],
];
let mut all_docs = vectors_per_doc.clone();
let mut rng: u32 = 777;
for _ in 0..50 {
let mut doc_vecs = Vec::new();
for _ in 0..3 {
rng = rng.wrapping_mul(1103515245).wrapping_add(12345);
let dim = 100 + (rng % 200);
rng = rng.wrapping_mul(1103515245).wrapping_add(12345);
let w = (rng % 50) as f32 / 100.0 + 0.01;
doc_vecs.push(vec![(dim, w)]);
}
all_docs.push(doc_vecs);
}
for vectors in &all_docs {
let mut doc_bmp = Document::new();
doc_bmp.add_text(_title_bmp, "doc");
for v in vectors {
doc_bmp.add_sparse_vector(sparse_bmp, v.clone());
}
writer_bmp.add_document(doc_bmp).unwrap();
let mut doc_ms = Document::new();
doc_ms.add_text(_title_ms, "doc");
for v in vectors {
doc_ms.add_sparse_vector(sparse_ms, v.clone());
}
writer_ms.add_document(doc_ms).unwrap();
}
writer_bmp.commit().await.unwrap();
writer_ms.commit().await.unwrap();
let index_bmp = Index::open(dir_bmp, config.clone()).await.unwrap();
let index_ms = Index::open(dir_ms, config.clone()).await.unwrap();
let query_dims = vec![(10, 1.0), (20, 0.5)];
let reader_bmp = index_bmp.reader().await.unwrap();
let searcher_bmp = reader_bmp.searcher().await.unwrap();
let q_bmp = SparseVectorQuery::new(sparse_bmp, query_dims.clone());
let results_bmp = searcher_bmp.search(&q_bmp, 10).await.unwrap();
let reader_ms = index_ms.reader().await.unwrap();
let searcher_ms = reader_ms.searcher().await.unwrap();
let q_ms = SparseVectorQuery::new(sparse_ms, query_dims);
let results_ms = searcher_ms.search(&q_ms, 10).await.unwrap();
assert_eq!(
results_bmp.len(),
results_ms.len(),
"Multi-ordinal: BMP returned {} results, MaxScore returned {}",
results_bmp.len(),
results_ms.len()
);
let n = results_bmp.len().min(5);
for i in 0..n {
let bmp_doc = results_bmp[i].doc_id;
let ms_doc = results_ms[i].doc_id;
let bmp_score = results_bmp[i].score;
let ms_score = results_ms[i].score;
let diff = (bmp_score - ms_score).abs();
assert!(
diff < 0.25 * ms_score.max(0.01),
"Multi-ordinal rank {}: BMP doc_id={} score={:.4}, MS doc_id={} score={:.4}, diff={:.4}",
i,
bmp_doc,
bmp_score,
ms_doc,
ms_score,
diff
);
}
}
#[tokio::test]
async fn test_bmp_many_blocks() {
let (schema, title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
.await
.unwrap();
for i in 0..500 {
let mut doc = Document::new();
doc.add_text(title, format!("Doc {}", i));
let dim = (i % 20) as u32;
let weight = 0.1 + (i as f32 / 500.0);
doc.add_sparse_vector(sparse, vec![(dim, weight), (100, 0.05)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
let index = Index::open(dir, config).await.unwrap();
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
let query = SparseVectorQuery::new(sparse, vec![(100, 1.0)]);
let results = searcher.search(&query, 600).await.unwrap();
assert!(
results.len() >= 400,
"Dim 100 should match most docs, got {}",
results.len()
);
let query = SparseVectorQuery::new(sparse, vec![(5, 1.0)]);
let results = searcher.search(&query, 100).await.unwrap();
assert!(
results.len() >= 20,
"Dim 5 should match ~25 docs, got {}",
results.len()
);
}
#[tokio::test]
async fn test_bmp_merge_exact_doc_ids() {
let (schema, title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
.await
.unwrap();
const DOCS_PER_SEG: usize = 100;
const NUM_SEGS: usize = 5;
for seg in 0..NUM_SEGS {
for i in 0..DOCS_PER_SEG {
let unique_dim = (seg * DOCS_PER_SEG + i) as u32;
let mut doc = Document::new();
doc.add_text(title, format!("seg{} doc{}", seg, i));
doc.add_sparse_vector(sparse, vec![(unique_dim, 1.0), (9999, 0.1)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
}
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.force_merge().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
assert_eq!(index.segment_readers().await.unwrap().len(), 1);
assert_eq!(
index.num_docs().await.unwrap() as usize,
DOCS_PER_SEG * NUM_SEGS
);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
for seg in 0..NUM_SEGS {
for i in 0..DOCS_PER_SEG {
let unique_dim = (seg * DOCS_PER_SEG + i) as u32;
let expected_title = format!("seg{} doc{}", seg, i);
let query = SparseVectorQuery::new(sparse, vec![(unique_dim, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
assert_eq!(
results.len(),
1,
"dim {} should match exactly 1 doc, got {}",
unique_dim,
results.len()
);
let doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
let got_title = doc.get_first(title).unwrap().as_text().unwrap().to_string();
assert_eq!(
got_title, expected_title,
"dim {} returned wrong doc: got '{}', expected '{}'",
unique_dim, got_title, expected_title
);
}
}
}
#[tokio::test]
async fn test_maxscore_merge() {
let (schema, title, sparse) = maxscore_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
.await
.unwrap();
const DOCS_PER_SEG: usize = 50;
const NUM_SEGS: usize = 3;
for seg in 0..NUM_SEGS {
for i in 0..DOCS_PER_SEG {
let unique_dim = (seg * DOCS_PER_SEG + i) as u32;
let mut doc = Document::new();
doc.add_text(title, format!("seg{} doc{}", seg, i));
doc.add_sparse_vector(sparse, vec![(unique_dim, 1.0), (9999, 0.1)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
}
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.force_merge().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
assert_eq!(index.segment_readers().await.unwrap().len(), 1);
assert_eq!(
index.num_docs().await.unwrap() as usize,
DOCS_PER_SEG * NUM_SEGS
);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
for seg in 0..NUM_SEGS {
for i in 0..DOCS_PER_SEG {
let unique_dim = (seg * DOCS_PER_SEG + i) as u32;
let expected_title = format!("seg{} doc{}", seg, i);
let query = SparseVectorQuery::new(sparse, vec![(unique_dim, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
assert_eq!(
results.len(),
1,
"MaxScore merge: dim {} should match exactly 1 doc, got {}",
unique_dim,
results.len()
);
let doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
let got_title = doc.get_first(title).unwrap().as_text().unwrap().to_string();
assert_eq!(
got_title, expected_title,
"MaxScore merge: dim {} returned wrong doc: got '{}', expected '{}'",
unique_dim, got_title, expected_title
);
}
}
let query = SparseVectorQuery::new(sparse, vec![(9999, 1.0)]);
let results = searcher.search(&query, 200).await.unwrap();
assert_eq!(
results.len(),
DOCS_PER_SEG * NUM_SEGS,
"MaxScore merge: dim 9999 should match all {} docs, got {}",
DOCS_PER_SEG * NUM_SEGS,
results.len()
);
}
#[tokio::test]
async fn test_bmp_multi_round_merge() {
let (schema, title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
.await
.unwrap();
for batch in 0..3 {
for i in 0..20 {
let mut doc = Document::new();
doc.add_text(title, format!("r1 b{} d{}", batch, i));
doc.add_sparse_vector(
sparse,
vec![(0, 0.5), ((batch * 10 + i % 5 + 1) as u32, 0.8)],
);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
}
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.force_merge().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
assert_eq!(index.num_docs().await.unwrap(), 60);
assert_eq!(index.segment_readers().await.unwrap().len(), 1);
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
for batch in 0..2 {
for i in 0..20 {
let mut doc = Document::new();
doc.add_text(title, format!("r2 b{} d{}", batch, i));
doc.add_sparse_vector(sparse, vec![(0, 0.3), (999, 0.9)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
}
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.force_merge().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
assert_eq!(index.num_docs().await.unwrap(), 100);
assert_eq!(index.segment_readers().await.unwrap().len(), 1);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
let query = SparseVectorQuery::new(sparse, vec![(0, 1.0)]);
let results = searcher.search(&query, 200).await.unwrap();
assert!(
results.len() >= 90,
"Dim 0 should match most docs after 2 merges, got {}",
results.len()
);
let query = SparseVectorQuery::new(sparse, vec![(999, 1.0)]);
let results = searcher.search(&query, 100).await.unwrap();
assert!(
results.len() >= 35,
"Dim 999 should match ~40 docs, got {}",
results.len()
);
}
#[tokio::test]
async fn test_bmp_merge_correctness() {
let mut sb = SchemaBuilder::default();
let title = sb.add_text_field("title", true, true);
let sparse = sb.add_sparse_vector_field_with_config("sparse", true, true, bmp_config());
let schema = sb.build();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
.await
.unwrap();
const DOCS_PER_SEG: usize = 100;
const NUM_SEGS: usize = 5;
for seg in 0..NUM_SEGS {
for i in 0..DOCS_PER_SEG {
let unique_dim = (seg * DOCS_PER_SEG + i) as u32;
let mut doc = Document::new();
doc.add_text(title, format!("seg{} doc{}", seg, i));
let topic_dim = 10000 + (seg as u32 * 100);
doc.add_sparse_vector(
sparse,
vec![(unique_dim, 1.0), (9999, 0.1), (topic_dim, 0.5)],
);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
}
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let segments = index.segment_readers().await.unwrap();
assert!(
segments.len() >= 5,
"Should have >= 5 segments before merge"
);
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.force_merge().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
assert_eq!(index.segment_readers().await.unwrap().len(), 1);
assert_eq!(
index.num_docs().await.unwrap() as usize,
DOCS_PER_SEG * NUM_SEGS
);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
let mut failures = Vec::new();
for seg in 0..NUM_SEGS {
for i in 0..DOCS_PER_SEG {
let unique_dim = (seg * DOCS_PER_SEG + i) as u32;
let expected_title = format!("seg{} doc{}", seg, i);
let query = SparseVectorQuery::new(sparse, vec![(unique_dim, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
if results.len() != 1 {
failures.push(format!(
"dim {}: expected 1 result, got {}",
unique_dim,
results.len()
));
continue;
}
let doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
let got_title = doc.get_first(title).unwrap().as_text().unwrap().to_string();
if got_title != expected_title {
failures.push(format!(
"dim {}: got '{}', expected '{}'",
unique_dim, got_title, expected_title
));
}
}
}
assert!(
failures.is_empty(),
"Merge correctness: {} failures:\n{}",
failures.len(),
failures[..failures.len().min(20)].join("\n")
);
let query = SparseVectorQuery::new(sparse, vec![(9999, 1.0)]);
let results = searcher.search(&query, 600).await.unwrap();
assert_eq!(
results.len(),
DOCS_PER_SEG * NUM_SEGS,
"Merge: dim 9999 should match all {} docs, got {}",
DOCS_PER_SEG * NUM_SEGS,
results.len()
);
}
#[tokio::test]
async fn test_bmp_merge_large() {
let mut sb = SchemaBuilder::default();
let title = sb.add_text_field("title", true, true);
let sparse = sb.add_sparse_vector_field_with_config("sparse", true, true, bmp_config());
let schema = sb.build();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
.await
.unwrap();
const DOCS_PER_SEG: usize = 500;
const NUM_SEGS: usize = 3;
for seg in 0..NUM_SEGS {
for i in 0..DOCS_PER_SEG {
let unique_dim = (seg * DOCS_PER_SEG + i) as u32;
let topic = i / 50;
let topic_dim = 20000 + (topic as u32 * 10);
let topic_dim2 = 20001 + (topic as u32 * 10);
let mut doc = Document::new();
doc.add_text(title, format!("s{}d{}", seg, i));
doc.add_sparse_vector(
sparse,
vec![
(unique_dim, 1.0),
(9999, 0.1),
(topic_dim, 0.8),
(topic_dim2, 0.5),
],
);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
}
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.force_merge().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
assert_eq!(index.segment_readers().await.unwrap().len(), 1);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
let mut failures = Vec::new();
for seg in 0..NUM_SEGS {
for i in 0..DOCS_PER_SEG {
let unique_dim = (seg * DOCS_PER_SEG + i) as u32;
let expected = format!("s{}d{}", seg, i);
let query = SparseVectorQuery::new(sparse, vec![(unique_dim, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
if results.len() != 1 {
failures.push(format!(
"dim {}: expected 1 result, got {}",
unique_dim,
results.len()
));
continue;
}
let doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
let got = doc.get_first(title).unwrap().as_text().unwrap().to_string();
if got != expected {
failures.push(format!(
"dim {}: got '{}', expected '{}'",
unique_dim, got, expected
));
}
}
}
assert!(
failures.is_empty(),
"Merge large: {} failures (of {}):\n{}",
failures.len(),
DOCS_PER_SEG * NUM_SEGS,
failures[..failures.len().min(30)].join("\n")
);
let query = SparseVectorQuery::new(sparse, vec![(20000, 1.0), (20001, 0.5)]);
let results = searcher.search(&query, 200).await.unwrap();
assert!(
results.len() >= 100,
"Topic query should match >=100 docs, got {}",
results.len()
);
}
#[tokio::test]
async fn test_bmp_reorder_accepts_segment_without_sparse_values() {
let (schema, title, _sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
let mut doc = Document::new();
doc.add_text(title, "text only");
writer.add_document(doc).unwrap();
writer.commit().await.unwrap();
writer.reorder().await.unwrap();
let index = Index::open(dir, config).await.unwrap();
assert_eq!(index.num_docs().await.unwrap(), 1);
}
#[tokio::test]
async fn test_bmp_reorder_standalone() {
let mut sb = SchemaBuilder::default();
let title = sb.add_text_field("title", true, true);
let sparse = sb.add_sparse_vector_field_with_config("sparse", true, true, bmp_config());
sb.set_reorder(sparse, true);
let schema = sb.build();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
.await
.unwrap();
const NUM_DOCS: usize = 200;
for i in 0..NUM_DOCS {
let unique_dim = i as u32;
let topic = i / 50;
let topic_dim = 10000 + (topic as u32 * 10);
let mut doc = Document::new();
doc.add_text(title, format!("doc{}", i));
doc.add_sparse_vector(
sparse,
vec![(unique_dim, 1.0), (9999, 0.1), (topic_dim, 0.5)],
);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
assert_eq!(index.segment_readers().await.unwrap().len(), 1);
assert_eq!(index.num_docs().await.unwrap() as usize, NUM_DOCS);
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.reorder().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
assert_eq!(index.segment_readers().await.unwrap().len(), 1);
assert_eq!(index.num_docs().await.unwrap() as usize, NUM_DOCS);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
let mut failures = Vec::new();
for i in 0..NUM_DOCS {
let unique_dim = i as u32;
let expected_title = format!("doc{}", i);
let query = SparseVectorQuery::new(sparse, vec![(unique_dim, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
if results.len() != 1 {
failures.push(format!(
"dim {}: expected 1 result, got {}",
unique_dim,
results.len()
));
continue;
}
let doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
let got_title = doc.get_first(title).unwrap().as_text().unwrap().to_string();
if got_title != expected_title {
failures.push(format!(
"dim {}: got '{}', expected '{}'",
unique_dim, got_title, expected_title
));
}
}
assert!(
failures.is_empty(),
"Standalone reorder: {} failures:\n{}",
failures.len(),
failures[..failures.len().min(20)].join("\n")
);
let query = SparseVectorQuery::new(sparse, vec![(9999, 1.0)]);
let results = searcher.search(&query, 300).await.unwrap();
assert_eq!(
results.len(),
NUM_DOCS,
"Reorder: dim 9999 should match all {} docs, got {}",
NUM_DOCS,
results.len()
);
}
#[tokio::test]
async fn test_bmp_reorder_after_merge_keeps_interior_padded_docs() {
let (schema, title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
.await
.unwrap();
const DOCS_PER_SEGMENT: usize = 100;
for seg in 0..2 {
for i in 0..DOCS_PER_SEGMENT {
let global = seg * DOCS_PER_SEGMENT + i;
let mut doc = Document::new();
doc.add_text(title, format!("doc{}", global));
doc.add_sparse_vector(sparse, vec![(global as u32, 1.0), (9999, 0.1)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
}
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.force_merge().await.unwrap();
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.reorder().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
assert_eq!(index.segment_readers().await.unwrap().len(), 1);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
let total = 2 * DOCS_PER_SEGMENT;
let mut failures = Vec::new();
for i in 0..total {
let query = SparseVectorQuery::new(sparse, vec![(i as u32, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
if results.len() != 1 {
failures.push(format!(
"dim {}: expected 1 result, got {}",
i,
results.len()
));
continue;
}
let doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
let got = doc.get_first(title).unwrap().as_text().unwrap().to_string();
if got != format!("doc{}", i) {
failures.push(format!("dim {}: got '{}', expected 'doc{}'", i, got, i));
}
}
assert!(
failures.is_empty(),
"Reorder after merge lost docs past interior padding: {} failures:\n{}",
failures.len(),
failures[..failures.len().min(20)].join("\n")
);
let query = SparseVectorQuery::new(sparse, vec![(9999, 1.0)]);
let results = searcher.search(&query, 300).await.unwrap();
assert_eq!(
results.len(),
total,
"dim 9999 should match all {} docs after reorder, got {}",
total,
results.len()
);
}
#[tokio::test]
async fn test_bmp_merge_with_reorder_on_merge() {
let mut sb = SchemaBuilder::default();
let title = sb.add_text_field("title", true, true);
let sparse = sb.add_sparse_vector_field_with_config("sparse", true, true, bmp_config());
sb.set_reorder(sparse, true);
sb.set_reorder_on_merge(true);
let schema = sb.build();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
.await
.unwrap();
const DOCS_PER_SEGMENT: usize = 100;
for seg in 0..2 {
for i in 0..DOCS_PER_SEGMENT {
let global = seg * DOCS_PER_SEGMENT + i;
let mut doc = Document::new();
doc.add_text(title, format!("doc{}", global));
doc.add_sparse_vector(sparse, vec![(global as u32, 1.0), (9999, 0.1)]);
if global.is_multiple_of(10) {
doc.add_sparse_vector(sparse, vec![(5000 + global as u32, 1.0)]);
}
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
}
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.force_merge().await.unwrap();
let metadata = crate::index::IndexMetadata::load(&dir).await.unwrap();
assert_eq!(metadata.segment_metas.len(), 1);
assert!(
metadata.segment_metas.values().all(|m| m.reordered),
"merge with reorder_on_merge must mark the output segment reordered"
);
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let total = 2 * DOCS_PER_SEGMENT;
assert_eq!(index.num_docs().await.unwrap() as usize, total);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
let mut failures = Vec::new();
for i in 0..total {
let mut dims_to_check = vec![i as u32];
if i.is_multiple_of(10) {
dims_to_check.push(5000 + i as u32);
}
for dim in dims_to_check {
let query = SparseVectorQuery::new(sparse, vec![(dim, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
if results.len() != 1 {
failures.push(format!(
"dim {}: expected 1 result, got {}",
dim,
results.len()
));
continue;
}
let doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
let got = doc.get_first(title).unwrap().as_text().unwrap().to_string();
if got != format!("doc{}", i) {
failures.push(format!("dim {}: got '{}', expected 'doc{}'", dim, got, i));
}
}
}
assert!(
failures.is_empty(),
"Merge-time reorder: {} failures:\n{}",
failures.len(),
failures[..failures.len().min(20)].join("\n")
);
let query = SparseVectorQuery::new(sparse, vec![(9999, 1.0)]);
let results = searcher.search(&query, 300).await.unwrap();
assert_eq!(
results.len(),
total,
"dim 9999 should match all {} docs after merge-time reorder, got {}",
total,
results.len()
);
}
#[tokio::test]
async fn test_bmp_reorder_skips_field_without_reorder_attribute() {
let mut sb = SchemaBuilder::default();
let title = sb.add_text_field("title", true, true);
let ordered = sb.add_sparse_vector_field_with_config("ordered", true, true, bmp_config());
let frozen = sb.add_sparse_vector_field_with_config("frozen", true, true, bmp_config());
sb.set_reorder(ordered, true);
let schema = sb.build();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
.await
.unwrap();
const NUM_DOCS: usize = 300;
for i in 0..NUM_DOCS {
let mut doc = Document::new();
doc.add_text(title, format!("doc{}", i));
let topic_dim = 10000 + (i % 2) as u32 * 10;
doc.add_sparse_vector(ordered, vec![(i as u32, 1.0), (topic_dim, 0.5)]);
doc.add_sparse_vector(frozen, vec![(i as u32, 1.0), (9999, 0.1)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let readers = index.segment_readers().await.unwrap();
assert_eq!(readers.len(), 1);
let frozen_doc_map_before: Vec<u8> = readers[0]
.bmp_indexes()
.get(&frozen.0)
.expect("frozen BMP index")
.doc_map_ids_slice()
.to_vec();
let ordered_doc_map_before: Vec<u8> = readers[0]
.bmp_indexes()
.get(&ordered.0)
.expect("ordered BMP index")
.doc_map_ids_slice()
.to_vec();
drop(readers);
drop(index);
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.reorder().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let readers = index.segment_readers().await.unwrap();
assert_eq!(readers.len(), 1);
let frozen_doc_map_after: Vec<u8> = readers[0]
.bmp_indexes()
.get(&frozen.0)
.expect("frozen BMP index after reorder")
.doc_map_ids_slice()
.to_vec();
let ordered_doc_map_after: Vec<u8> = readers[0]
.bmp_indexes()
.get(&ordered.0)
.expect("ordered BMP index after reorder")
.doc_map_ids_slice()
.to_vec();
drop(readers);
assert_eq!(
frozen_doc_map_before, frozen_doc_map_after,
"field without `reorder` attribute must be copied byte-identically"
);
assert_ne!(
ordered_doc_map_before, ordered_doc_map_after,
"field with `reorder` attribute must actually be BP-reordered"
);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
for (field, name) in [(ordered, "ordered"), (frozen, "frozen")] {
for i in [0usize, 63, 64, 150, NUM_DOCS - 1] {
let query = SparseVectorQuery::new(field, vec![(i as u32, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
assert_eq!(results.len(), 1, "{} dim {}: expected 1 result", name, i);
let doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
assert_eq!(
doc.get_first(title).unwrap().as_text().unwrap(),
format!("doc{}", i),
"{} dim {} returned wrong doc",
name,
i
);
}
}
}
#[tokio::test]
async fn test_bmp_reorder_multi_field() {
let mut sb = SchemaBuilder::default();
let title = sb.add_text_field("title", true, true);
let sparse_a = sb.add_sparse_vector_field_with_config("sparse_a", true, true, bmp_config());
let sparse_b = sb.add_sparse_vector_field_with_config("sparse_b", true, true, bmp_config());
sb.set_reorder(sparse_a, true);
sb.set_reorder(sparse_b, true);
let schema = sb.build();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
.await
.unwrap();
const NUM_DOCS: usize = 100;
for i in 0..NUM_DOCS {
let mut doc = Document::new();
doc.add_text(title, format!("doc{}", i));
doc.add_sparse_vector(sparse_a, vec![(i as u32, 1.0), (9999, 0.1)]);
doc.add_sparse_vector(sparse_b, vec![(1000 + i as u32, 1.0), (19999, 0.1)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.reorder().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
assert_eq!(index.segment_readers().await.unwrap().len(), 1);
assert_eq!(index.num_docs().await.unwrap() as usize, NUM_DOCS);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
let mut failures = Vec::new();
for i in 0..NUM_DOCS {
let query = SparseVectorQuery::new(sparse_a, vec![(i as u32, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
if results.len() != 1 {
failures.push(format!(
"field_a dim {}: expected 1 result, got {}",
i,
results.len()
));
continue;
}
let doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
let got = doc.get_first(title).unwrap().as_text().unwrap();
if got != format!("doc{}", i) {
failures.push(format!("field_a dim {}: got '{}'", i, got));
}
}
for i in 0..NUM_DOCS {
let query = SparseVectorQuery::new(sparse_b, vec![(1000 + i as u32, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
if results.len() != 1 {
failures.push(format!(
"field_b dim {}: expected 1 result, got {}",
1000 + i,
results.len()
));
continue;
}
let doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
let got = doc.get_first(title).unwrap().as_text().unwrap();
if got != format!("doc{}", i) {
failures.push(format!("field_b dim {}: got '{}'", 1000 + i, got));
}
}
assert!(
failures.is_empty(),
"Multi-field reorder: {} failures:\n{}",
failures.len(),
failures[..failures.len().min(20)].join("\n")
);
let query = SparseVectorQuery::new(sparse_a, vec![(9999, 1.0)]);
let results = searcher.search(&query, 200).await.unwrap();
assert_eq!(results.len(), NUM_DOCS);
let query = SparseVectorQuery::new(sparse_b, vec![(19999, 1.0)]);
let results = searcher.search(&query, 200).await.unwrap();
assert_eq!(results.len(), NUM_DOCS);
}
#[tokio::test]
async fn test_bmp_multi_ordinal_clustering() {
let mut sb = SchemaBuilder::default();
let title = sb.add_text_field("title", true, true);
let sparse = sb.add_sparse_vector_field_with_config("sparse", true, true, bmp_config());
let schema = sb.build();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema.clone(), config.clone())
.await
.unwrap();
const NUM_DOCS: usize = 200;
let mut rng: u32 = 42;
for i in 0..NUM_DOCS / 2 {
let mut doc = Document::new();
doc.add_text(title, format!("doc{}", i));
let unique_a = 5000 + i as u32;
rng = rng.wrapping_mul(1103515245).wrapping_add(12345);
let topic_a_dim = 1001 + (rng % 99);
rng = rng.wrapping_mul(1103515245).wrapping_add(12345);
let w = 0.3 + (rng % 70) as f32 / 100.0;
doc.add_sparse_vector(sparse, vec![(1000, 0.5), (topic_a_dim, w), (unique_a, 1.0)]);
let unique_b = 6000 + i as u32;
rng = rng.wrapping_mul(1103515245).wrapping_add(12345);
let topic_b_dim = 2001 + (rng % 99);
rng = rng.wrapping_mul(1103515245).wrapping_add(12345);
let w = 0.3 + (rng % 70) as f32 / 100.0;
doc.add_sparse_vector(sparse, vec![(2000, 0.5), (topic_b_dim, w), (unique_b, 1.0)]);
let unique_c = 7000 + i as u32;
rng = rng.wrapping_mul(1103515245).wrapping_add(12345);
let topic_c_dim = 3001 + (rng % 99);
rng = rng.wrapping_mul(1103515245).wrapping_add(12345);
let w = 0.3 + (rng % 70) as f32 / 100.0;
doc.add_sparse_vector(sparse, vec![(3000, 0.5), (topic_c_dim, w), (unique_c, 1.0)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
for i in NUM_DOCS / 2..NUM_DOCS {
let mut doc = Document::new();
doc.add_text(title, format!("doc{}", i));
let unique_a = 5000 + i as u32;
rng = rng.wrapping_mul(1103515245).wrapping_add(12345);
let topic_a_dim = 1001 + (rng % 99);
rng = rng.wrapping_mul(1103515245).wrapping_add(12345);
let w = 0.3 + (rng % 70) as f32 / 100.0;
doc.add_sparse_vector(sparse, vec![(1000, 0.5), (topic_a_dim, w), (unique_a, 1.0)]);
let unique_b = 6000 + i as u32;
rng = rng.wrapping_mul(1103515245).wrapping_add(12345);
let topic_b_dim = 2001 + (rng % 99);
rng = rng.wrapping_mul(1103515245).wrapping_add(12345);
let w = 0.3 + (rng % 70) as f32 / 100.0;
doc.add_sparse_vector(sparse, vec![(2000, 0.5), (topic_b_dim, w), (unique_b, 1.0)]);
let unique_c = 7000 + i as u32;
rng = rng.wrapping_mul(1103515245).wrapping_add(12345);
let topic_c_dim = 3001 + (rng % 99);
rng = rng.wrapping_mul(1103515245).wrapping_add(12345);
let w = 0.3 + (rng % 70) as f32 / 100.0;
doc.add_sparse_vector(sparse, vec![(3000, 0.5), (topic_c_dim, w), (unique_c, 1.0)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.force_merge().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
assert_eq!(index.segment_readers().await.unwrap().len(), 1);
assert_eq!(index.num_docs().await.unwrap() as usize, NUM_DOCS);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
let mut failures = Vec::new();
for i in 0..NUM_DOCS {
let expected = format!("doc{}", i);
let query = SparseVectorQuery::new(sparse, vec![(5000 + i as u32, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
if results.len() != 1 {
failures.push(format!(
"doc{} ord0 unique_dim={}: got {} results",
i,
5000 + i,
results.len()
));
} else {
let doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
let got = doc.get_first(title).unwrap().as_text().unwrap();
if got != expected {
failures.push(format!(
"doc{} ord0: got '{}', expected '{}'",
i, got, expected
));
}
}
let query = SparseVectorQuery::new(sparse, vec![(6000 + i as u32, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
if results.len() != 1 {
failures.push(format!(
"doc{} ord1 unique_dim={}: got {} results",
i,
6000 + i,
results.len()
));
} else {
let doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
let got = doc.get_first(title).unwrap().as_text().unwrap();
if got != expected {
failures.push(format!(
"doc{} ord1: got '{}', expected '{}'",
i, got, expected
));
}
}
let query = SparseVectorQuery::new(sparse, vec![(7000 + i as u32, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
if results.len() != 1 {
failures.push(format!(
"doc{} ord2 unique_dim={}: got {} results",
i,
7000 + i,
results.len()
));
} else {
let doc = searcher
.doc(results[0].segment_id, results[0].doc_id)
.await
.unwrap()
.unwrap();
let got = doc.get_first(title).unwrap().as_text().unwrap();
if got != expected {
failures.push(format!(
"doc{} ord2: got '{}', expected '{}'",
i, got, expected
));
}
}
}
assert!(
failures.is_empty(),
"Multi-ordinal clustering: {} failures:\n{}",
failures.len(),
failures[..failures.len().min(20)].join("\n")
);
let query = SparseVectorQuery::new(sparse, vec![(1000, 1.0)]);
let results_a = searcher.search(&query, 300).await.unwrap();
assert_eq!(
results_a.len(),
NUM_DOCS,
"Topic A anchor dim 1000 should match all {} docs, got {}",
NUM_DOCS,
results_a.len()
);
let query = SparseVectorQuery::new(sparse, vec![(1000, 0.8), (2000, 0.8)]);
let results_cross = searcher.search(&query, 300).await.unwrap();
assert_eq!(
results_cross.len(),
NUM_DOCS,
"Cross-topic query should match all {} docs, got {}",
NUM_DOCS,
results_cross.len()
);
}
#[cfg(feature = "metrics")]
#[tokio::test]
async fn test_bmp_query_emits_prometheus_metrics() {
use metrics_util::debugging::DebuggingRecorder;
let recorder = DebuggingRecorder::new();
let snapshotter = recorder.snapshotter();
recorder.install().expect("install debugging recorder");
let (mut schema, _title, sparse) = bmp_schema();
schema.set_index_name("metrics_test");
let tmp = tempfile::tempdir().unwrap();
let dir = crate::directories::MmapDirectory::new(tmp.path());
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
for seg in 0..2u32 {
for i in 0..200u32 {
let mut doc = Document::new();
doc.add_sparse_vector(sparse, vec![(seg * 200 + i, 1.0), (9999, 0.1)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
}
writer.force_merge().await.unwrap();
drop(writer);
let index = Index::open(dir, config).await.unwrap();
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
let query = SparseVectorQuery::new(sparse, vec![(9999, 1.0), (5, 0.5)]);
let results = searcher.search(&query, 10).await.unwrap();
assert!(!results.is_empty());
let snapshot = snapshotter.snapshot().into_vec();
for expected in [
"hermes_bmp_query_duration_seconds",
"hermes_bmp_blocks_scored_total",
"hermes_bmp_superblocks_visited_total",
"hermes_bmp_docmap_lookups_total",
"hermes_bmp_docmap_lookups_per_query",
"hermes_cold_write_bytes_total",
] {
assert!(
snapshot.iter().any(|(key, _, _, _)| {
let key = key.key();
key.name() == expected
&& key
.labels()
.any(|l| l.key() == "index" && l.value() == "metrics_test")
}),
"metric '{}' not emitted with index=\"metrics_test\"",
expected,
);
}
}
#[tokio::test]
async fn test_budgeted_reorder_marks_unconverged_and_stays_searchable() {
use crate::segment::BpBudget;
let (schema, title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
const NUM_DOCS: usize = 300;
for i in 0..NUM_DOCS {
let mut doc = Document::new();
doc.add_text(title, format!("doc{}", i));
let topic_dim = 10000 + (i % 2) as u32 * 10;
doc.add_sparse_vector(sparse, vec![(i as u32, 1.0), (topic_dim, 0.5)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
drop(writer);
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let seg_id = index.segment_readers().await.unwrap()[0].meta().id;
let sm = std::sync::Arc::clone(index.segment_manager());
let budget = BpBudget {
min_partition_docs: None,
time_budget: Some(std::time::Duration::ZERO),
};
let reordered = sm
.reorder_single_segment(&format!("{:032x}", seg_id), None, budget)
.await
.unwrap();
assert!(reordered);
let metadata = crate::index::IndexMetadata::load(&dir).await.unwrap();
assert_eq!(metadata.segment_metas.len(), 1);
let info = metadata.segment_metas.values().next().unwrap();
assert!(info.reordered, "budgeted pass must mark segment reordered");
assert!(
!info.bp_converged,
"zero-budget pass must be marked unconverged for the optimizer to deepen"
);
assert_eq!(info.bp_unconverged_passes, 1);
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
for i in [0usize, 63, 150, NUM_DOCS - 1] {
let query = SparseVectorQuery::new(sparse, vec![(i as u32, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
assert_eq!(results.len(), 1, "dim {} lost after budgeted reorder", i);
}
let new_seg_id = metadata.segment_ids()[0].clone();
let sm = std::sync::Arc::clone(index.segment_manager());
let reordered = sm
.reorder_single_segment(&new_seg_id, None, budget)
.await
.unwrap();
assert!(reordered);
let metadata = crate::index::IndexMetadata::load(&dir).await.unwrap();
let info = metadata.segment_metas.values().next().unwrap();
assert!(!info.bp_converged);
assert_eq!(info.bp_unconverged_passes, 2);
let new_seg_id = metadata.segment_ids()[0].clone();
let reordered = sm
.reorder_single_segment(&new_seg_id, None, BpBudget::full())
.await
.unwrap();
assert!(reordered);
let metadata = crate::index::IndexMetadata::load(&dir).await.unwrap();
let info = metadata.segment_metas.values().next().unwrap();
assert!(
info.bp_converged,
"full-budget follow-up pass must converge"
);
assert_eq!(info.bp_unconverged_passes, 0);
}
#[tokio::test]
async fn test_bmp_reorder_small_segment_clusters_low_df_topics() {
let (schema, _title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
const NUM_DOCS: usize = 200;
for i in 0..NUM_DOCS {
let mut doc = Document::new();
let topic_dim = 10000 + (i % 2) as u32 * 10;
doc.add_sparse_vector(sparse, vec![(i as u32, 1.0), (topic_dim, 0.5)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
drop(writer);
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let doc_map_before: Vec<u8> = index.segment_readers().await.unwrap()[0]
.bmp_indexes()
.get(&sparse.0)
.unwrap()
.doc_map_ids_slice()
.to_vec();
drop(index);
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.reorder().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let doc_map_after: Vec<u8> = index.segment_readers().await.unwrap()[0]
.bmp_indexes()
.get(&sparse.0)
.unwrap()
.doc_map_ids_slice()
.to_vec();
assert_ne!(
doc_map_before, doc_map_after,
"small-segment reorder must cluster low-df topic dims, not degenerate to identity"
);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
for i in [0usize, 99, NUM_DOCS - 1] {
let query = SparseVectorQuery::new(sparse, vec![(i as u32, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
assert_eq!(
results.len(),
1,
"dim {} lost after small-segment reorder",
i
);
}
}
#[tokio::test]
async fn test_auto_reorder_uses_blockwise_for_coherent_merged_segments() {
let (schema, _title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
const DOCS_PER_SEGMENT: usize = 4096;
for seg in 0..2 {
for i in 0..DOCS_PER_SEGMENT {
let global = seg * DOCS_PER_SEGMENT + i;
let frac = i * 100 / DOCS_PER_SEGMENT;
let topic: u32 = if seg == 0 {
match frac {
0..=39 => 0,
40..=69 => 1,
70..=89 => 2,
_ => 3,
}
} else {
match frac {
0..=39 => 3,
40..=69 => 2,
70..=89 => 1,
_ => 0,
}
};
let mut entries: Vec<(u32, f32)> =
(0..32).map(|t| (50000 + topic * 100 + t, 0.5)).collect();
entries.push((global as u32, 1.0));
let mut doc = Document::new();
doc.add_sparse_vector(sparse, entries);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
}
writer.reorder().await.unwrap();
writer.force_merge().await.unwrap();
drop(writer);
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let readers = index.segment_readers().await.unwrap();
assert_eq!(readers.len(), 1);
let bmp = readers[0].bmp_indexes().get(&sparse.0).unwrap();
let d = bmp.total_postings() as f32 / bmp.total_terms() as f32;
assert!(
d >= 4.0,
"test setup: expected coherent blocks after record reorder + merge, got d={:.2}",
d
);
let ids_before = bmp.doc_map_ids_slice().to_vec();
drop(readers);
drop(index);
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.reorder().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let readers = index.segment_readers().await.unwrap();
let bmp = readers[0].bmp_indexes().get(&sparse.0).unwrap();
let ids_after = bmp.doc_map_ids_slice().to_vec();
const CHUNK_BYTES: usize = 64 * 4;
let chunks = |v: &[u8]| -> Vec<Vec<u8>> {
let mut c: Vec<Vec<u8>> = v.chunks(CHUNK_BYTES).map(|x| x.to_vec()).collect();
c.sort();
c
};
assert_eq!(
chunks(&ids_before),
chunks(&ids_after),
"blockwise reorder must move blocks as intact units"
);
assert_ne!(
ids_before, ids_after,
"blockwise reorder should actually permute block order"
);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
for i in [0usize, 4095, 4096, 8191] {
let query = SparseVectorQuery::new(sparse, vec![(i as u32, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
assert_eq!(results.len(), 1, "dim {} lost after blockwise reorder", i);
}
}
fn bmp_schema_with_config(
config: SparseVectorConfig,
) -> (crate::dsl::Schema, crate::dsl::Field, crate::dsl::Field) {
let mut sb = SchemaBuilder::default();
let title = sb.add_text_field("title", true, true);
let sparse = sb.add_sparse_vector_field_with_config("sparse", true, true, config);
sb.set_reorder(sparse, true);
(sb.build(), title, sparse)
}
#[tokio::test]
async fn test_bmp_block_size_256_build_merge_reorder_search() {
let (schema, _title, sparse) = bmp_schema_with_config(SparseVectorConfig {
bmp_block_size: 256,
..bmp_config()
});
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
const DOCS_PER_SEGMENT: usize = 600; for seg in 0..2 {
for i in 0..DOCS_PER_SEGMENT {
let global = (seg * DOCS_PER_SEGMENT + i) as u32;
let mut doc = Document::new();
doc.add_sparse_vector(sparse, vec![(global, 1.0), (900_000 + global % 4, 0.5)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
}
writer.force_merge().await.unwrap();
writer.reorder().await.unwrap();
drop(writer);
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let readers = index.segment_readers().await.unwrap();
assert_eq!(readers.len(), 1);
let bmp = readers[0].bmp_indexes().get(&sparse.0).unwrap();
assert_eq!(bmp.bmp_block_size, 256);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
for i in [0usize, 255, 256, 599, 600, 2 * DOCS_PER_SEGMENT - 1] {
let query = SparseVectorQuery::new(sparse, vec![(i as u32, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
assert_eq!(results.len(), 1, "dim {} lost at block_size 256", i);
}
}
const RARE_DIM_CORPUS_DOCS: usize = 8192;
fn add_rare_dim_clustered_corpus(
writer: &mut IndexWriter<RamDirectory>,
sparse: crate::dsl::Field,
) {
for i in 0..RARE_DIM_CORPUS_DOCS {
let group = i / 4;
let topic = (group % 32) as u32;
let mut entries: Vec<(u32, f32)> = (0..4)
.map(|k| (100_000 + (group as u32) * 4 + k, 0.5))
.collect();
entries.push((90_000 + topic, 0.5));
entries.push((i as u32, 1.0));
let mut doc = Document::new();
doc.add_sparse_vector(sparse, entries);
writer.add_document(doc).unwrap();
}
}
fn doc_map_chunks(v: &[u8]) -> Vec<Vec<u8>> {
const CHUNK_BYTES: usize = 64 * 4;
let mut c: Vec<Vec<u8>> = v.chunks(CHUNK_BYTES).map(|x| x.to_vec()).collect();
c.sort();
c
}
#[tokio::test]
async fn test_auto_reorder_rare_dim_clustered_low_d_picks_blockwise() {
let (schema, _title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
add_rare_dim_clustered_corpus(&mut writer, sparse);
writer.commit().await.unwrap();
drop(writer);
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let readers = index.segment_readers().await.unwrap();
assert_eq!(readers.len(), 1);
let bmp = readers[0].bmp_indexes().get(&sparse.0).unwrap();
let d = bmp.total_postings() as f32 / bmp.total_terms() as f32;
assert!(
d < 4.0,
"test setup: expected rare-dim corpus with low absolute d, got d={:.2}",
d
);
let ids_before = bmp.doc_map_ids_slice().to_vec();
drop(readers);
drop(index);
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.reorder().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let readers = index.segment_readers().await.unwrap();
let bmp = readers[0].bmp_indexes().get(&sparse.0).unwrap();
let ids_after = bmp.doc_map_ids_slice().to_vec();
assert_eq!(
doc_map_chunks(&ids_before),
doc_map_chunks(&ids_after),
"coherent rare-dim segment must take the blockwise path"
);
assert_ne!(
ids_before, ids_after,
"blockwise reorder should actually permute block order"
);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
for i in [0usize, 63, 4096, 8191] {
let query = SparseVectorQuery::new(sparse, vec![(i as u32, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
assert_eq!(results.len(), 1, "dim {} lost after blockwise reorder", i);
}
}
#[tokio::test]
async fn test_deepening_pass_on_unconverged_segment_forces_record_level() {
use crate::segment::BpBudget;
let (schema, _title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
add_rare_dim_clustered_corpus(&mut writer, sparse);
writer.commit().await.unwrap();
drop(writer);
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let seg_id = index.segment_readers().await.unwrap()[0].meta().id;
let sm = std::sync::Arc::clone(index.segment_manager());
let budget = BpBudget {
min_partition_docs: None,
time_budget: Some(std::time::Duration::ZERO),
};
assert!(
sm.reorder_single_segment(&format!("{:032x}", seg_id), None, budget)
.await
.unwrap()
);
let metadata = crate::index::IndexMetadata::load(&dir).await.unwrap();
let info = metadata.segment_metas.values().next().unwrap();
assert!(info.reordered);
assert!(
!info.bp_converged,
"zero-budget pass must be marked unconverged"
);
drop(index);
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let ids_partial = index.segment_readers().await.unwrap()[0]
.bmp_indexes()
.get(&sparse.0)
.unwrap()
.doc_map_ids_slice()
.to_vec();
let sm = std::sync::Arc::clone(index.segment_manager());
let new_seg_id = metadata.segment_ids()[0].clone();
assert!(
sm.reorder_single_segment(&new_seg_id, None, BpBudget::full())
.await
.unwrap()
);
let index2 = Index::open(dir.clone(), config.clone()).await.unwrap();
let ids_deepened = index2.segment_readers().await.unwrap()[0]
.bmp_indexes()
.get(&sparse.0)
.unwrap()
.doc_map_ids_slice()
.to_vec();
assert_ne!(
doc_map_chunks(&ids_partial),
doc_map_chunks(&ids_deepened),
"deepening pass must run record-level BP (scatter records), not blockwise"
);
let metadata = crate::index::IndexMetadata::load(&dir).await.unwrap();
let info = metadata.segment_metas.values().next().unwrap();
assert!(
info.bp_converged,
"full-budget deepening pass must converge"
);
let reader = index2.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
for i in [0usize, 63, 4096, 8191] {
let query = SparseVectorQuery::new(sparse, vec![(i as u32, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
assert_eq!(results.len(), 1, "dim {} lost after deepening pass", i);
}
}
#[tokio::test]
async fn test_blockwise_budget_depth_cap_scales_to_block_units() {
use crate::segment::BpBudget;
let (schema, _title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
add_rare_dim_clustered_corpus(&mut writer, sparse);
writer.commit().await.unwrap();
drop(writer);
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let seg_id = index.segment_readers().await.unwrap()[0].meta().id;
let ids_before = index.segment_readers().await.unwrap()[0]
.bmp_indexes()
.get(&sparse.0)
.unwrap()
.doc_map_ids_slice()
.to_vec();
let sm = std::sync::Arc::clone(index.segment_manager());
let budget = BpBudget {
min_partition_docs: Some(4096),
time_budget: None,
};
assert!(
sm.reorder_single_segment(&format!("{:032x}", seg_id), None, budget)
.await
.unwrap()
);
let index2 = Index::open(dir.clone(), config.clone()).await.unwrap();
let ids_after = index2.segment_readers().await.unwrap()[0]
.bmp_indexes()
.get(&sparse.0)
.unwrap()
.doc_map_ids_slice()
.to_vec();
assert_eq!(
doc_map_chunks(&ids_before),
doc_map_chunks(&ids_after),
"coherent corpus must take the blockwise path"
);
assert_ne!(
ids_before, ids_after,
"a 4096-doc depth cap = superblock depth for blockwise BP — the pass must actually permute blocks"
);
let metadata = crate::index::IndexMetadata::load(&dir).await.unwrap();
let info = metadata.segment_metas.values().next().unwrap();
assert!(
!info.bp_converged,
"depth-capped pass must be recorded unconverged for the deepening ladder"
);
}
#[tokio::test]
async fn test_budgeted_merge_marks_unconverged_and_deepens() {
use crate::segment::BpBudget;
let mut sb = SchemaBuilder::default();
let _title = sb.add_text_field("title", true, true);
let sparse = sb.add_sparse_vector_field_with_config("sparse", true, true, bmp_config());
sb.set_reorder(sparse, true);
sb.set_reorder_on_merge(true);
let schema = sb.build();
let dir = RamDirectory::new();
let config = IndexConfig {
merge_bp_time_budget: Some(std::time::Duration::ZERO),
..IndexConfig::default()
};
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
const DOCS_PER_SEGMENT: usize = 300;
for seg in 0..2 {
for i in 0..DOCS_PER_SEGMENT {
let global = (seg * DOCS_PER_SEGMENT + i) as u32;
let topic = global % 4;
let mut entries: Vec<(u32, f32)> =
(0..32).map(|t| (50_000 + topic * 100 + t, 0.5)).collect();
entries.push((global, 1.0));
let mut doc = Document::new();
doc.add_sparse_vector(sparse, entries);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
}
writer.force_merge().await.unwrap();
drop(writer);
let metadata = crate::index::IndexMetadata::load(&dir).await.unwrap();
assert_eq!(metadata.segment_metas.len(), 1);
let info = metadata.segment_metas.values().next().unwrap();
assert!(info.reordered, "merged segment must be marked reordered");
assert!(
!info.bp_converged,
"budget-truncated merge BP must be recorded unconverged"
);
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
for i in [0usize, 299, 300, 2 * DOCS_PER_SEGMENT - 1] {
let query = SparseVectorQuery::new(sparse, vec![(i as u32, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
assert_eq!(results.len(), 1, "dim {} lost after budgeted merge", i);
}
let sm = std::sync::Arc::clone(index.segment_manager());
let seg_id = metadata.segment_ids()[0].clone();
assert!(
sm.reorder_single_segment(&seg_id, None, BpBudget::full())
.await
.unwrap()
);
let metadata = crate::index::IndexMetadata::load(&dir).await.unwrap();
let info = metadata.segment_metas.values().next().unwrap();
assert!(
info.bp_converged,
"full-depth deepening pass must converge the merged segment"
);
}
#[tokio::test]
async fn test_auto_reorder_interleaved_high_d_picks_records() {
let (schema, _title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
const NUM_DOCS: usize = 8192;
for i in 0..NUM_DOCS {
let topic = (i % 4) as u32;
let mut entries: Vec<(u32, f32)> =
(0..32).map(|t| (50_000 + topic * 100 + t, 0.5)).collect();
entries.push((i as u32, 1.0));
let mut doc = Document::new();
doc.add_sparse_vector(sparse, entries);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
drop(writer);
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let readers = index.segment_readers().await.unwrap();
assert_eq!(readers.len(), 1);
let bmp = readers[0].bmp_indexes().get(&sparse.0).unwrap();
let d = bmp.total_postings() as f32 / bmp.total_terms() as f32;
assert!(
d >= 4.0,
"test setup: expected interleaved corpus with high absolute d, got d={:.2}",
d
);
let ids_before = bmp.doc_map_ids_slice().to_vec();
drop(readers);
drop(index);
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.reorder().await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let readers = index.segment_readers().await.unwrap();
let bmp = readers[0].bmp_indexes().get(&sparse.0).unwrap();
let ids_after = bmp.doc_map_ids_slice().to_vec();
assert_ne!(
doc_map_chunks(&ids_before),
doc_map_chunks(&ids_after),
"interleaved segment must take the record-level path (blockwise cannot fix intra-block scramble)"
);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
for i in [0usize, 63, 4096, 8191] {
let query = SparseVectorQuery::new(sparse, vec![(i as u32, 1.0)]);
let results = searcher.search(&query, 5).await.unwrap();
assert_eq!(
results.len(),
1,
"dim {} lost after record-level reorder",
i
);
}
}
#[cfg(feature = "metrics")]
#[tokio::test(flavor = "multi_thread")]
#[ignore = "quantization experiment — run manually in release"]
async fn bench_aggressive_quantization() {
use crate::directories::{Directory, DirectoryWriter};
use metrics_util::debugging::{DebugValue, DebuggingRecorder};
struct XorShift(u64);
impl XorShift {
fn next(&mut self) -> u64 {
let mut x = self.0;
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
self.0 = x;
x
}
fn next_f64(&mut self) -> f64 {
(self.next() >> 11) as f64 / (1u64 << 53) as f64
}
}
let recorder = DebuggingRecorder::new();
let snapshotter = recorder.snapshotter();
recorder.install().expect("install debugging recorder");
let counters = || -> (u64, u64, u64) {
let mut scored = 0u64;
let mut skipped = 0u64;
let mut sbs = 0u64;
for (k, _, _, v) in snapshotter.snapshot().into_vec() {
if let DebugValue::Counter(c) = v {
match k.key().name() {
"hermes_bmp_blocks_scored_total" => scored += c,
"hermes_bmp_blocks_skipped_total" => skipped += c,
"hermes_bmp_superblocks_visited_total" => sbs += c,
_ => {}
}
}
}
(scored, skipped, sbs)
};
const DOCS: usize = 400_000;
const DIMS: usize = 100_000;
const TOPICS: usize = 256;
const TOPIC_DIMS: usize = 32; const MAX_W: f32 = 5.0;
const K: usize = 10;
const QUERIES: usize = 100;
let bg_base = TOPICS * TOPIC_DIMS;
let bg_dims = DIMS - bg_base;
let mut cum = Vec::with_capacity(bg_dims);
let mut acc = 0.0f64;
for i in 0..bg_dims {
acc += 1.0 / (i + 1) as f64;
cum.push(acc);
}
let total = acc;
let gen_doc = |rng: &mut XorShift, snap4: bool| -> Vec<(u32, f32)> {
let topic = (rng.next() % TOPICS as u64) as usize;
let mut by_dim: std::collections::BTreeMap<u32, f32> = std::collections::BTreeMap::new();
for _ in 0..10 {
let d = (topic * TOPIC_DIMS + (rng.next() % TOPIC_DIMS as u64) as usize) as u32;
let w = 1.5 + 3.5 * rng.next_f64() as f32;
let e = by_dim.entry(d).or_insert(0.0);
if w > *e {
*e = w;
}
}
for _ in 0..40 {
let r = rng.next_f64() * total;
let d = (bg_base + cum.partition_point(|&c| c < r).min(bg_dims - 1)) as u32;
let idf = ((bg_dims as f64 / ((d as usize - bg_base) as f64 + 1.0)).ln()
/ (bg_dims as f64).ln()) as f32;
let w = 1.2 * idf * (0.3 + 0.7 * rng.next_f64() as f32) + 0.05;
let e = by_dim.entry(d).or_insert(0.0);
if w > *e {
*e = w;
}
}
by_dim
.into_iter()
.map(|(d, w)| {
let w = if snap4 {
let q = (w / MAX_W * 255.0).round().clamp(0.0, 255.0);
let m = (q / 17.0).round() * 17.0;
m / 255.0 * MAX_W
} else {
w
};
(d, w)
})
.filter(|&(_, w)| w > 0.0)
.collect()
};
type GenDoc<'a> = dyn Fn(&mut XorShift, bool) -> Vec<(u32, f32)> + 'a;
async fn build(
dir: &RamDirectory,
sparse: crate::dsl::Field,
snap4: bool,
gen_doc: &GenDoc<'_>,
) {
let (mut schema, _t, _s) = bmp_schema();
schema.set_index_name("quant");
let config = IndexConfig {
merge_policy: Box::new(crate::merge::NoMergePolicy),
num_indexing_threads: 1,
max_indexing_memory_bytes: 8 * 1024 * 1024 * 1024,
..IndexConfig::default()
};
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
let mut rng = XorShift(0x5EEDFACE0BADC0DE); for _ in 0..DOCS {
let entries = gen_doc(&mut rng, snap4);
loop {
let mut doc = Document::new();
doc.add_sparse_vector(sparse, entries.clone());
match writer.add_document(doc) {
Ok(()) => break,
Err(crate::Error::QueueFull) => {
tokio::time::sleep(std::time::Duration::from_millis(3)).await;
}
Err(e) => panic!("{e:?}"),
}
}
}
writer.commit().await.unwrap();
writer.force_merge().await.unwrap();
writer.reorder().await.unwrap();
drop(writer);
}
let mut qrng = XorShift(0x0123456789ABCDEF);
let queries: Vec<Vec<(u32, f32)>> = (0..QUERIES)
.map(|_| {
let topic = (qrng.next() % TOPICS as u64) as usize;
let mut q: Vec<(u32, f32)> = (0..8)
.map(|_| {
let d =
(topic * TOPIC_DIMS + (qrng.next() % TOPIC_DIMS as u64) as usize) as u32;
(d, 1.5 + 3.5 * qrng.next_f64() as f32)
})
.collect();
for _ in 0..4 {
let r = qrng.next_f64() * total;
let d = (bg_base + cum.partition_point(|&c| c < r).min(bg_dims - 1)) as u32;
q.push((d, 0.5 + qrng.next_f64() as f32));
}
q.sort_by_key(|&(d, _)| d);
q.dedup_by_key(|&mut (d, _)| d);
q
})
.collect();
let bg_queries: Vec<Vec<(u32, f32)>> = (0..50)
.map(|_| {
let mut q: Vec<(u32, f32)> = (0..12)
.map(|_| {
let r = qrng.next_f64() * total;
let d = (bg_base + cum.partition_point(|&c| c < r).min(bg_dims - 1)) as u32;
let idf = ((bg_dims as f64 / ((d as usize - bg_base) as f64 + 1.0)).ln()
/ (bg_dims as f64).ln()) as f32;
(d, 2.0 * idf * (0.3 + 0.7 * qrng.next_f64() as f32) + 0.1)
})
.collect();
q.sort_by_key(|&(d, _)| d);
q.dedup_by_key(|&mut (d, _)| d);
q
})
.collect();
let run_queries = |bmps: Vec<crate::segment::BmpIndex>,
queries: &[Vec<(u32, f32)>]|
-> (Vec<Vec<(u32, f32)>>, f64) {
let t = std::time::Instant::now();
let mut out = Vec::with_capacity(queries.len());
for q in queries {
let mut all: Vec<(u32, f32)> = Vec::new();
for bmp in &bmps {
let r =
crate::query::bmp::execute_bmp(bmp, "quant", "sparse", q, K, 1.0, 0).unwrap();
all.extend(r.iter().map(|d| (d.doc_id, d.score)));
}
all.sort_by(|a, b| b.1.total_cmp(&a.1));
all.truncate(K);
out.push(all);
}
(out, t.elapsed().as_secs_f64())
};
let dir = RamDirectory::new();
let (_s0, _t0, sparse) = bmp_schema();
build(&dir, sparse, false, &gen_doc).await;
let config = IndexConfig::default();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let segs = index.segment_readers().await.unwrap();
assert_eq!(segs.len(), 1);
let seg_id = segs[0].meta().id;
let bmp = segs[0].bmp_indexes().get(&sparse.0).unwrap().clone();
let num_blocks = bmp.num_blocks as u64;
drop(segs);
drop(index);
let sparse_path = crate::segment::SegmentFiles::new(seg_id).sparse;
let original = dir
.open_read(&sparse_path)
.await
.unwrap()
.read_bytes()
.await
.unwrap()
.as_slice()
.to_vec();
let flen = original.len();
assert_eq!(
u32::from_le_bytes(original[flen - 4..].try_into().unwrap()),
crate::segment::format::SPARSE_FOOTER_MAGIC
);
let blob_len = u64::from_le_bytes(original[flen - 24..flen - 16].try_into().unwrap()) as usize;
let bfoot = blob_len - 64;
assert_eq!(
u32::from_le_bytes(original[bfoot + 60..bfoot + 64].try_into().unwrap()),
crate::segment::format::BMP_BLOB_MAGIC_V14
);
let grid_start =
u64::from_le_bytes(original[bfoot + 8..bfoot + 16].try_into().unwrap()) as usize;
let grid_end =
u64::from_le_bytes(original[bfoot + 16..bfoot + 24].try_into().unwrap()) as usize;
println!(
"\ncorpus: {DOCS} docs, {TOPICS} topics, {} blocks | grid {:.1} MB | k={K}, {QUERIES} queries, exact",
num_blocks,
(grid_end - grid_start) as f64 / 1e6
);
let patch = |levels: &[u8]| -> Vec<u8> {
let map = |n: u8| -> u8 { *levels.iter().find(|&&l| l >= n).unwrap() };
let mut bytes = original.clone();
for b in &mut bytes[grid_start..grid_end] {
*b = map(*b & 0x0F) | (map(*b >> 4) << 4);
}
bytes
};
type TopK = Vec<Vec<(u32, f32)>>;
let mut baseline: Option<(TopK, TopK)> = None;
for (name, levels) in [
("grid 4-bit (shipped)", (0u8..=15).collect::<Vec<_>>()),
(
"grid 3-bit-ish (9 lvl)",
vec![0, 2, 4, 6, 8, 10, 12, 14, 15],
),
("grid 2-bit (4 lvl)", vec![0, 5, 10, 15]),
] {
dir.write(&sparse_path, &patch(&levels)).await.unwrap();
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let segs = index.segment_readers().await.unwrap();
let bmps: Vec<_> = segs
.iter()
.map(|s| s.bmp_indexes().get(&sparse.0).unwrap().clone())
.collect();
drop(segs);
let (s0, k0, v0) = counters();
let (topk, secs) = run_queries(bmps.clone(), &queries);
let (s1, k1, v1) = counters();
let (bg_topk, _bg_secs) = run_queries(bmps, &bg_queries);
let (s2, k2, _v2) = counters();
match &baseline {
None => baseline = Some((topk, bg_topk)),
Some((base, bg_base_topk)) => {
for (qi, (a, b)) in base.iter().zip(topk.iter()).enumerate() {
assert_eq!(a, b, "q{qi}: grid quantization must not change exact top-k");
}
for (qi, (a, b)) in bg_base_topk.iter().zip(bg_topk.iter()).enumerate() {
assert_eq!(
a, b,
"bg q{qi}: grid quantization must not change exact top-k"
);
}
}
}
println!(
"{name:<24} topical: blocks/q {:>6.1} skip {:>5.1}% sbs/q {:>5.1} {:>5.2} ms/q | background: blocks/q {:>7.1} skip {:>5.1}%",
(s1 - s0) as f64 / QUERIES as f64,
100.0 * (k1 - k0) as f64 / ((s1 - s0) + (k1 - k0)).max(1) as f64,
(v1 - v0) as f64 / QUERIES as f64,
secs * 1000.0 / QUERIES as f64,
(s2 - s1) as f64 / 50.0,
100.0 * (k2 - k1) as f64 / ((s2 - s1) + (k2 - k1)).max(1) as f64,
);
}
dir.write(&sparse_path, &original).await.unwrap();
let dir4 = RamDirectory::new();
build(&dir4, sparse, true, &gen_doc).await;
let index4 = Index::open(dir4.clone(), config.clone()).await.unwrap();
let segs4 = index4.segment_readers().await.unwrap();
let bmps4: Vec<_> = segs4
.iter()
.map(|s| s.bmp_indexes().get(&sparse.0).unwrap().clone())
.collect();
drop(segs4);
let (s0, k0, _v0) = counters();
let (topk4, secs4) = run_queries(bmps4, &queries);
let (s1, k1, _v1) = counters();
let (base, _bg) = baseline.as_ref().unwrap();
let mut overlap = 0usize;
let mut denom = 0usize;
for (a, b) in base.iter().zip(topk4.iter()) {
let set: std::collections::HashSet<u32> = a.iter().map(|&(d, _)| d).collect();
overlap += b.iter().filter(|&&(d, _)| set.contains(&d)).count();
denom += a.len();
}
println!(
"{:<24} blocks scored/q {:>7.1} | skipped/q {:>7.1} | skip {:>5.1}% | recall@{K} vs 8-bit {:>5.1}% | {:>6.2} ms/q",
"weights 4-bit",
(s1 - s0) as f64 / QUERIES as f64,
(k1 - k0) as f64 / QUERIES as f64,
100.0 * (k1 - k0) as f64 / ((s1 - s0) + (k1 - k0)).max(1) as f64,
100.0 * overlap as f64 / denom.max(1) as f64,
secs4 * 1000.0 / QUERIES as f64,
);
}
#[tokio::test]
async fn test_bmp_grid_bits_2_exact_parity_and_roundtrip() {
let build = |grid_bits: u8| {
let mut cfg = bmp_config();
cfg.bmp_grid_bits = grid_bits;
cfg
};
let corpus = |writer: &mut IndexWriter<RamDirectory>, sparse: crate::dsl::Field| {
const DOCS: usize = 600;
for seg in 0..2 {
for i in 0..DOCS {
let global = (seg * DOCS + i) as u32;
let topic = global % 4;
let mut entries: Vec<(u32, f32)> = (0..32)
.map(|t| (50_000 + topic * 100 + t, 0.5 + (t as f32) * 0.01))
.collect();
entries.push((global, 1.0));
let mut doc = Document::new();
doc.add_sparse_vector(sparse, entries);
writer.add_document(doc).unwrap();
}
}
};
let mut results: Vec<Vec<Vec<(u32, f32)>>> = Vec::new();
let mut grid_sizes: Vec<usize> = Vec::new();
for grid_bits in [4u8, 2u8] {
let mut sb = SchemaBuilder::default();
let _t = sb.add_text_field("title", true, true);
let sparse = sb.add_sparse_vector_field_with_config("sparse", true, true, build(grid_bits));
sb.set_reorder(sparse, true);
let schema = sb.build();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
corpus(&mut writer, sparse);
writer.commit().await.unwrap();
writer.force_merge().await.unwrap(); writer.reorder().await.unwrap(); drop(writer);
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let readers = index.segment_readers().await.unwrap();
assert_eq!(readers.len(), 1);
let bmp = readers[0].bmp_indexes().get(&sparse.0).unwrap();
assert_eq!(bmp.grid_bits(), grid_bits, "footer must carry grid_bits");
grid_sizes.push(bmp.packed_row_size());
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
let mut per_query = Vec::new();
for topic in 0..4u32 {
let q: Vec<(u32, f32)> = (0..8).map(|t| (50_000 + topic * 100 + t, 1.0)).collect();
let query = SparseVectorQuery::new(sparse, q);
let r = searcher.search(&query, 10).await.unwrap();
assert_eq!(r.len(), 10);
per_query.push(r.iter().map(|d| (d.doc_id, d.score)).collect::<Vec<_>>());
}
for needle in [0u32, 599, 600, 1199] {
let query = SparseVectorQuery::new(sparse, vec![(needle, 1.0)]);
let r = searcher.search(&query, 5).await.unwrap();
assert_eq!(r.len(), 1, "needle {needle} lost at grid_bits={grid_bits}");
per_query.push(r.iter().map(|d| (d.doc_id, d.score)).collect::<Vec<_>>());
}
results.push(per_query);
}
for (q4, q2) in results[0].iter().zip(results[1].iter()) {
let s4: Vec<f32> = q4.iter().map(|&(_, s)| s).collect();
let s2: Vec<f32> = q2.iter().map(|&(_, s)| s).collect();
assert_eq!(s4.len(), s2.len());
for (a, b) in s4.iter().zip(s2.iter()) {
assert!((a - b).abs() < 1e-3, "score mismatch {a} vs {b}");
}
}
assert_eq!(grid_sizes[1], grid_sizes[0].div_ceil(2));
}
#[tokio::test]
async fn test_blockwise_reorder_preserves_sb_grid_scale() {
let (schema, _title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
add_rare_dim_clustered_corpus(&mut writer, sparse);
writer.commit().await.unwrap();
drop(writer);
let sb_max = |dir: RamDirectory, config: IndexConfig| async move {
let index = Index::open(dir, config).await.unwrap();
let readers = index.segment_readers().await.unwrap();
let bmp = readers[0].bmp_indexes().get(&sparse.0).unwrap();
bmp.sb_grid_slice().iter().copied().max().unwrap_or(0)
};
let before = sb_max(dir.clone(), config.clone()).await;
assert!(
before > 15,
"test setup: builder sb_grid must carry u8 impacts (got max {before})"
);
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
writer.reorder().await.unwrap();
drop(writer);
let after = sb_max(dir.clone(), config.clone()).await;
assert_eq!(
after, before,
"blockwise reorder must preserve sb_grid impact scale (0-255) — \
cell-scale values (≤15) deflate superblock UBs ~17× (unsafe pruning)"
);
}
#[tokio::test]
#[ignore]
async fn bench_forward_index_build() {
use crate::segment::{
build_forward_index_from_blocks, build_forward_index_from_bmps, graph_bisection,
};
struct XorShift(u64);
impl XorShift {
fn next(&mut self) -> u64 {
let mut x = self.0;
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
self.0 = x;
x
}
}
const DOCS: usize = 300_000;
const DIMS: u32 = 30_000;
const TERMS_PER_DOC: usize = 48;
let config_sparse = SparseVectorConfig {
format: SparseFormat::Bmp,
weight_quantization: WeightQuantization::UInt8,
bmp_block_size: 64,
dims: Some(DIMS),
max_weight: Some(5.0),
..SparseVectorConfig::default()
};
let (schema, _title, sparse) = bmp_schema_with_config(config_sparse);
let dir = RamDirectory::new();
let config = IndexConfig {
merge_policy: Box::new(crate::merge::NoMergePolicy),
num_indexing_threads: 1,
max_indexing_memory_bytes: 8 * 1024 * 1024 * 1024,
..IndexConfig::default()
};
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
let mut rng = XorShift(0xC0FFEE1234567890);
for _ in 0..DOCS {
let mut entries: Vec<(u32, f32)> = (0..TERMS_PER_DOC)
.map(|t| {
let d = if t % 2 == 0 {
(rng.next() % 2_000) as u32
} else {
2_000 + (rng.next() % (DIMS as u64 - 2_000)) as u32
};
(d, 0.5 + (rng.next() % 100) as f32 / 25.0)
})
.collect();
entries.sort_by_key(|&(d, _)| d);
entries.dedup_by_key(|&mut (d, _)| d);
loop {
let mut doc = Document::new();
doc.add_sparse_vector(sparse, entries.clone());
match writer.add_document(doc) {
Ok(()) => break,
Err(crate::Error::QueueFull) => {
tokio::time::sleep(std::time::Duration::from_millis(3)).await;
}
Err(e) => panic!("{e:?}"),
}
}
}
writer.commit().await.unwrap();
drop(writer);
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let readers = index.segment_readers().await.unwrap();
assert_eq!(readers.len(), 1);
let bmp = readers[0].bmp_indexes().get(&sparse.0).unwrap().clone();
println!(
"corpus: {} docs, {} blocks, {} postings",
DOCS,
bmp.num_blocks,
bmp.total_postings()
);
let min_df = (DOCS / 5000).clamp(2, 128);
let max_df = ((DOCS as f64) * 0.9) as usize;
let budget = 2usize * 1024 * 1024 * 1024;
for round in 0..3 {
let t = std::time::Instant::now();
let (fwd, _) = build_forward_index_from_bmps(&[&bmp], min_df, max_df, budget).unwrap();
println!(
"record-level fwd build round {}: {:.1}ms ({} terms, {} postings)",
round,
t.elapsed().as_secs_f64() * 1000.0,
fwd.num_terms,
fwd.total_postings(),
);
}
for round in 0..3 {
let t = std::time::Instant::now();
let fwd = build_forward_index_from_blocks(&[&bmp], budget);
println!(
"block-level fwd build round {}: {:.1}ms ({} terms, {} postings)",
round,
t.elapsed().as_secs_f64() * 1000.0,
fwd.num_terms,
fwd.total_postings(),
);
}
let (fwd, _) = build_forward_index_from_bmps(&[&bmp], min_df, max_df, budget).unwrap();
let t = std::time::Instant::now();
let (_perm, converged) = graph_bisection(&fwd, 64, 20, crate::segment::BpBudget::full());
println!(
"graph_bisection (record-level, full depth): {:.1}ms (converged={})",
t.elapsed().as_secs_f64() * 1000.0,
converged,
);
drop(fwd);
drop(bmp);
drop(readers);
drop(index);
let mut writer = IndexWriter::open(dir.clone(), config.clone())
.await
.unwrap();
let t = std::time::Instant::now();
writer.reorder().await.unwrap();
println!(
"writer.reorder() end-to-end: {:.1}ms",
t.elapsed().as_secs_f64() * 1000.0,
);
}
#[tokio::test]
async fn test_bmp_block_posting_overflow_256() {
let config = SparseVectorConfig {
format: SparseFormat::Bmp,
weight_quantization: WeightQuantization::UInt8,
bmp_block_size: 256,
dims: Some(100_000),
max_weight: Some(5.0),
..SparseVectorConfig::default()
};
let (schema, _title, sparse) = bmp_schema_with_config(config);
let dir = RamDirectory::new();
let config = IndexConfig::default();
let mut writer = IndexWriter::create(dir.clone(), schema, config.clone())
.await
.unwrap();
for d in 0..256u32 {
let mut entries: Vec<(u32, f32)> = (0..300u32)
.map(|t| (t, 0.5 + (t % 40) as f32 * 0.1))
.collect();
entries.push((50_000 + d, 4.0)); let mut doc = Document::new();
doc.add_sparse_vector(sparse, entries);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
drop(writer);
let index = Index::open(dir.clone(), config.clone()).await.unwrap();
let readers = index.segment_readers().await.unwrap();
assert_eq!(readers.len(), 1);
let bmp = readers[0].bmp_indexes().get(&sparse.0).unwrap();
assert!(
bmp.total_postings() > u16::MAX as u64,
"test setup must overflow u16 postings per block (got {})",
bmp.total_postings()
);
let reader = index.reader().await.unwrap();
let searcher = reader.searcher().await.unwrap();
for d in [0u32, 100, 255] {
let query = SparseVectorQuery::new(sparse, vec![(50_000 + d, 1.0)]);
let r = searcher.search(&query, 5).await.unwrap();
assert_eq!(r.len(), 1, "needle {d} lost past the 65,535th posting");
assert_eq!(r[0].doc_id, d);
}
let query = SparseVectorQuery::new(sparse, vec![(0, 1.0), (299, 1.0)]);
let r = searcher.search(&query, 10).await.unwrap();
assert_eq!(r.len(), 10);
let (fwd, _) =
crate::segment::build_forward_index_from_bmps(&[bmp], 2, 1_000_000, 2 * 1024 * 1024 * 1024)
.unwrap();
assert_eq!(fwd.total_postings(), 300 * 256);
}
#[tokio::test]
async fn test_reorder_skips_segment_consumed_by_merge() {
let (schema, _title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig {
merge_policy: Box::new(crate::merge::NoMergePolicy),
..IndexConfig::default()
};
let index = Index::create(dir.clone(), schema, config.clone())
.await
.unwrap();
let mut writer = index.writer();
for seg in 0..2 {
for i in 0..300u32 {
let mut doc = Document::new();
doc.add_sparse_vector(sparse, vec![(seg * 1000 + i, 1.0), (50_000 + i % 7, 0.5)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
}
let segs = index.segment_readers().await.unwrap();
assert_eq!(segs.len(), 2);
let victim = crate::segment::SegmentId(segs[0].meta().id).to_hex();
let total_docs: u32 = segs.iter().map(|s| s.num_docs()).sum();
drop(segs);
let reader = index.reader().await.unwrap();
let _searcher = reader.searcher().await.unwrap();
writer.force_merge().await.unwrap();
let swept = index
.segment_manager()
.cleanup_orphan_segments()
.await
.unwrap();
assert_eq!(swept, 0, "snapshot-deferred sources are not orphans");
use crate::directories::Directory;
let files = dir.list_files(std::path::Path::new("")).await.unwrap();
assert!(
files
.iter()
.any(|path| path.to_string_lossy().contains(&victim)),
"snapshot must keep the retired source files alive"
);
let reordered = index
.segment_manager()
.reorder_single_segment(&victim, None, crate::segment::BpBudget::full())
.await
.unwrap();
assert!(
!reordered,
"stale reorder of a merged-away segment must skip"
);
let index2 = Index::open(dir.clone(), config.clone()).await.unwrap();
let segs2 = index2.segment_readers().await.unwrap();
let total_after: u32 = segs2.iter().map(|s| s.num_docs()).sum();
assert_eq!(
total_after, total_docs,
"reordering a merged-away segment duplicated its documents"
);
}
#[tokio::test]
async fn test_failed_reorder_leaves_no_orphan_files() {
use crate::directories::Directory;
let (schema, _title, sparse) = bmp_schema();
let dir = RamDirectory::new();
let config = IndexConfig {
merge_policy: Box::new(crate::merge::NoMergePolicy),
..IndexConfig::default()
};
let index = Index::create(dir.clone(), schema, config.clone())
.await
.unwrap();
let mut writer = index.writer();
for i in 0..200u32 {
let mut doc = Document::new();
doc.add_sparse_vector(sparse, vec![(i, 1.0), (50_000 + i % 5, 0.5)]);
writer.add_document(doc).unwrap();
}
writer.commit().await.unwrap();
async fn seg_files(dir: &RamDirectory) -> Vec<String> {
let mut ids: Vec<String> = dir
.list_files(std::path::Path::new(""))
.await
.unwrap()
.into_iter()
.filter_map(|p| {
let n = p.to_string_lossy().to_string();
n.starts_with("seg_").then(|| n[4..36].to_string())
})
.collect();
ids.sort();
ids.dedup();
ids
}
let victim = {
let segs = index.segment_readers().await.unwrap();
crate::segment::SegmentId(segs[0].meta().id).to_hex()
};
crate::segment::delete_segment(&dir, crate::segment::SegmentId::from_hex(&victim).unwrap())
.await
.unwrap();
let before = seg_files(&dir).await;
let err = index
.segment_manager()
.reorder_single_segment(&victim, None, crate::segment::BpBudget::full())
.await;
assert!(err.is_err(), "reorder of a fileless segment must fail");
assert_eq!(
seg_files(&dir).await,
before,
"failed reorder must delete its partial output files"
);
let orphan_hex = "00000000000000000000000000000abc";
let orphan_id = crate::segment::SegmentId::from_hex(orphan_hex).unwrap();
use crate::directories::DirectoryWriter;
dir.write(
std::path::Path::new(&format!("seg_{orphan_hex}.store")),
b"junk",
)
.await
.unwrap();
dir.write(
std::path::Path::new(&format!("seg_{orphan_hex}.sparse")),
b"junk",
)
.await
.unwrap();
let tracker = index.segment_manager().tracker();
tracker.register(orphan_hex);
let scheduled = tracker.mark_for_deletion(&[orphan_hex.to_string()]);
assert_eq!(scheduled, vec![orphan_id]);
assert!(tracker.is_deletion_protected(orphan_hex));
let swept = index
.segment_manager()
.cleanup_orphan_segments()
.await
.unwrap();
assert_eq!(
swept, 0,
"orphan sweep must not race an already scheduled deletion"
);
assert!(
seg_files(&dir).await.contains(&orphan_hex.to_string()),
"scheduled files must remain untouched by the sweeper"
);
tracker.complete_deletion(&scheduled);
let swept = index
.segment_manager()
.cleanup_orphan_segments()
.await
.unwrap();
assert!(swept >= 1, "orphan sweep must delete stray segment files");
let after = seg_files(&dir).await;
assert!(
!after.contains(&orphan_hex.to_string()),
"orphan files must be gone"
);
}