diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a9125b1..5e3a35a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -196,16 +196,10 @@ jobs: key: short-fuzz-${{ matrix.target }} - name: 安装 cargo-fuzz - uses: taiki-e/install-action@v2 - with: - tool: cargo-fuzz - continue-on-error: true - - - name: Fallback 安装 cargo-fuzz - run: | - if ! cargo fuzz --version 2>/dev/null; then - cargo install cargo-fuzz --locked - fi + # install-action 会回退到 musl 预编译包,使 cargo-fuzz 错把 CI 主机识别为 + # x86_64-unknown-linux-musl;ASan 与静态 musl 不兼容,因此从源码为实际 + # x86_64-unknown-linux-gnu 主机安装。 + run: cargo install cargo-fuzz --locked - name: 运行短时模糊测试 (${{ matrix.target }}) run: | @@ -293,8 +287,9 @@ jobs: # (重型 benchmark 会导致 CI 超时且不贡献库覆盖率) # --exclude triviumdb-cli: CLI/TUI 工具是独立交付物,不计入核心库 80% 门禁 cargo llvm-cov --workspace --exclude triviumdb-cli --lib --tests --lcov --output-path lcov.info --fail-under-lines 80 -- --test-threads=1 - cargo llvm-cov report --workspace --exclude triviumdb-cli --html --output-dir target/llvm-cov/html - cargo llvm-cov report --workspace --exclude triviumdb-cli + # report 子命令复用上一步已收集的覆盖率数据,不接受 workspace package 选择参数。 + cargo llvm-cov report --html --output-dir target/llvm-cov/html + cargo llvm-cov report - name: 上传覆盖率报告 uses: actions/upload-artifact@v4 diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml index 00c8945..7ed2aca 100644 --- a/fuzz/Cargo.toml +++ b/fuzz/Cargo.toml @@ -31,3 +31,6 @@ doc = false name = "fuzz_filter_parse" path = "fuzz_targets/fuzz_filter_parse.rs" doc = false + +# cargo-fuzz 作为独立 workspace 构建,避免被父级 workspace 发现后拒绝运行。 +[workspace] diff --git a/src/database/pipeline.rs b/src/database/pipeline.rs index 96481d4..c9e32d4 100644 --- a/src/database/pipeline.rs +++ b/src/database/pipeline.rs @@ -127,7 +127,7 @@ pub(crate) fn execute_pipeline( config.force_brute_force || (config.enable_advanced_pipeline && config.enable_sparse_residual); mt.ensure_vectors_cache(need_flat); recall_text(&mt, config, query_text, &mut seed_map); - recall_vector(&mt, config, query_vector, &mut seed_map); + recall_vector(&mut mt, config, query_vector, &mut seed_map); recall_residual(&mt, config, query_vector, &mut seed_map); } @@ -257,7 +257,7 @@ fn recall_text( /// L2 + L3: 向量稠密召回(自适应路由 + 布隆预过滤) fn recall_vector( - mt: &MemTable, + mt: &mut MemTable, config: &SearchConfig, query_vector: Option<&[T]>, seed_map: &mut std::collections::HashMap, @@ -269,15 +269,6 @@ fn recall_vector( let dim = mt.dim(); - // 构建 payload 过滤闭包 - let filter_ref = config.payload_filter.as_ref(); - let passes_filter = |id: NodeId| -> bool { - match filter_ref { - None => true, - Some(f) => mt.get_payload(id).is_some_and(|p| f.matches(p)), - } - }; - // ═══════════════════════════════════════════════════════ // 动态引擎路由: // 1. QuIVer Vamana 图搜索(N >= 10,000 时由 ensure_vectors_cache 自动构建) @@ -286,11 +277,26 @@ fn recall_vector( // —— 需要连续 flat 数组,由 ensure_vectors_cache 在该路径下构建 merged 缓存 // ═══════════════════════════════════════════════════════ let vector_hits: Vec = if !config.force_brute_force && mt.quiver().is_some() { - quiver_pipeline(mt, config, query_vector, &passes_filter) + let approximate_hits = quiver_pipeline(mt, config, query_vector); + if config.payload_filter.is_some() && approximate_hits.len() < config.top_k { + // QuIVer 的 BQ beam 候选池本身不感知 Payload。高选择性过滤可能令 + // `ef_search` 池中不足 top_k,即使池外仍有匹配节点。仅在确实欠填时 + // 物化连续缓存并回退精确扫描,兼顾常规冷路径与过滤结果完整性。 + tracing::debug!( + returned = approximate_hits.len(), + requested = config.top_k, + "QuIVer 过滤后结果不足,回退精确暴力扫描 (filtered QuIVer underfilled; falling back to exact scan)" + ); + mt.ensure_vectors_cache(true); + let vectors = mt.flat_vectors(); + brute_force_pipeline(mt, config, query_vector, vectors, dim) + } else { + approximate_hits + } } else { // ensure_vectors_cache() 已在 execute_pipeline 中按需构建好 merged 缓存 let vectors = mt.flat_vectors(); - brute_force_pipeline(mt, config, query_vector, vectors, dim, &passes_filter) + brute_force_pipeline(mt, config, query_vector, vectors, dim) }; for hit in vector_hits { @@ -305,30 +311,19 @@ fn brute_force_pipeline( query_vector: &[T], vectors: &[T], dim: usize, - passes_filter: &(dyn Fn(NodeId) -> bool + Sync), ) -> Vec { let bloom_mask = config .payload_filter .as_ref() .map(|f| f.extract_must_have_mask()) .unwrap_or(0); - let fast_tags = mt.fast_tags_slice(); - brute_force::search( + brute_force::search_filter_map( query_vector, vectors, dim, config.top_k, config.min_score, - |idx| { - let id = mt.get_id_by_index(idx); - if bloom_mask != 0 - && idx < fast_tags.len() - && (fast_tags[idx] & bloom_mask) != bloom_mask - { - return 0; // True Negative - } - if passes_filter(id) { id } else { 0 } - }, + |idx| eligible_node_id(mt, config.payload_filter.as_ref(), bloom_mask, idx), ) } @@ -342,7 +337,6 @@ fn quiver_pipeline( mt: &MemTable, config: &SearchConfig, query_vector: &[T], - passes_filter: &(dyn Fn(NodeId) -> bool + Sync), ) -> Vec { let quiver = mt.quiver().unwrap(); let q_f32: Vec = query_vector.iter().map(|x| x.to_f32()).collect(); @@ -355,19 +349,31 @@ fn quiver_pipeline( rerank_limit: None, }; + let filter_ref = config.payload_filter.as_ref(); + let bloom_mask = filter_ref + .map(|filter| filter.extract_must_have_mask()) + .unwrap_or(0); + // 冷热分离精排回调:按 MemTable slot 索引**按需**从 mmap 零拷贝取回单条向量, // 转为 f32 写入复用缓冲区。整个查询只触达 ~ef 个候选, // 不再物化全量 f32 数组(不触发 merged 缓存),冷数据始终留在 OS PageCache。 let vec_pool = mt.vec_pool(); let raw_results = quiver.search( &q_f32, - |slot, buf| match vec_pool.get(slot) { - Some(v) => { - buf.clear(); - buf.extend(v.iter().map(|x| x.to_f32())); - true + |slot, buf| { + // QuIVer 仍可穿过不匹配节点进行图导航,但只有满足过滤条件的候选 + // 才进入 f32 精排与内部 top_k,避免固定 2× 过召回被无效候选占满。 + if eligible_node_id(mt, filter_ref, bloom_mask, slot).is_none() { + return false; + } + match vec_pool.get(slot) { + Some(v) => { + buf.clear(); + buf.extend(v.iter().map(|x| x.to_f32())); + true + } + None => false, } - None => false, }, &search_cfg, ); @@ -375,7 +381,9 @@ fn quiver_pipeline( // 应用 payload 过滤 + min_score 阈值 let mut hits: Vec = raw_results .into_iter() - .filter(|&(id, score)| score >= config.min_score && passes_filter(id)) + .filter(|&(id, score)| { + score >= config.min_score && matches_payload_filter(mt, filter_ref, id) + }) .map(|(id, score)| SearchHit { id, score, @@ -410,13 +418,24 @@ fn recall_residual( None => return, }; + let filter_ref = config.payload_filter.as_ref(); + let bloom_mask = filter_ref + .map(|filter| filter.extract_must_have_mask()) + .unwrap_or(0); + let entity_vecs: Vec> = seed_map .keys() .filter_map(|&id| { + if !matches_payload_filter(mt, filter_ref, id) { + return None; + } mt.get_vector(id) .map(|v| v.iter().map(|&x| x.to_f32()).collect()) }) .collect(); + if entity_vecs.is_empty() { + return; + } let q_f32: Vec = query_vector.iter().map(|&x| x.to_f32()).collect(); let (_, residual, residual_norm) = @@ -431,13 +450,13 @@ fn recall_residual( ); let r_orig: Vec = residual.iter().map(|&x| T::from_f32(x)).collect(); let dim = mt.dim(); - let shadow_hits = brute_force::search( + let shadow_hits = brute_force::search_filter_map( &r_orig, mt.flat_vectors(), dim, config.top_k, config.min_score, - |idx| mt.get_id_by_index(idx), + |idx| eligible_node_id(mt, filter_ref, bloom_mask, idx), ); for sh in shadow_hits { *seed_map.entry(sh.id).or_insert(0.0) += sh.score * 0.8; // 影子抑制衰减 @@ -454,18 +473,12 @@ fn aggregate_seeds( ) { let filter_ref = config.payload_filter.as_ref(); for (&id, &score) in seed_map { - if score >= config.min_score { - let passes = match filter_ref { - None => mt.contains(id), - Some(f) => mt.get_payload(id).is_some_and(|p| f.matches(p)), - }; - if passes { - let payload = mt - .get_payload(id) - .cloned() - .unwrap_or(serde_json::Value::Null); - anchor_hits.push(SearchHit { id, score, payload }); - } + if score >= config.min_score && matches_payload_filter(mt, filter_ref, id) { + let payload = mt + .get_payload(id) + .cloned() + .unwrap_or(serde_json::Value::Null); + anchor_hits.push(SearchHit { id, score, payload }); } } anchor_hits.sort_by(|a, b| { @@ -476,6 +489,56 @@ fn aggregate_seeds( anchor_hits.truncate(config.top_k.max(15)); } +/// 检查一个逻辑节点是否仍然活跃并满足精确 Payload 条件。 +#[inline] +fn matches_payload_filter( + mt: &MemTable, + filter: Option<&crate::filter::Filter>, + id: NodeId, +) -> bool { + if id == 0 { + return false; + } + match filter { + None => mt.contains(id), + Some(filter) => mt + .get_payload(id) + .is_some_and(|payload| filter.matches(payload)), + } +} + +/// 将物理向量槽位映射为可参与召回的逻辑节点。 +/// +/// 顺序固定为:墓碑检查 → Bloom true-negative → 精确 Payload 判断。 +/// 返回 `None` 的槽位不会计算相似度,也不会占用任何 `top_k` 名额。 +#[inline] +fn eligible_node_id( + mt: &MemTable, + filter: Option<&crate::filter::Filter>, + bloom_mask: u64, + idx: usize, +) -> Option { + let id = *mt.internal_indices().get(idx)?; + if id == 0 { + return None; + } + if bloom_mask != 0 + && mt + .fast_tags_slice() + .get(idx) + .is_some_and(|tag| (*tag & bloom_mask) != bloom_mask) + { + return None; + } + match filter { + None => Some(id), + Some(filter) => mt + .get_payload(id) + .filter(|payload| filter.matches(payload)) + .map(|_| id), + } +} + /// L9: DPP 多样性采样 fn apply_dpp( mt: &MemTable, @@ -539,6 +602,36 @@ mod tests { mt } + fn filter_regression_nodes() -> Vec<(u64, Vec, serde_json::Value)> { + vec![ + ( + 10, + vec![1.0, 0.0], + serde_json::json!({"tenant": "drop", "rank": 0}), + ), + ( + 11, + vec![0.8, 0.6], + serde_json::json!({"tenant": "drop", "rank": 0}), + ), + ( + 12, + vec![0.6, 0.8], + serde_json::json!({"tenant": "keep", "rank": 1}), + ), + ( + 13, + vec![0.0, 1.0], + serde_json::json!({"tenant": "keep", "rank": 2}), + ), + ( + 14, + vec![-0.6, 0.8], + serde_json::json!({"tenant": "keep", "rank": 3}), + ), + ] + } + fn wrap(mt: MemTable) -> Arc>> { Arc::new(Mutex::new(mt)) } @@ -668,7 +761,7 @@ mod tests { let query: Vec = vec![1.0, 0.0, 0.0]; let mut seed_map = std::collections::HashMap::new(); - recall_vector(&mt, &cfg, Some(&query), &mut seed_map); + recall_vector(&mut mt, &cfg, Some(&query), &mut seed_map); assert!(!seed_map.is_empty(), "应召回至少一个节点"); // 节点 1 与 query 完全对齐,得分最高 @@ -686,7 +779,7 @@ mod tests { mt.ensure_vectors_cache(true); let cfg = default_config(); let mut seed_map = std::collections::HashMap::new(); - recall_vector(&mt, &cfg, None, &mut seed_map); + recall_vector(&mut mt, &cfg, None, &mut seed_map); assert!(seed_map.is_empty()); } @@ -709,7 +802,7 @@ mod tests { }; let query = vec![1.0, 0.0, 0.0]; let mut seed_map = std::collections::HashMap::new(); - recall_vector(&mt, &cfg, Some(&query), &mut seed_map); + recall_vector(&mut mt, &cfg, Some(&query), &mut seed_map); assert!(seed_map.contains_key(&1)); assert!( @@ -718,6 +811,152 @@ mod tests { ); } + #[test] + fn test_recall_vector_filter_before_top_k() { + let mut mt = make_memtable(2, &filter_regression_nodes()); + mt.ensure_vectors_cache(true); + + let cfg = SearchConfig { + top_k: 2, + min_score: -1.0, + force_brute_force: true, + payload_filter: Some(Filter::eq("tenant", serde_json::json!("keep"))), + ..Default::default() + }; + let mut seed_map = std::collections::HashMap::new(); + recall_vector(&mut mt, &cfg, Some(&[1.0, 0.0]), &mut seed_map); + + assert_eq!(seed_map.len(), 2); + assert!(seed_map.contains_key(&12)); + assert!(seed_map.contains_key(&13)); + assert!(!seed_map.contains_key(&0)); + assert!(!seed_map.contains_key(&10)); + assert!(!seed_map.contains_key(&11)); + } + + #[test] + fn test_recall_vector_non_bloom_filter_before_top_k() { + let mut mt = make_memtable(2, &filter_regression_nodes()); + mt.ensure_vectors_cache(true); + + let cfg = SearchConfig { + top_k: 2, + min_score: -1.0, + force_brute_force: true, + payload_filter: Some(Filter::gt("rank", 0.0)), + ..Default::default() + }; + let mut seed_map = std::collections::HashMap::new(); + recall_vector(&mut mt, &cfg, Some(&[1.0, 0.0]), &mut seed_map); + + assert_eq!(seed_map.len(), 2); + assert!(seed_map.contains_key(&12)); + assert!(seed_map.contains_key(&13)); + } + + #[test] + fn test_recall_vector_tombstone_does_not_consume_top_k() { + let mut mt = make_memtable( + 2, + &[ + (1, vec![1.0, 0.0], serde_json::json!({})), + (2, vec![-1.0, 0.0], serde_json::json!({})), + (3, vec![-0.8, 0.6], serde_json::json!({})), + ], + ); + mt.delete(1).unwrap(); + mt.ensure_vectors_cache(true); + + let cfg = SearchConfig { + top_k: 1, + min_score: -1.0, + force_brute_force: true, + ..Default::default() + }; + let mut seed_map = std::collections::HashMap::new(); + recall_vector(&mut mt, &cfg, Some(&[1.0, 0.0]), &mut seed_map); + + assert_eq!(seed_map.len(), 1); + assert!(seed_map.contains_key(&3)); + assert!(!seed_map.contains_key(&0)); + } + + #[test] + fn test_recall_vector_quiver_filters_before_internal_top_k() { + let mut mt = make_memtable(2, &filter_regression_nodes()); + mt.build_quiver(&crate::index::quiver::QuIVerConfig::default()); + assert!(mt.quiver().is_some()); + + let cfg = SearchConfig { + top_k: 2, + min_score: -1.0, + payload_filter: Some(Filter::eq("tenant", serde_json::json!("keep"))), + ..Default::default() + }; + let mut seed_map = std::collections::HashMap::new(); + recall_vector(&mut mt, &cfg, Some(&[1.0, 0.0]), &mut seed_map); + + assert_eq!(seed_map.len(), 2); + assert!(seed_map.contains_key(&12)); + assert!(seed_map.contains_key(&13)); + } + + #[test] + fn test_recall_vector_quiver_high_selectivity_falls_back_to_exact_scan() { + let mut nodes: Vec<(u64, Vec, serde_json::Value)> = (1..=40) + .map(|id| (id, vec![1.0, 0.0], serde_json::json!({"tenant": "drop"}))) + .collect(); + nodes.extend([ + (100, vec![0.0, 1.0], serde_json::json!({"tenant": "keep"})), + (101, vec![-0.8, 0.6], serde_json::json!({"tenant": "keep"})), + (102, vec![-1.0, 0.0], serde_json::json!({"tenant": "keep"})), + ]); + let mut mt = make_memtable(2, &nodes); + mt.build_quiver(&crate::index::quiver::QuIVerConfig::default()); + + let cfg = SearchConfig { + top_k: 2, + min_score: -1.0, + payload_filter: Some(Filter::eq("tenant", serde_json::json!("keep"))), + ..Default::default() + }; + let approximate_hits = quiver_pipeline(&mt, &cfg, &[1.0, 0.0]); + assert!( + approximate_hits.len() < cfg.top_k, + "该数据集必须先复现 QuIVer 未过滤候选池欠填" + ); + let mut seed_map = std::collections::HashMap::new(); + recall_vector(&mut mt, &cfg, Some(&[1.0, 0.0]), &mut seed_map); + + assert_eq!(seed_map.len(), 2); + assert!(seed_map.contains_key(&100)); + assert!(seed_map.contains_key(&101)); + } + + #[test] + fn test_recall_vector_bloom_preserves_signed_zero_equality() { + let mut mt = make_memtable(2, &[(1, vec![1.0, 0.0], serde_json::json!({"zero": -0.0}))]); + mt.ensure_vectors_cache(true); + + let filter = Filter::eq("zero", serde_json::json!(0.0)); + let bloom_mask = filter.extract_must_have_mask(); + assert_ne!(bloom_mask, 0); + assert_eq!(mt.fast_tags_slice()[0] & bloom_mask, bloom_mask); + + let cfg = SearchConfig { + top_k: 1, + min_score: -1.0, + force_brute_force: true, + payload_filter: Some(filter), + ..Default::default() + }; + let mut seed_map = std::collections::HashMap::new(); + recall_vector(&mut mt, &cfg, Some(&[1.0, 0.0]), &mut seed_map); + + assert_eq!(seed_map.len(), 1); + assert!(seed_map.contains_key(&1)); + } + // ════════ recall_text ════════ #[test] @@ -783,6 +1022,40 @@ mod tests { assert!(seed_map.is_empty()); } + #[test] + fn test_recall_residual_shadow_filters_before_top_k() { + let mut mt = make_memtable( + 2, + &[ + (1, vec![0.0, 1.0], serde_json::json!({"tenant": "keep"})), + (2, vec![1.0, 0.0], serde_json::json!({"tenant": "drop"})), + (3, vec![0.8, 0.6], serde_json::json!({"tenant": "drop"})), + (4, vec![0.6, 0.8], serde_json::json!({"tenant": "keep"})), + (5, vec![-0.6, 0.8], serde_json::json!({"tenant": "keep"})), + ], + ); + mt.ensure_vectors_cache(true); + + let cfg = SearchConfig { + top_k: 2, + min_score: -1.0, + enable_advanced_pipeline: true, + enable_sparse_residual: true, + fista_threshold: 0.0, + payload_filter: Some(Filter::eq("tenant", serde_json::json!("keep"))), + ..Default::default() + }; + let mut seed_map = std::collections::HashMap::from([(1, 1.0)]); + + recall_residual(&mt, &cfg, Some(&[1.0, 0.0]), &mut seed_map); + + assert_eq!(seed_map.len(), 2); + assert!(seed_map.contains_key(&1)); + assert!(seed_map.contains_key(&4)); + assert!(!seed_map.contains_key(&2)); + assert!(!seed_map.contains_key(&3)); + } + // ════════ apply_dpp ════════ #[test] @@ -951,6 +1224,31 @@ mod tests { assert_eq!(results[0].id, 1, "最相似节点应排第一"); } + #[test] + fn test_execute_pipeline_payload_filter_fills_top_k() { + let mt = wrap(make_memtable(2, &filter_regression_nodes())); + let hook: Arc = Arc::new(NoopHook); + let cfg = SearchConfig { + top_k: 2, + min_score: -1.0, + expand_depth: 0, + force_brute_force: true, + payload_filter: Some(Filter::eq("tenant", serde_json::json!("keep"))), + ..Default::default() + }; + let mut ctx = HookContext::new(); + + let results = + execute_pipeline(&mt, &hook, None, Some(&[1.0, 0.0]), &cfg, &mut ctx).unwrap(); + + assert_eq!(results.len(), 2); + assert_eq!( + results.iter().map(|hit| hit.id).collect::>(), + vec![12, 13] + ); + assert!(results.iter().all(|hit| hit.payload["tenant"] == "keep")); + } + #[test] fn test_execute_pipeline_respects_top_k() { let nodes: Vec<(u64, Vec, serde_json::Value)> = (1..=10) diff --git a/src/database/transaction.rs b/src/database/transaction.rs index e350f5d..c0262db 100644 --- a/src/database/transaction.rs +++ b/src/database/transaction.rs @@ -9,7 +9,7 @@ use crate::VectorType; use crate::database::Database; use crate::error::Result; use crate::node::NodeId; -use crate::storage::memtable::MemTable; +use crate::storage::memtable::{MemTable, checked_next_node_id}; use crate::storage::wal::WalEntry; use super::lock_or_recover; @@ -35,7 +35,7 @@ pub(crate) fn replay_entry(mt: &mut MemTable, entry: WalEntry< let _ = mt.raw_insert(id, &vector, payload_val); } // 无论是否跳过,都必须推进 next_id 防止后续 insert 复用已物化的 ID - mt.advance_next_id(id + 1); + mt.advance_next_id(id.checked_add(1).unwrap_or(NodeId::MAX)); } WalEntry::Link { src, @@ -335,11 +335,14 @@ impl Database }); } } - pre_assigned_ids.push(Some(sim_next_id)); - pending_ids.insert(sim_next_id); - sim_next_id += 1; + let assigned_id = sim_next_id; + let next_id = checked_next_node_id(assigned_id)?; + pre_assigned_ids.push(Some(assigned_id)); + pending_ids.insert(assigned_id); + sim_next_id = next_id; } TxOp::InsertWithId { id, vector, .. } => { + let next_id = checked_next_node_id(*id)?; if check_exists!(id) { return Err(crate::error::TriviumError::NodeAlreadyExists(*id)); } @@ -360,7 +363,7 @@ impl Database pre_assigned_ids.push(Some(*id)); pending_ids.insert(*id); if *id >= sim_next_id { - sim_next_id = *id + 1; + sim_next_id = next_id; } } TxOp::Link { src, dst, .. } => { diff --git a/src/filter.rs b/src/filter.rs index eda2490..cfbf472 100644 --- a/src/filter.rs +++ b/src/filter.rs @@ -1,5 +1,24 @@ use serde_json::Value; +/// 将可安全写入行级 Bloom 签名的标量规范化为稳定字符串。 +/// +/// `serde_json::Number` 认为浮点 `+0.0` 与 `-0.0` 相等,但它们的 Display +/// 字符串不同;这里统一成 `0.0`,确保“精确相等”必然得到相同 Bloom 位。 +pub(crate) fn bloom_scalar_repr(value: &Value) -> Option { + match value { + Value::String(value) => Some(value.clone()), + Value::Bool(value) => Some(value.to_string()), + Value::Number(value) => { + if value.is_f64() && value.as_f64() == Some(0.0) { + Some("0.0".to_string()) + } else { + Some(value.to_string()) + } + } + Value::Null | Value::Array(_) | Value::Object(_) => None, + } +} + /// 过滤条件表达式 /// 支持: $eq, $ne, $gt, $gte, $lt, $lte, $in, $nin, $startsWith, $contains, /// $exists, $size, $all, $type, $and, $or @@ -119,13 +138,15 @@ impl Filter { pub fn extract_must_have_mask(&self) -> u64 { match self { Filter::Eq(key, val) => { + // 行级签名只为标量叶子写入 `key:value` 位。 + // 数组、对象和 Null 的展开方式无法安全表示“整体相等”,若仍生成查询位 + // 会把真实匹配错误判成 Bloom true-negative,因此必须退化到精确过滤。 + let Some(val_str) = bloom_scalar_repr(val) else { + return 0; + }; let mut hasher = std::collections::hash_map::DefaultHasher::new(); use std::hash::{Hash, Hasher}; // Consistent with how fast_tags hashes values - let val_str = match val { - Value::String(s) => s.clone(), - v => v.to_string(), - }; format!("{}:{}", key, val_str).hash(&mut hasher); 1u64 << (hasher.finish() % 64) } diff --git a/src/index/brute_force.rs b/src/index/brute_force.rs index 1ec7503..960ed54 100644 --- a/src/index/brute_force.rs +++ b/src/index/brute_force.rs @@ -19,7 +19,7 @@ pub fn search( // 使用 rayon 的 par_chunks 将向量池按 dim 分块,多核并行计算每个块的余弦相似度。 // 每个线程独立收集自己的命中结果,最后合并排序。 // 这是纯安全代码,不涉及任何 unsafe 操作。 - let mut hits: Vec = flat_db_vectors + let hits: Vec = flat_db_vectors .par_chunks(dim) .enumerate() .filter_map(|(i, vec_slice)| { @@ -36,6 +36,52 @@ pub fn search( }) .collect(); + sort_and_truncate(hits, top_k) +} + +/// 带候选过滤的并行暴力搜索。 +/// +/// `id_filter_map` 在相似度计算和 `top_k` 截断之前执行: +/// - `Some(id)` 表示该物理槽位是有效候选; +/// - `None` 表示墓碑、Payload 不匹配或其他应跳过的槽位。 +/// +/// 该函数仅供引擎内部检索管线使用。公开的 [`search`] 保持原有签名和调用语义, +/// 避免用 `NodeId == 0` 充当过滤哨兵,也避免破坏外部 Rust 调用者。 +pub(crate) fn search_filter_map( + query: &[T], + flat_db_vectors: &[T], + dim: usize, + top_k: usize, + min_score: f32, + id_filter_map: impl Fn(usize) -> Option + Sync, +) -> Vec { + if flat_db_vectors.is_empty() || dim == 0 { + return Vec::new(); + } + + let hits: Vec = flat_db_vectors + .par_chunks(dim) + .enumerate() + .filter_map(|(i, vec_slice)| { + // 候选资格必须先于相似度计算与 top_k 截断判定,避免无效高分候选占位。 + let id = id_filter_map(i)?; + let score = T::similarity(query, vec_slice); + if score >= min_score { + Some(SearchHit { + id, + score, + payload: serde_json::Value::Null, + }) + } else { + None + } + }) + .collect(); + + sort_and_truncate(hits, top_k) +} + +fn sort_and_truncate(mut hits: Vec, top_k: usize) -> Vec { hits.sort_by(|a, b| { b.score .partial_cmp(&a.score) @@ -44,3 +90,30 @@ pub fn search( hits.truncate(top_k); hits } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn filter_map_excludes_candidates_before_top_k() { + let flat = vec![ + 1.0f32, 0.0, // id 10, score 1.0, rejected + 0.8, 0.6, // id 11, score 0.8, rejected + 0.6, 0.8, // id 12, score 0.6, accepted + 0.0, 1.0, // id 13, score 0.0, accepted + -0.6, 0.8, // id 14, score -0.6, accepted + ]; + + let hits = search_filter_map(&[1.0, 0.0], &flat, 2, 2, -1.0, |idx| match idx { + 0 | 1 => None, + _ => Some(10 + idx as u64), + }); + + assert_eq!(hits.len(), 2); + assert_eq!( + hits.iter().map(|hit| hit.id).collect::>(), + vec![12, 13] + ); + } +} diff --git a/src/storage/file_format.rs b/src/storage/file_format.rs index 49965f4..6cf1d23 100644 --- a/src/storage/file_format.rs +++ b/src/storage/file_format.rs @@ -158,7 +158,7 @@ fn save_mmap(memtable: &mut MemTable, path: &str) -> Result<() .map(|m| m.len()) .unwrap_or(0); let marker_path = flush_ok_path_from_db(path); - let marker_tmp = format!("{}.tmp", &marker_path); + let marker_tmp = format!("{}.tmp", marker_path); { let mut f = File::create(&marker_tmp)?; f.write_all(&tdb_size.to_le_bytes())?; diff --git a/src/storage/memtable.rs b/src/storage/memtable.rs index d98bf0a..9e5ebdd 100644 --- a/src/storage/memtable.rs +++ b/src/storage/memtable.rs @@ -7,6 +7,23 @@ use crate::node::{Edge, NodeId}; use crate::storage::vec_pool::VecPool; use std::collections::{HashMap, HashSet}; +/// 返回可安全用于自增分配器的下一个节点 ID。 +/// +/// `0` 是墓碑哨兵,`NodeId::MAX` 没有可表示的后继值,因此二者都不能作为 +/// 可分配 ID。集中校验可避免 debug panic、release 回绕以及由此产生的真实 ID 0。 +pub(crate) fn checked_next_node_id(id: NodeId) -> Result { + match id { + 0 => Err(TriviumError::InvalidInput( + "节点 ID 0 为墓碑保留值 (Node ID 0 is reserved for tombstones)".into(), + )), + NodeId::MAX => Err(TriviumError::InvalidInput( + "节点 ID 空间已耗尽,不能分配 u64::MAX (Node ID space exhausted; u64::MAX cannot be assigned)" + .into(), + )), + _ => Ok(id + 1), + } +} + /// 计算给定 JSON 对象的行级特征布隆签名(共 64 位) fn calculate_json_signature(value: &serde_json::Value) -> u64 { let mut sig = 0u64; @@ -32,19 +49,13 @@ fn flatten_and_hash_json(prefix: &str, value: &serde_json::Value, sig: &mut u64) flatten_and_hash_json(prefix, v, sig); } } - serde_json::Value::String(s) => { - let mut hasher = std::collections::hash_map::DefaultHasher::new(); - format!("{}:{}", prefix, s).hash(&mut hasher); - *sig |= 1u64 << (hasher.finish() % 64); - } - serde_json::Value::Bool(b) => { - let mut hasher = std::collections::hash_map::DefaultHasher::new(); - format!("{}:{}", prefix, b).hash(&mut hasher); - *sig |= 1u64 << (hasher.finish() % 64); - } - serde_json::Value::Number(n) => { + serde_json::Value::String(_) + | serde_json::Value::Bool(_) + | serde_json::Value::Number(_) => { let mut hasher = std::collections::hash_map::DefaultHasher::new(); - format!("{}:{}", prefix, n).hash(&mut hasher); + let scalar = crate::filter::bloom_scalar_repr(value) + .expect("string/bool/number must have a Bloom representation"); + format!("{}:{}", prefix, scalar).hash(&mut hasher); *sig |= 1u64 << (hasher.finish() % 64); } serde_json::Value::Null => {} @@ -287,7 +298,8 @@ impl MemTable { Self::validate_vector(vector)?; let id = self.next_id; - self.next_id += 1; + let next_id = checked_next_node_id(id)?; + self.next_id = next_id; // 1. 记录向量(优先尝试从空闲槽复活,否则推入尾部增量层) let sig = calculate_json_signature(&payload); @@ -333,13 +345,15 @@ impl MemTable { } /// 使用外部指定的 ID 插入节点(例如从外部知识库导入数据)。 - /// 如果 ID 已存在会返回错误,并且会自动更新内部的 next_id 以免未来冲突。 + /// 合法范围是 `1..NodeId::MAX`:0 保留为墓碑,MAX 没有可表示的后继 ID。 + /// 如果 ID 已存在会返回错误,并自动更新 next_id 以免未来冲突。 pub fn insert_with_id( &mut self, id: NodeId, vector: &[T], payload: serde_json::Value, ) -> Result<()> { + let next_id = checked_next_node_id(id)?; if self.payloads.contains_key(&id) { return Err(TriviumError::NodeAlreadyExists(id)); } @@ -373,7 +387,7 @@ impl MemTable { // 防御性推进分配器指针,避免后续普通 insert 撞车 if id >= self.next_id { - self.next_id = id + 1; + self.next_id = next_id; } // 增量更新 QuIVer 索引(如果已构建且未暂停同步) diff --git a/tests/hook_and_search.rs b/tests/hook_and_search.rs index fa78fe5..3ac173a 100644 --- a/tests/hook_and_search.rs +++ b/tests/hook_and_search.rs @@ -480,6 +480,7 @@ fn COV3_18_payload_filter() { ..Default::default() }; let hits = db.search_advanced(&[10.0, 0.0, 0.0, 0.0], &config).unwrap(); + assert_eq!(hits.len(), 10, "应填满全部 10 个 group=A 的匹配结果"); for h in &hits { let p = db.get_payload(h.id).unwrap(); assert_eq!(p["group"], "A", "payload_filter 应只返回 group=A"); diff --git a/tests/unit/database.rs b/tests/unit/database.rs index fe06937..2d6b1aa 100644 --- a/tests/unit/database.rs +++ b/tests/unit/database.rs @@ -8,7 +8,7 @@ //! find_nodes_by_field, register_property_index use serde_json::json; -use triviumdb::Database; +use triviumdb::{Database, TriviumError}; fn temp_db(name: &str) -> String { let dir = std::env::temp_dir().join(format!("tdb_unit_{}", name)); @@ -72,6 +72,38 @@ fn insert_with_id_和_get() { assert_eq!(node.id, 42); } +#[test] +fn insert_with_id_拒绝墓碑保留ID零() { + let mut db = open_db("insert_zero_id"); + let err = db.insert_with_id(0, &[1.0, 0.0, 0.0], json!({})); + + assert!(matches!(err, Err(TriviumError::InvalidInput(_)))); + assert_eq!(db.node_count(), 0); +} + +#[test] +fn insert_with_id_拒绝无后继的最大ID() { + let mut db = open_db("insert_max_id"); + let err = db.insert_with_id(u64::MAX, &[1.0, 0.0, 0.0], json!({})); + + assert!(matches!(err, Err(TriviumError::InvalidInput(_)))); + assert_eq!(db.node_count(), 0); + assert_eq!(db.insert(&[1.0, 0.0, 0.0], json!({})).unwrap(), 1); +} + +#[test] +fn insert_在最大可用ID后返回空间耗尽且不写入() { + let mut db = open_db("insert_exhausted_id_space"); + db.insert_with_id(u64::MAX - 1, &[1.0, 0.0, 0.0], json!({})) + .unwrap(); + + let err = db.insert(&[0.0, 1.0, 0.0], json!({})); + assert!(matches!(err, Err(TriviumError::InvalidInput(_)))); + assert_eq!(db.node_count(), 1); + assert!(db.contains(u64::MAX - 1)); + assert!(!db.contains(0)); +} + #[test] fn delete_操作() { let mut db = open_db("delete"); @@ -432,6 +464,41 @@ fn tx_insert_with_id() { assert!(db.contains(100)); } +#[test] +fn tx_insert_with_id_拒绝墓碑保留ID零() { + let mut db = open_db("tx_insert_zero_id"); + let mut tx = db.begin_tx(); + tx.insert_with_id(0, &[1.0, 0.0, 0.0], json!({})); + + assert!(matches!(tx.commit(), Err(TriviumError::InvalidInput(_)))); + assert_eq!(db.node_count(), 0); +} + +#[test] +fn tx_insert_with_id_拒绝无后继的最大ID且保持原子性() { + let mut db = open_db("tx_insert_max_id"); + let mut tx = db.begin_tx(); + tx.insert_with_id(u64::MAX, &[1.0, 0.0, 0.0], json!({})); + tx.insert(&[0.0, 1.0, 0.0], json!({})); + + assert!(matches!(tx.commit(), Err(TriviumError::InvalidInput(_)))); + assert_eq!(db.node_count(), 0); + assert_eq!(db.insert(&[1.0, 0.0, 0.0], json!({})).unwrap(), 1); +} + +#[test] +fn tx_最大可用ID后自动分配失败且保持原子性() { + let mut db = open_db("tx_exhausted_id_space"); + let mut tx = db.begin_tx(); + tx.insert_with_id(u64::MAX - 1, &[1.0, 0.0, 0.0], json!({})); + tx.insert(&[0.0, 1.0, 0.0], json!({})); + + assert!(matches!(tx.commit(), Err(TriviumError::InvalidInput(_)))); + assert_eq!(db.node_count(), 0); + assert!(!db.contains(u64::MAX - 1)); + assert!(!db.contains(0)); +} + #[test] fn tx_link_和_unlink() { let mut db = open_db("tx_link"); diff --git a/tests/unit/filter.rs b/tests/unit/filter.rs index 6050a5b..1541499 100644 --- a/tests/unit/filter.rs +++ b/tests/unit/filter.rs @@ -252,6 +252,37 @@ fn bloom_mask_eq产生非零掩码() { assert_eq!(mask.count_ones(), 1); } +#[test] +fn bloom_mask_eq仅为安全标量生成掩码() { + assert_ne!( + Filter::eq("enabled", json!(true)).extract_must_have_mask(), + 0 + ); + assert_ne!(Filter::eq("count", json!(3)).extract_must_have_mask(), 0); + + assert_eq!( + Filter::eq("items", json!([1, 2])).extract_must_have_mask(), + 0 + ); + assert_eq!( + Filter::eq("meta", json!({"kind": "x"})).extract_must_have_mask(), + 0 + ); + assert_eq!( + Filter::eq("missing", json!(null)).extract_must_have_mask(), + 0 + ); +} + +#[test] +fn bloom_mask_浮点正负零使用相同规范表示() { + let positive_zero = Filter::eq("zero", json!(0.0)).extract_must_have_mask(); + let negative_zero = Filter::eq("zero", json!(-0.0)).extract_must_have_mask(); + + assert_ne!(positive_zero, 0); + assert_eq!(positive_zero, negative_zero); +} + #[test] fn bloom_mask_and合并多个掩码() { let f = Filter::and(vec![ diff --git a/tests/unit/memtable.rs b/tests/unit/memtable.rs index 2cb604a..4bdd809 100644 --- a/tests/unit/memtable.rs +++ b/tests/unit/memtable.rs @@ -7,6 +7,7 @@ //! register_node, register_tombstone, advance_next_id, 等全部公开方法 use serde_json::json; +use triviumdb::TriviumError; use triviumdb::storage::memtable::MemTable; const DIM: usize = 3; @@ -93,6 +94,36 @@ fn insert_with_id_基础() { assert!(mt.next_id_value() > 42); } +#[test] +fn insert_with_id_拒绝墓碑保留ID零() { + let mut mt = make_mt(); + let err = mt.insert_with_id(0, &[1.0, 2.0, 3.0], json!({})); + + assert!(matches!(err, Err(TriviumError::InvalidInput(_)))); + assert_eq!(mt.node_count(), 0); +} + +#[test] +fn insert_with_id_拒绝无后继的最大ID() { + let mut mt = make_mt(); + let err = mt.insert_with_id(u64::MAX, &[1.0, 2.0, 3.0], json!({})); + + assert!(matches!(err, Err(TriviumError::InvalidInput(_)))); + assert_eq!(mt.node_count(), 0); + assert_eq!(mt.next_id_value(), 1); + assert_eq!(mt.insert(&[1.0, 2.0, 3.0], json!({})).unwrap(), 1); +} + +#[test] +fn insert_ID空间耗尽时原子失败() { + let mut mt = MemTable::::new_with_next_id(DIM, u64::MAX); + let err = mt.insert(&[1.0, 2.0, 3.0], json!({})); + + assert!(matches!(err, Err(TriviumError::InvalidInput(_)))); + assert_eq!(mt.node_count(), 0); + assert_eq!(mt.next_id_value(), u64::MAX); +} + #[test] fn insert_with_id_重复ID报错() { let mut mt = make_mt(); diff --git a/tests/vector_types.rs b/tests/vector_types.rs index f7824c3..9f003f0 100644 --- a/tests/vector_types.rs +++ b/tests/vector_types.rs @@ -12,7 +12,8 @@ //! - 精度边界与极端值 use half::f16; -use triviumdb::Database; +use triviumdb::database::SearchConfig; +use triviumdb::{Database, Filter}; // ════════════════════════════════════════════════════════════════ // 公共基础设施 @@ -114,6 +115,62 @@ fn F16_搜索_余弦相似度正确排序() { cleanup(&path); } +#[test] +fn F16_Payload过滤在TopK截断前生效() { + let path = tmp_db("f16_filter_before_topk"); + cleanup(&path); + + let mut db = Database::::open(&path, F16_DIM).unwrap(); + db.insert( + &f16_vec(&[1.0, 0.0, 0.0, 0.0]), + serde_json::json!({"tenant": "drop"}), + ) + .unwrap(); + db.insert( + &f16_vec(&[0.8, 0.6, 0.0, 0.0]), + serde_json::json!({"tenant": "drop"}), + ) + .unwrap(); + let first_keep = db + .insert( + &f16_vec(&[0.6, 0.8, 0.0, 0.0]), + serde_json::json!({"tenant": "keep"}), + ) + .unwrap(); + let second_keep = db + .insert( + &f16_vec(&[0.0, 1.0, 0.0, 0.0]), + serde_json::json!({"tenant": "keep"}), + ) + .unwrap(); + db.insert( + &f16_vec(&[-0.6, 0.8, 0.0, 0.0]), + serde_json::json!({"tenant": "keep"}), + ) + .unwrap(); + + let config = SearchConfig { + top_k: 2, + expand_depth: 0, + min_score: -1.0, + force_brute_force: true, + payload_filter: Some(Filter::eq("tenant", serde_json::json!("keep"))), + ..Default::default() + }; + let results = db + .search_advanced(&f16_vec(&[1.0, 0.0, 0.0, 0.0]), &config) + .unwrap(); + + assert_eq!(results.len(), 2); + assert_eq!( + results.iter().map(|hit| hit.id).collect::>(), + vec![first_keep, second_keep] + ); + assert!(results.iter().all(|hit| hit.payload["tenant"] == "keep")); + + cleanup(&path); +} + #[test] fn F16_持久化_flush后重新加载数据完整() { let path = tmp_db("f16_persist"); @@ -332,6 +389,57 @@ fn U64_搜索_汉明相似度正确排序() { cleanup(&path); } +#[test] +fn U64_Payload过滤在TopK截断前生效() { + let path = tmp_db("u64_filter_before_topk"); + cleanup(&path); + + let mut db = Database::::open(&path, U64_DIM).unwrap(); + db.insert(&[u64::MAX, u64::MAX], serde_json::json!({"tenant": "drop"})) + .unwrap(); + db.insert( + &[u64::MAX - 1, u64::MAX], + serde_json::json!({"tenant": "drop"}), + ) + .unwrap(); + let first_keep = db + .insert( + &[u64::MAX - 3, u64::MAX], + serde_json::json!({"tenant": "keep"}), + ) + .unwrap(); + let second_keep = db + .insert( + &[u64::MAX - 7, u64::MAX], + serde_json::json!({"tenant": "keep"}), + ) + .unwrap(); + db.insert( + &[u64::MAX - 15, u64::MAX], + serde_json::json!({"tenant": "keep"}), + ) + .unwrap(); + + let config = SearchConfig { + top_k: 2, + expand_depth: 0, + min_score: 0.0, + force_brute_force: true, + payload_filter: Some(Filter::eq("tenant", serde_json::json!("keep"))), + ..Default::default() + }; + let results = db.search_advanced(&[u64::MAX, u64::MAX], &config).unwrap(); + + assert_eq!(results.len(), 2); + assert_eq!( + results.iter().map(|hit| hit.id).collect::>(), + vec![first_keep, second_keep] + ); + assert!(results.iter().all(|hit| hit.payload["tenant"] == "keep")); + + cleanup(&path); +} + #[test] fn U64_持久化_flush后重新加载() { let path = tmp_db("u64_persist");