hermes_core/segment/
mod.rs

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

hermes_core/segment/
mod.rs
