hermes_core/segment/
mod.rs

1#[cfg(feature = "native")]
2mod builder;
3#[cfg(feature = "native")]
4mod merger;
5mod reader;
6mod store;
7#[cfg(feature = "native")]
8mod tracker;
9mod types;
10mod vector_data;
11
12#[cfg(feature = "native")]
13pub use builder::{MemoryBreakdown, SegmentBuilder, SegmentBuilderConfig, SegmentBuilderStats};
14#[cfg(feature = "native")]
15pub use merger::{MergeStats, SegmentMerger, TrainedVectorStructures, delete_segment};
16pub use reader::{AsyncSegmentReader, SegmentReader, SparseIndex, VectorIndex, VectorSearchResult};
17pub use store::*;
18#[cfg(feature = "native")]
19pub use tracker::{SegmentSnapshot, SegmentTracker};
20pub use types::{FieldStats, SegmentFiles, SegmentId, SegmentMeta};
21pub use vector_data::{FlatVectorData, IVFRaBitQIndexData, ScaNNIndexData};
22
23#[cfg(test)]
24#[cfg(feature = "native")]
25mod tests {
26    use super::*;
27    use crate::directories::RamDirectory;
28    use crate::dsl::SchemaBuilder;
29    use std::sync::Arc;
30
31    #[tokio::test]
32    async fn test_async_segment_reader() {
33        let mut schema_builder = SchemaBuilder::default();
34        let title = schema_builder.add_text_field("title", true, true);
35        let schema = Arc::new(schema_builder.build());
36
37        let dir = RamDirectory::new();
38        let segment_id = SegmentId::new();
39
40        // Build segment using sync builder
41        let config = SegmentBuilderConfig::default();
42        let mut builder = SegmentBuilder::new((*schema).clone(), config).unwrap();
43
44        let mut doc = crate::dsl::Document::new();
45        doc.add_text(title, "Hello World");
46        builder.add_document(doc).unwrap();
47
48        let mut doc = crate::dsl::Document::new();
49        doc.add_text(title, "Goodbye World");
50        builder.add_document(doc).unwrap();
51
52        builder.build(&dir, segment_id).await.unwrap();
53
54        // Open with async reader
55        let reader = AsyncSegmentReader::open(&dir, segment_id, schema.clone(), 0, 16)
56            .await
57            .unwrap();
58
59        assert_eq!(reader.num_docs(), 2);
60
61        // Test postings lookup
62        let postings = reader.get_postings(title, b"hello").await.unwrap();
63        assert!(postings.is_some());
64        assert_eq!(postings.unwrap().doc_count(), 1);
65
66        let postings = reader.get_postings(title, b"world").await.unwrap();
67        assert!(postings.is_some());
68        assert_eq!(postings.unwrap().doc_count(), 2);
69
70        // Test document retrieval
71        let doc = reader.doc(0).await.unwrap().unwrap();
72        assert_eq!(doc.get_first(title).unwrap().as_text(), Some("Hello World"));
73    }
74
75    #[tokio::test]
76    async fn test_dense_vector_ordinal_tracking() {
77        use crate::query::MultiValueCombiner;
78
79        let mut schema_builder = SchemaBuilder::default();
80        // Use simple add method - defaults to Flat index
81        let embedding = schema_builder.add_dense_vector_field("embedding", 4, true, true);
82        let schema = Arc::new(schema_builder.build());
83
84        let dir = RamDirectory::new();
85        let segment_id = SegmentId::new();
86
87        let config = SegmentBuilderConfig::default();
88        let mut builder = SegmentBuilder::new((*schema).clone(), config).unwrap();
89
90        // Doc 0: single vector
91        let mut doc = crate::dsl::Document::new();
92        doc.add_dense_vector(embedding, vec![1.0, 0.0, 0.0, 0.0]);
93        builder.add_document(doc).unwrap();
94
95        // Doc 1: multi-valued vectors (2 vectors)
96        let mut doc = crate::dsl::Document::new();
97        doc.add_dense_vector(embedding, vec![0.0, 1.0, 0.0, 0.0]);
98        doc.add_dense_vector(embedding, vec![0.0, 0.0, 1.0, 0.0]);
99        builder.add_document(doc).unwrap();
100
101        // Doc 2: single vector
102        let mut doc = crate::dsl::Document::new();
103        doc.add_dense_vector(embedding, vec![0.0, 0.0, 0.0, 1.0]);
104        builder.add_document(doc).unwrap();
105
106        builder.build(&dir, segment_id).await.unwrap();
107
108        let reader = AsyncSegmentReader::open(&dir, segment_id, schema.clone(), 0, 16)
109            .await
110            .unwrap();
111
112        // Query close to doc 1's first vector
113        let query = vec![0.0, 0.9, 0.1, 0.0];
114        let results = reader
115            .search_dense_vector(embedding, &query, 10, 1, MultiValueCombiner::Max)
116            .unwrap();
117
118        // Doc 1 should be in results with ordinal tracking
119        let doc1_result = results.iter().find(|r| r.doc_id == 1);
120        assert!(doc1_result.is_some(), "Doc 1 should be in results");
121
122        let doc1 = doc1_result.unwrap();
123        // Should have 2 ordinals (0 and 1) for the two vectors
124        assert!(
125            doc1.ordinals.len() <= 2,
126            "Doc 1 should have at most 2 ordinals, got {}",
127            doc1.ordinals.len()
128        );
129
130        // Check ordinals are valid (0 or 1)
131        for (ordinal, _score) in &doc1.ordinals {
132            assert!(*ordinal <= 1, "Ordinal should be 0 or 1, got {}", ordinal);
133        }
134    }
135
136    #[tokio::test]
137    async fn test_sparse_vector_ordinal_tracking() {
138        use crate::query::MultiValueCombiner;
139
140        let mut schema_builder = SchemaBuilder::default();
141        let sparse = schema_builder.add_sparse_vector_field("sparse", true, true);
142        let schema = Arc::new(schema_builder.build());
143
144        let dir = RamDirectory::new();
145        let segment_id = SegmentId::new();
146
147        let config = SegmentBuilderConfig::default();
148        let mut builder = SegmentBuilder::new((*schema).clone(), config).unwrap();
149
150        // Doc 0: single sparse vector
151        let mut doc = crate::dsl::Document::new();
152        doc.add_sparse_vector(sparse, vec![(0, 1.0), (1, 0.5)]);
153        builder.add_document(doc).unwrap();
154
155        // Doc 1: multi-valued sparse vectors (2 vectors)
156        let mut doc = crate::dsl::Document::new();
157        doc.add_sparse_vector(sparse, vec![(0, 0.8), (2, 0.3)]);
158        doc.add_sparse_vector(sparse, vec![(1, 0.9), (3, 0.4)]);
159        builder.add_document(doc).unwrap();
160
161        // Doc 2: single sparse vector
162        let mut doc = crate::dsl::Document::new();
163        doc.add_sparse_vector(sparse, vec![(2, 1.0), (3, 0.5)]);
164        builder.add_document(doc).unwrap();
165
166        builder.build(&dir, segment_id).await.unwrap();
167
168        let reader = AsyncSegmentReader::open(&dir, segment_id, schema.clone(), 0, 16)
169            .await
170            .unwrap();
171
172        // Query matching dimension 0
173        let query = vec![(0u32, 1.0f32)];
174        let results = reader
175            .search_sparse_vector(sparse, &query, 10, MultiValueCombiner::Sum)
176            .await
177            .unwrap();
178
179        // Both doc 0 and doc 1 have dimension 0
180        assert!(results.len() >= 2, "Should have at least 2 results");
181
182        // Check doc 1 has ordinal tracking
183        let doc1_result = results.iter().find(|r| r.doc_id == 1);
184        assert!(doc1_result.is_some(), "Doc 1 should be in results");
185
186        let doc1 = doc1_result.unwrap();
187        // Doc 1's first sparse vector has dim 0, so ordinal should be 0
188        assert!(
189            !doc1.ordinals.is_empty(),
190            "Doc 1 should have ordinal information"
191        );
192
193        // Check ordinals are valid (0 or 1)
194        for (ordinal, _score) in &doc1.ordinals {
195            assert!(*ordinal <= 1, "Ordinal should be 0 or 1, got {}", ordinal);
196        }
197    }
198}
hermes_core/segment/mod.rs

hermes_core/segment/
mod.rs
