rebuild BM25 from nodes and harden index invariants

Make BM25 explicitly node-only by removing vector indexing, rebuild schema v2 from node state on startup, and fail loudly when forward, reverse, and doc-length state diverge.

Author: xav-db

helix-db/src/helix_engine/bm25/bm25.rs

@@ -1,11 +1,12 @@
 use crate::{
     helix_engine::{
+        bm25::bm25::{BM25, BM25_SCHEMA_VERSION, build_bm25_payload},
         storage_core::HelixGraphStorage,
         types::GraphError,
         vector_core::{vector::HVector, vector_core},
     },
     protocol::value::Value,
-    utils::properties::ImmutablePropertiesMap,
+    utils::{items::Node, properties::ImmutablePropertiesMap},
 };
 use bincode::Options;
 use itertools::Itertools;
@@ -38,10 +39,79 @@ pub fn migrate(storage: &mut HelixGraphStorage) -> Result<(), GraphError> {
 
     verify_vectors_and_repair(storage)?;
     remove_orphaned_vector_edges(storage)?;
+    migrate_bm25(storage)?;
 
     Ok(())
 }
 
+fn migrate_bm25(storage: &mut HelixGraphStorage) -> Result<(), GraphError> {
+    const BATCH_SIZE: usize = 1024;
+
+    let Some(bm25) = storage.bm25.as_ref() else {
+        return Ok(());
+    };
+
+    let current_schema_version = {
+        let txn = storage.graph_env.read_txn()?;
+        bm25.schema_version(&txn)?
+    };
+
+    if current_schema_version == Some(BM25_SCHEMA_VERSION) {
+        return Ok(());
+    }
+
+    {
+        let mut txn = storage.graph_env.write_txn()?;
+        bm25.clear_all(&mut txn)?;
+        txn.commit()?;
+    }
+
+    let read_txn = storage.graph_env.read_txn()?;
+    let mut batch = Vec::with_capacity(BATCH_SIZE);
+
+    for kv in storage.nodes_db.iter(&read_txn)? {
+        let (id, value) = kv?;
+        batch.push((id, value.to_vec()));
+
+        if batch.len() == BATCH_SIZE {
+            rebuild_bm25_batch(storage, bm25, &batch)?;
+            batch.clear();
+        }
+    }
+
+    drop(read_txn);
+
+    if !batch.is_empty() {
+        rebuild_bm25_batch(storage, bm25, &batch)?;
+    }
+
+    let mut txn = storage.graph_env.write_txn()?;
+    bm25.write_schema_version(&mut txn, BM25_SCHEMA_VERSION)?;
+    txn.commit()?;
+
+    Ok(())
+}
+
+fn rebuild_bm25_batch(
+    storage: &HelixGraphStorage,
+    bm25: &crate::helix_engine::bm25::bm25::HBM25Config,
+    batch: &[(u128, Vec<u8>)],
+) -> Result<(), GraphError> {
+    let arena = bumpalo::Bump::new();
+    let mut txn = storage.graph_env.write_txn()?;
+
+    for (id, value) in batch {
+        let node = Node::from_bincode_bytes(*id, value, &arena)?;
+        if let Some(properties) = node.properties.as_ref() {
+            let data = build_bm25_payload(properties, node.label);
+            bm25.insert_doc(&mut txn, *id, &data)?;
+        }
+    }
+
+    txn.commit()?;
+    Ok(())
+}
+
 pub(crate) fn migrate_pre_metadata_to_native_vector_endianness(
     storage: &mut HelixGraphStorage,
 ) -> Result<StorageMetadata, GraphError> {

helix-db/src/helix_engine/bm25/bm25_tests.rs

@@ -9,19 +9,29 @@
 //! - Performance tests for large datasets
 
 use super::{
-    metadata::{StorageMetadata, VectorEndianness, NATIVE_VECTOR_ENDIANNESS},
+    HelixGraphStorage,
+    metadata::{NATIVE_VECTOR_ENDIANNESS, StorageMetadata, VectorEndianness},
     storage_migration::{
         convert_all_vector_properties, convert_old_vector_properties_to_new_format,
         convert_vector_endianness, migrate,
     },
-    HelixGraphStorage,
 };
 use crate::{
     helix_engine::{
-        storage_core::version_info::VersionInfo, traversal_core::config::Config, types::GraphError,
+        bm25::bm25::{
+            BM25, BM25_SCHEMA_VERSION, BM25_SCHEMA_VERSION_KEY, BM25Metadata, METADATA_KEY,
+        },
+        storage_core::version_info::VersionInfo,
+        traversal_core::{
+            config::Config,
+            ops::{g::G, source::add_n::AddNAdapter},
+        },
+        types::GraphError,
     },
     protocol::value::Value,
+    utils::{items::Node, properties::ImmutablePropertiesMap},
 };
+use bumpalo::Bump;
 use std::collections::HashMap;
 use tempfile::TempDir;
 
@@ -169,6 +179,45 @@ fn clear_metadata(storage: &mut HelixGraphStorage) -> Result<(), GraphError> {
     Ok(())
 }
 
+fn add_test_node(
+    storage: &HelixGraphStorage,
+    label: &'static str,
+    properties: &[(&'static str, Value)],
+) -> u128 {
+    let arena = Bump::new();
+    let properties = if properties.is_empty() {
+        None
+    } else {
+        Some(ImmutablePropertiesMap::new(
+            properties.len(),
+            properties.iter().map(|(key, value)| (*key, value.clone())),
+            &arena,
+        ))
+    };
+    let mut txn = storage.graph_env.write_txn().unwrap();
+    let node = G::new_mut(storage, &arena, &mut txn)
+        .add_n(label, properties, None)
+        .collect_to_obj()
+        .unwrap();
+    let node_id = node.id();
+    txn.commit().unwrap();
+    node_id
+}
+
+fn bm25_search_ids(storage: &HelixGraphStorage, query: &str) -> Vec<u128> {
+    let arena = Bump::new();
+    let txn = storage.graph_env.read_txn().unwrap();
+    storage
+        .bm25
+        .as_ref()
+        .unwrap()
+        .search(&txn, query, 10, &arena)
+        .unwrap()
+        .into_iter()
+        .map(|(id, _)| id)
+        .collect()
+}
+
 // ============================================================================
 // Unit Tests: Endianness Conversion
 // ============================================================================
@@ -960,6 +1009,140 @@ fn test_error_handling_graceful_failure() {
     assert_eq!(count, 11); // 10 valid + 1 invalid
 }
 
+#[test]
+fn test_bm25_migration_rerun_is_noop_once_schema_written() {
+    let (mut storage, _temp_dir) = setup_test_storage();
+    let node_id = add_test_node(&storage, "person", &[("name", Value::from("stable_term"))]);
+
+    let before_metadata = {
+        let txn = storage.graph_env.read_txn().unwrap();
+        let bm25 = storage.bm25.as_ref().unwrap();
+        assert_eq!(
+            bm25.schema_version(&txn).unwrap(),
+            Some(BM25_SCHEMA_VERSION)
+        );
+        bm25.metadata_db
+            .get(&txn, METADATA_KEY)
+            .unwrap()
+            .map(|bytes| bytes.to_vec())
+    };
+    let before_results = bm25_search_ids(&storage, "stable_term");
+    assert_eq!(before_results, vec![node_id]);
+
+    migrate(&mut storage).unwrap();
+
+    let after_metadata = {
+        let txn = storage.graph_env.read_txn().unwrap();
+        let bm25 = storage.bm25.as_ref().unwrap();
+        assert_eq!(
+            bm25.schema_version(&txn).unwrap(),
+            Some(BM25_SCHEMA_VERSION)
+        );
+        bm25.metadata_db
+            .get(&txn, METADATA_KEY)
+            .unwrap()
+            .map(|bytes| bytes.to_vec())
+    };
+    let after_results = bm25_search_ids(&storage, "stable_term");
+
+    assert_eq!(after_results, vec![node_id]);
+    assert_eq!(before_results, after_results);
+    assert_eq!(before_metadata, after_metadata);
+}
+
+#[test]
+fn test_bm25_migration_repairs_stale_node_index() {
+    let (mut storage, _temp_dir) = setup_test_storage();
+    let node_id = add_test_node(&storage, "person", &[("name", Value::from("legacyalpha"))]);
+
+    assert_eq!(bm25_search_ids(&storage, "legacyalpha"), vec![node_id]);
+    assert!(bm25_search_ids(&storage, "freshomega").is_empty());
+
+    {
+        let arena = Bump::new();
+        let mut txn = storage.graph_env.write_txn().unwrap();
+        let node_bytes = storage
+            .nodes_db
+            .get(&txn, &node_id)
+            .unwrap()
+            .unwrap()
+            .to_vec();
+        let mut node = Node::from_bincode_bytes(node_id, &node_bytes, &arena).unwrap();
+        node.properties = Some(ImmutablePropertiesMap::new(
+            1,
+            std::iter::once(("name", Value::from("freshomega"))),
+            &arena,
+        ));
+
+        let updated_bytes = node.to_bincode_bytes().unwrap();
+        storage
+            .nodes_db
+            .put(&mut txn, &node_id, &updated_bytes)
+            .unwrap();
+        storage
+            .bm25
+            .as_ref()
+            .unwrap()
+            .metadata_db
+            .put(&mut txn, BM25_SCHEMA_VERSION_KEY, &0u64.to_le_bytes())
+            .unwrap();
+        txn.commit().unwrap();
+    }
+
+    assert_eq!(bm25_search_ids(&storage, "legacyalpha"), vec![node_id]);
+    assert!(bm25_search_ids(&storage, "freshomega").is_empty());
+
+    migrate(&mut storage).unwrap();
+
+    assert!(bm25_search_ids(&storage, "legacyalpha").is_empty());
+    assert_eq!(bm25_search_ids(&storage, "freshomega"), vec![node_id]);
+
+    let txn = storage.graph_env.read_txn().unwrap();
+    assert_eq!(
+        storage.bm25.as_ref().unwrap().schema_version(&txn).unwrap(),
+        Some(BM25_SCHEMA_VERSION)
+    );
+}
+
+#[test]
+fn test_bm25_migration_drops_legacy_direct_docs() {
+    let (mut storage, _temp_dir) = setup_test_storage();
+    let node_id = add_test_node(&storage, "person", &[("name", Value::from("nodeonlyterm"))]);
+
+    {
+        let mut txn = storage.graph_env.write_txn().unwrap();
+        let bm25 = storage.bm25.as_ref().unwrap();
+        bm25.insert_doc(&mut txn, 999u128, "legacyvectorterm")
+            .unwrap();
+        bm25.metadata_db
+            .put(&mut txn, BM25_SCHEMA_VERSION_KEY, &0u64.to_le_bytes())
+            .unwrap();
+        txn.commit().unwrap();
+    }
+
+    assert_eq!(bm25_search_ids(&storage, "legacyvectorterm"), vec![999u128]);
+    assert_eq!(bm25_search_ids(&storage, "nodeonlyterm"), vec![node_id]);
+
+    migrate(&mut storage).unwrap();
+
+    assert!(bm25_search_ids(&storage, "legacyvectorterm").is_empty());
+    assert_eq!(bm25_search_ids(&storage, "nodeonlyterm"), vec![node_id]);
+
+    let txn = storage.graph_env.read_txn().unwrap();
+    let metadata: BM25Metadata = bincode::deserialize(
+        storage
+            .bm25
+            .as_ref()
+            .unwrap()
+            .metadata_db
+            .get(&txn, METADATA_KEY)
+            .unwrap()
+            .unwrap(),
+    )
+    .unwrap();
+    assert_eq!(metadata.total_docs, 1);
+}
+
 // ============================================================================
 // Performance Tests
 // ============================================================================

helix-db/src/helix_engine/storage_core/storage_migration.rs

@@ -9,6 +9,7 @@ use crate::{
         storage_core::HelixGraphStorage,
         traversal_core::{
             ops::{
+                bm25::search_bm25::SearchBM25Adapter,
                 g::G,
                 source::{add_n::AddNAdapter, n_from_id::NFromIdAdapter},
                 util::update::UpdateAdapter,
@@ -90,3 +91,43 @@ fn test_update_node() {
         other => panic!("unexpected traversal value: {other:?}"),
     }
 }
+
+#[test]
+fn test_update_node_without_prior_bm25_doc_becomes_searchable() {
+    let (_temp_dir, storage) = setup_test_db();
+    let arena = Bump::new();
+    let mut txn = storage.graph_env.write_txn().unwrap();
+
+    let node = G::new_mut(&storage, &arena, &mut txn)
+        .add_n("person", None, None)
+        .collect_to_obj()
+        .unwrap();
+    txn.commit().unwrap();
+
+    let arena = Bump::new();
+    let txn = storage.graph_env.read_txn().unwrap();
+    let traversal = G::new(&storage, &txn, &arena)
+        .n_from_id(&node.id())
+        .collect::<Result<Vec<_>, _>>()
+        .unwrap();
+    drop(txn);
+
+    let arena = Bump::new();
+    let mut txn = storage.graph_env.write_txn().unwrap();
+    G::new_mut_from_iter(&storage, &mut txn, traversal.into_iter(), &arena)
+        .update(&[("name", Value::from("bm25_searchable"))])
+        .collect::<Result<Vec<_>, _>>()
+        .unwrap();
+    txn.commit().unwrap();
+
+    let arena = Bump::new();
+    let txn = storage.graph_env.read_txn().unwrap();
+    let results = G::new(&storage, &txn, &arena)
+        .search_bm25("person", "bm25_searchable", 10)
+        .unwrap()
+        .collect::<Result<Vec<_>, _>>()
+        .unwrap();
+
+    assert_eq!(results.len(), 1);
+    assert_eq!(results[0].id(), node.id());
+}