xz-embed 0.1.1

文本向量嵌入与向量存储抽象层
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177

use std::collections::HashMap;
use xz_embed::{SearchResult, rrf_fusion};

// ── helpers ──

fn make_search_result(id: &str, score: f32) -> SearchResult {
    SearchResult {
        id: id.to_string(),
        score,
        metadata: HashMap::new(),
        content: None,
        channel: None,
    }
}

// ═══════════════════════════════════════════════════════════════
// rrf_fusion tests
// ═══════════════════════════════════════════════════════════════

/// Documents appearing in both vector and keyword results get a higher
/// fused score than documents appearing in only one channel.
#[test]
fn test_rrf_overlap() {
    let vector_results = vec![make_search_result("doc1", 0.9), make_search_result("doc2", 0.8)];
    let keyword_results = vec![("doc1".to_string(), 0.85), ("doc3".to_string(), 0.75)];

    let results = rrf_fusion(&vector_results, &keyword_results, 60.0);

    // doc1 appears in both channels → should rank first
    assert_eq!(results[0].id, "doc1", "overlap doc should rank first");

    // doc2 appears only in vector, doc3 only in keyword
    assert_eq!(results.len(), 3, "should have 3 unique docs");

    // doc1's fused score > both single-channel scores
    let doc1 = results.iter().find(|r| r.id == "doc1").unwrap();
    let doc2 = results.iter().find(|r| r.id == "doc2").unwrap();
    let doc3 = results.iter().find(|r| r.id == "doc3").unwrap();
    assert!(
        doc1.fused_score > doc2.fused_score,
        "overlap score {} should exceed single-vector score {}",
        doc1.fused_score,
        doc2.fused_score
    );
    assert!(
        doc1.fused_score > doc3.fused_score,
        "overlap score {} should exceed single-keyword score {}",
        doc1.fused_score,
        doc3.fused_score
    );

    // doc1 has non-zero scores from both channels
    assert!(doc1.vector_score > 0.0, "doc1 should have vector_score");
    assert!(doc1.keyword_score > 0.0, "doc1 should have keyword_score");

    // doc2 has vector_score but no keyword_score (0.0)
    assert!(doc2.vector_score > 0.0, "doc2 should have vector_score");
    assert_eq!(doc2.keyword_score, 0.0, "doc2 should have no keyword score");

    // doc3 has keyword_score but no vector_score (0.0)
    assert_eq!(doc3.vector_score, 0.0, "doc3 should have no vector score");
    assert!(doc3.keyword_score > 0.0, "doc3 should have keyword_score");
}

/// Disjoint result sets (no overlapping IDs) should still produce a
/// unified ranking where higher-ranked items in each channel get
/// higher fused scores.
#[test]
fn test_rrf_disjoint() {
    let vector_results = vec![
        make_search_result("v1", 0.9),
        make_search_result("v2", 0.8),
        make_search_result("v3", 0.7),
    ];
    let keyword_results = vec![("k1".to_string(), 0.9), ("k2".to_string(), 0.8)];

    let results = rrf_fusion(&vector_results, &keyword_results, 60.0);

    assert_eq!(results.len(), 5, "all 5 disjoint docs should appear");

    // Results sorted by fused_score descending
    for w in results.windows(2) {
        assert!(
            w[0].fused_score >= w[1].fused_score,
            "results should be sorted descending by fused_score: {} >= {}",
            w[0].fused_score,
            w[1].fused_score
        );
    }

    // Top rank = 1 from vector + rank 1 from keyword (same RRF contribution
    // since both are rank 1 with identical k). However, there is a subtlety:
    // v1's RRF = 1/(60+0+1) ≈ 1/61 ≈ 0.01639, k1's RRF = 1/(60+0+1) ≈ 0.01639.
    // So they tie. The order among ties is unspecified but we can check
    // that the top 2 fused_scores are equal.
    let top = &results[0];
    let second = &results[1];
    assert!(
        (top.fused_score - second.fused_score).abs() < f32::EPSILON,
        "top two disjoint results with same rank should have nearly equal fused_score"
    );

    // v1 should have non-zero vector_score and zero keyword_score
    let v1 = results.iter().find(|r| r.id == "v1").unwrap();
    assert!(v1.vector_score > 0.0, "v1 should have vector_score");
    assert_eq!(v1.keyword_score, 0.0, "v1 should have no keyword score");

    // k1 should have zero vector_score and non-zero keyword_score
    let k1 = results.iter().find(|r| r.id == "k1").unwrap();
    assert_eq!(k1.vector_score, 0.0, "k1 should have no vector score");
    assert!(k1.keyword_score > 0.0, "k1 should have keyword_score");
}

/// When keyword_results is empty, fusion should return only vector
/// results ranked by their RRF contributions (which preserve the
/// original vector ordering).
#[test]
fn test_rrf_single_channel() {
    let vector_results = vec![
        make_search_result("a", 0.9),
        make_search_result("b", 0.7),
        make_search_result("c", 0.5),
    ];
    let keyword_results: Vec<(String, f32)> = vec![];

    let results = rrf_fusion(&vector_results, &keyword_results, 60.0);

    // All vector results should appear, no keyword results
    assert_eq!(results.len(), 3, "should have 3 results from vector channel only");

    // Preserve original order (a highest score → rank 1 → highest fused)
    assert_eq!(results[0].id, "a", "highest vector score should rank first");
    assert_eq!(results[1].id, "b", "middle vector score should rank second");
    assert_eq!(results[2].id, "c", "lowest vector score should rank third");

    // All results have non-zero vector score and zero keyword score
    for r in &results {
        assert!(r.vector_score > 0.0, "{} should have vector_score", r.id);
        assert_eq!(r.keyword_score, 0.0, "{} should have no keyword score", r.id);
    }
}

/// NaN scores should not cause panics or crashes. The function uses
/// `partial_cmp` with `unwrap_or(Ordering::Equal)` which is NaN-safe.
#[test]
fn test_rrf_nan_scores() {
    // NaN in vector results
    let vector_results =
        vec![make_search_result("nan_doc", f32::NAN), make_search_result("normal", 0.8)];
    let keyword_results = vec![("normal".to_string(), 0.7), ("extra".to_string(), 0.5)];

    // Should not panic
    let results = rrf_fusion(&vector_results, &keyword_results, 60.0);

    // Should still produce results
    assert!(!results.is_empty(), "should produce results even with NaN scores");

    // The NaN-scored doc should appear in results
    let nan_doc = results.iter().find(|r| r.id == "nan_doc");
    assert!(nan_doc.is_some(), "NaN-scored doc should be in fusion results");

    // vector_score for nan_doc should be NaN (preserved from input)
    let nd = nan_doc.unwrap();
    assert!(nd.vector_score.is_nan(), "NaN input score should remain NaN in output");
    assert!(nd.keyword_score == 0.0, "NaN doc has no keyword match");

    // Fusion should still work correctly for the non-NaN items
    let normal = results.iter().find(|r| r.id == "normal").unwrap();
    assert!(normal.vector_score > 0.0, "normal doc should have vector_score");
    assert!(normal.keyword_score > 0.0, "normal doc should have keyword_score");

    // NaN in keyword results too
    let kw_nan: Vec<(String, f32)> = vec![("nan_kw".to_string(), f32::NAN)];
    let results2 = rrf_fusion(&[], &kw_nan, 60.0);
    assert_eq!(results2.len(), 1, "NaN keyword result should appear");
    assert!(results2[0].keyword_score.is_nan(), "NaN keyword score preserved");
}