hermes_core/segment/
mod.rs

1pub(crate) mod ann_build;
2#[cfg(any(feature = "native", feature = "wasm"))]
3mod builder;
4pub(crate) mod format;
5#[cfg(feature = "native")]
6mod merger;
7pub(crate) mod reader;
8#[cfg(feature = "native")]
9pub(crate) mod reorder;
10mod store;
11#[cfg(feature = "native")]
12mod tracker;
13mod types;
14mod vector_data;
15
16#[cfg(any(feature = "native", feature = "wasm"))]
17pub use builder::{
18    BpBudget, MemoryBreakdown, SegmentBuilder, SegmentBuilderConfig, SegmentBuilderStats,
19};
20#[cfg(feature = "native")]
21pub use merger::{MergeStats, SegmentMerger, delete_segment};
22pub(crate) use reader::BmpIndex;
23pub(crate) use reader::bmp::BMP_SUPERBLOCK_SIZE;
24pub(crate) use reader::bmp::{
25    accumulate_u4_weighted, block_term_postings, compute_block_masks_4bit, find_dim_in_block_data,
26};
27pub(crate) use reader::combine_ordinal_results;
28#[cfg(feature = "native")]
29pub mod pin;
30pub use reader::{SegmentReader, SparseIndex, VectorIndex, VectorSearchResult};
31pub use store::*;
32#[cfg(feature = "native")]
33pub use tracker::{SegmentSnapshot, SegmentTracker};
34pub use types::{FieldStats, SegmentFiles, SegmentId, SegmentMeta, TrainedVectorStructures};
35pub use vector_data::{
36    FlatVectorData, IVFRaBitQIndexData, LazyFlatVectorData, ScaNNIndexData, dequantize_raw,
37};
38
39/// Format byte count as human-readable string
40#[cfg(any(feature = "native", feature = "wasm"))]
41pub(crate) fn format_bytes(bytes: usize) -> String {
42    if bytes >= 1024 * 1024 * 1024 {
43        format!("{:.2} GB", bytes as f64 / (1024.0 * 1024.0 * 1024.0))
44    } else if bytes >= 1024 * 1024 {
45        format!("{:.2} MB", bytes as f64 / (1024.0 * 1024.0))
46    } else if bytes >= 1024 {
47        format!("{:.2} KB", bytes as f64 / 1024.0)
48    } else {
49        format!("{} B", bytes)
50    }
51}
52
53/// Write adapter that tracks bytes written.
54///
55/// Concrete type so it works with generic `serialize<W: Write>` functions
56/// (unlike `dyn StreamingWriter` which isn't `Sized`).
57#[cfg(any(feature = "native", feature = "wasm"))]
58pub(crate) struct OffsetWriter {
59    inner: Box<dyn crate::directories::StreamingWriter>,
60    offset: u64,
61}
62
63#[cfg(any(feature = "native", feature = "wasm"))]
64impl OffsetWriter {
65    pub(crate) fn new(inner: Box<dyn crate::directories::StreamingWriter>) -> Self {
66        Self { inner, offset: 0 }
67    }
68
69    /// Current write position (total bytes written so far).
70    pub(crate) fn offset(&self) -> u64 {
71        self.offset
72    }
73
74    /// Finalize the underlying streaming writer.
75    pub(crate) fn finish(self) -> std::io::Result<()> {
76        self.inner.finish()
77    }
78}
79
80#[cfg(any(feature = "native", feature = "wasm"))]
81impl std::io::Write for OffsetWriter {
82    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
83        let n = self.inner.write(buf)?;
84        self.offset += n as u64;
85        Ok(n)
86    }
87
88    fn flush(&mut self) -> std::io::Result<()> {
89        self.inner.flush()
90    }
91}
92
93#[cfg(test)]
94#[cfg(feature = "native")]
95mod tests {
96    use super::*;
97    use crate::directories::RamDirectory;
98    use crate::dsl::SchemaBuilder;
99    use std::sync::Arc;
100
101    #[tokio::test]
102    async fn test_async_segment_reader() {
103        let mut schema_builder = SchemaBuilder::default();
104        let title = schema_builder.add_text_field("title", true, true);
105        let schema = Arc::new(schema_builder.build());
106
107        let dir = RamDirectory::new();
108        let segment_id = SegmentId::new();
109
110        // Build segment using sync builder
111        let config = SegmentBuilderConfig::default();
112        let mut builder = SegmentBuilder::new(Arc::clone(&schema), config).unwrap();
113
114        let mut doc = crate::dsl::Document::new();
115        doc.add_text(title, "Hello World");
116        builder.add_document(doc).unwrap();
117
118        let mut doc = crate::dsl::Document::new();
119        doc.add_text(title, "Goodbye World");
120        builder.add_document(doc).unwrap();
121
122        builder.build(&dir, segment_id, None).await.unwrap();
123
124        // Open with async reader
125        let reader = SegmentReader::open(&dir, segment_id, schema.clone(), 16)
126            .await
127            .unwrap();
128
129        assert_eq!(reader.num_docs(), 2);
130
131        // Test postings lookup
132        let postings = reader.get_postings(title, b"hello").await.unwrap();
133        assert!(postings.is_some());
134        assert_eq!(postings.unwrap().doc_count(), 1);
135
136        let postings = reader.get_postings(title, b"world").await.unwrap();
137        assert!(postings.is_some());
138        assert_eq!(postings.unwrap().doc_count(), 2);
139
140        // Test document retrieval
141        let doc = reader.doc(0).await.unwrap().unwrap();
142        assert_eq!(doc.get_first(title).unwrap().as_text(), Some("Hello World"));
143    }
144
145    #[tokio::test]
146    async fn test_dense_vector_ordinal_tracking() {
147        use crate::query::MultiValueCombiner;
148
149        let mut schema_builder = SchemaBuilder::default();
150        // Use simple add method - defaults to Flat index
151        let embedding = schema_builder.add_dense_vector_field("embedding", 4, true, true);
152        let schema = Arc::new(schema_builder.build());
153
154        let dir = RamDirectory::new();
155        let segment_id = SegmentId::new();
156
157        let config = SegmentBuilderConfig::default();
158        let mut builder = SegmentBuilder::new(Arc::clone(&schema), config).unwrap();
159
160        // Doc 0: single vector
161        let mut doc = crate::dsl::Document::new();
162        doc.add_dense_vector(embedding, vec![1.0, 0.0, 0.0, 0.0]);
163        builder.add_document(doc).unwrap();
164
165        // Doc 1: multi-valued vectors (2 vectors)
166        let mut doc = crate::dsl::Document::new();
167        doc.add_dense_vector(embedding, vec![0.0, 1.0, 0.0, 0.0]);
168        doc.add_dense_vector(embedding, vec![0.0, 0.0, 1.0, 0.0]);
169        builder.add_document(doc).unwrap();
170
171        // Doc 2: single vector
172        let mut doc = crate::dsl::Document::new();
173        doc.add_dense_vector(embedding, vec![0.0, 0.0, 0.0, 1.0]);
174        builder.add_document(doc).unwrap();
175
176        builder.build(&dir, segment_id, None).await.unwrap();
177
178        let reader = SegmentReader::open(&dir, segment_id, schema.clone(), 16)
179            .await
180            .unwrap();
181
182        // Query close to doc 1's first vector
183        let query = vec![0.0, 0.9, 0.1, 0.0];
184        let results = reader
185            .search_dense_vector(embedding, &query, 10, 0, 1.0, MultiValueCombiner::Max)
186            .await
187            .unwrap();
188
189        // Doc 1 should be in results with ordinal tracking
190        let doc1_result = results.iter().find(|r| r.doc_id == 1);
191        assert!(doc1_result.is_some(), "Doc 1 should be in results");
192
193        let doc1 = doc1_result.unwrap();
194        // Should have 2 ordinals (0 and 1) for the two vectors
195        assert!(
196            doc1.ordinals.len() <= 2,
197            "Doc 1 should have at most 2 ordinals, got {}",
198            doc1.ordinals.len()
199        );
200
201        // Check ordinals are valid (0 or 1)
202        for (ordinal, _score) in &doc1.ordinals {
203            assert!(*ordinal <= 1, "Ordinal should be 0 or 1, got {}", ordinal);
204        }
205    }
206
207    #[tokio::test]
208    async fn test_sparse_vector_ordinal_tracking() {
209        use crate::query::MultiValueCombiner;
210
211        let mut schema_builder = SchemaBuilder::default();
212        let sparse = schema_builder.add_sparse_vector_field("sparse", true, true);
213        let schema = Arc::new(schema_builder.build());
214
215        let dir = RamDirectory::new();
216        let segment_id = SegmentId::new();
217
218        let config = SegmentBuilderConfig::default();
219        let mut builder = SegmentBuilder::new(Arc::clone(&schema), config).unwrap();
220
221        // Doc 0: single sparse vector
222        let mut doc = crate::dsl::Document::new();
223        doc.add_sparse_vector(sparse, vec![(0, 1.0), (1, 0.5)]);
224        builder.add_document(doc).unwrap();
225
226        // Doc 1: multi-valued sparse vectors (2 vectors)
227        let mut doc = crate::dsl::Document::new();
228        doc.add_sparse_vector(sparse, vec![(0, 0.8), (2, 0.3)]);
229        doc.add_sparse_vector(sparse, vec![(1, 0.9), (3, 0.4)]);
230        builder.add_document(doc).unwrap();
231
232        // Doc 2: single sparse vector
233        let mut doc = crate::dsl::Document::new();
234        doc.add_sparse_vector(sparse, vec![(2, 1.0), (3, 0.5)]);
235        builder.add_document(doc).unwrap();
236
237        builder.build(&dir, segment_id, None).await.unwrap();
238
239        let reader = SegmentReader::open(&dir, segment_id, schema.clone(), 16)
240            .await
241            .unwrap();
242
243        // Query matching dimension 0 via SparseVectorQuery
244        let query = crate::query::SparseVectorQuery::new(sparse, vec![(0, 1.0)])
245            .with_combiner(MultiValueCombiner::Sum);
246        let mut collector = crate::query::TopKCollector::new(10);
247        crate::query::collect_segment(&reader, &query, &mut collector)
248            .await
249            .unwrap();
250        let top_docs = collector.into_sorted_results();
251
252        // Both doc 0 and doc 1 have dimension 0
253        assert!(top_docs.len() >= 2, "Should have at least 2 results");
254    }
255}
hermes_core/segment/mod.rs

hermes_core/segment/
mod.rs
