hermes_core/segment/
chunk_map.rs

1//! Virtual-id maps of chunked text fields (`seg_<id>.chunks`).
2//!
3//! A text field declared `chunked` indexes every value as its own scoring
4//! unit: term postings and positions are keyed by a dense, segment-local
5//! **virtual id** instead of the document id. This file maps each virtual id
6//! back to `(doc_id, ordinal)` and records the chunk's token count for BM25
7//! length normalisation. See `docs/chunked-text-fields.md`.
8//!
9//! ```text
10//! [magic "CHNK"][version u32 = 1][num_fields u32]
11//! TOC × num_fields: [field_id u32][num_chunks u32][total_tokens u64][data_offset u64]
12//! per field:        doc_ids u32 × n | ordinals u16 × n | lengths u16 × n
13//! ```
14//!
15//! Virtual ids are assigned in indexing order, and documents are indexed in
16//! doc-id order, so `doc_ids` is non-decreasing and `(doc_id, ordinal)` is
17//! strictly increasing. Merges concatenate sections and add the document
18//! offset to `doc_ids`; ordinals and lengths are copied verbatim.
19
20use std::io::{self, Write};
21
22use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
23use rustc_hash::FxHashMap;
24
25use crate::DocId;
26use crate::directories::OwnedBytes;
27
28const MAGIC: u32 = 0x4B4E_4843; // "CHNK"
29const VERSION: u32 = 1;
30const HEADER_SIZE: usize = 12;
31const TOC_ENTRY_SIZE: usize = 24;
32
33/// Token count stored per chunk; longer chunks saturate.
34pub const MAX_CHUNK_LENGTH: u32 = u16::MAX as u32;
35
36/// In-memory map of one chunked field while a segment is being built.
37#[derive(Debug, Default, Clone)]
38pub struct ChunkMapBuilder {
39    doc_ids: Vec<DocId>,
40    ordinals: Vec<u16>,
41    lengths: Vec<u16>,
42    total_tokens: u64,
43}
44
45impl ChunkMapBuilder {
46    /// Number of chunks so far (the next virtual id).
47    pub fn len(&self) -> usize {
48        self.doc_ids.len()
49    }
50
51    pub fn is_empty(&self) -> bool {
52        self.doc_ids.is_empty()
53    }
54
55    /// Register the next chunk. Returns its virtual id.
56    pub fn push(&mut self, doc_id: DocId, ordinal: u16, token_count: u32) -> io::Result<u32> {
57        let vid = u32::try_from(self.doc_ids.len()).map_err(|_| {
58            io::Error::new(
59                io::ErrorKind::InvalidData,
60                "chunked text field exceeds u32::MAX chunks in one segment",
61            )
62        })?;
63        self.doc_ids.push(doc_id);
64        self.ordinals.push(ordinal);
65        self.lengths.push(token_count.min(MAX_CHUNK_LENGTH) as u16);
66        self.total_tokens += u64::from(token_count);
67        Ok(vid)
68    }
69
70    /// Heap bytes held by this builder (memory-budget accounting).
71    pub fn estimated_bytes(&self) -> usize {
72        self.doc_ids.capacity() * 4 + self.ordinals.capacity() * 2 + self.lengths.capacity() * 2
73    }
74
75    fn section_bytes(&self) -> u64 {
76        self.doc_ids.len() as u64 * 8
77    }
78}
79
80/// Write every chunked field's map as one `.chunks` file.
81///
82/// `fields` must be sorted by field id and contain only non-empty builders.
83pub fn write_chunk_maps<W: Write + ?Sized>(
84    writer: &mut W,
85    fields: &[(u32, &ChunkMapBuilder)],
86) -> io::Result<u64> {
87    let mut offset = (HEADER_SIZE + TOC_ENTRY_SIZE * fields.len()) as u64;
88    writer.write_u32::<LittleEndian>(MAGIC)?;
89    writer.write_u32::<LittleEndian>(VERSION)?;
90    writer.write_u32::<LittleEndian>(fields.len() as u32)?;
91    for (field_id, map) in fields {
92        writer.write_u32::<LittleEndian>(*field_id)?;
93        writer.write_u32::<LittleEndian>(map.len() as u32)?;
94        writer.write_u64::<LittleEndian>(map.total_tokens)?;
95        writer.write_u64::<LittleEndian>(offset)?;
96        offset += map.section_bytes();
97    }
98    for (_, map) in fields {
99        for doc_id in &map.doc_ids {
100            writer.write_u32::<LittleEndian>(*doc_id)?;
101        }
102        for ordinal in &map.ordinals {
103            writer.write_u16::<LittleEndian>(*ordinal)?;
104        }
105        for length in &map.lengths {
106            writer.write_u16::<LittleEndian>(*length)?;
107        }
108    }
109    Ok(offset)
110}
111
112/// Read-only chunk map of one field, backed by the mapped `.chunks` file.
113#[derive(Debug, Clone)]
114pub struct ChunkMap {
115    doc_ids: OwnedBytes,
116    ordinals: OwnedBytes,
117    lengths: OwnedBytes,
118    num_chunks: u32,
119    total_tokens: u64,
120}
121
122impl ChunkMap {
123    /// Number of chunks (virtual ids) in this segment.
124    #[inline]
125    pub fn num_chunks(&self) -> u32 {
126        self.num_chunks
127    }
128
129    /// Sum of all chunk token counts.
130    pub fn total_tokens(&self) -> u64 {
131        self.total_tokens
132    }
133
134    /// Average chunk length in tokens (1.0 when empty).
135    pub fn avg_len(&self) -> f32 {
136        if self.num_chunks == 0 {
137            1.0
138        } else {
139            (self.total_tokens as f64 / f64::from(self.num_chunks)) as f32
140        }
141    }
142
143    /// Document owning virtual id `vid`.
144    #[inline]
145    pub fn doc_id(&self, vid: u32) -> DocId {
146        let at = vid as usize * 4;
147        let b = &self.doc_ids.as_slice()[at..at + 4];
148        u32::from_le_bytes([b[0], b[1], b[2], b[3]])
149    }
150
151    /// Ordinal (value index within the document) of virtual id `vid`.
152    #[inline]
153    pub fn ordinal(&self, vid: u32) -> u16 {
154        let at = vid as usize * 2;
155        let b = &self.ordinals.as_slice()[at..at + 2];
156        u16::from_le_bytes([b[0], b[1]])
157    }
158
159    /// Token count of virtual id `vid` (saturated at `MAX_CHUNK_LENGTH`).
160    #[inline]
161    pub fn length(&self, vid: u32) -> u32 {
162        let at = vid as usize * 2;
163        let b = &self.lengths.as_slice()[at..at + 2];
164        u32::from(u16::from_le_bytes([b[0], b[1]]))
165    }
166
167    /// `(doc_id, ordinal)` of virtual id `vid`.
168    #[inline]
169    pub fn resolve(&self, vid: u32) -> (DocId, u16) {
170        (self.doc_id(vid), self.ordinal(vid))
171    }
172
173    /// Raw little-endian document-id column (merge copy).
174    pub(crate) fn doc_id_bytes(&self) -> &[u8] {
175        self.doc_ids.as_slice()
176    }
177
178    /// Raw little-endian ordinal column (merge copy).
179    pub(crate) fn ordinal_bytes(&self) -> &[u8] {
180        self.ordinals.as_slice()
181    }
182
183    /// Raw little-endian length column (merge copy).
184    pub(crate) fn length_bytes(&self) -> &[u8] {
185        self.lengths.as_slice()
186    }
187}
188
189/// Parse a `.chunks` file into per-field maps.
190pub fn read_chunk_maps(bytes: OwnedBytes) -> io::Result<FxHashMap<u32, ChunkMap>> {
191    let data = bytes.as_slice();
192    if data.len() < HEADER_SIZE {
193        return Err(io::Error::new(
194            io::ErrorKind::InvalidData,
195            "chunk map file shorter than its header",
196        ));
197    }
198    let mut cursor = io::Cursor::new(data);
199    let magic = cursor.read_u32::<LittleEndian>()?;
200    if magic != MAGIC {
201        return Err(io::Error::new(
202            io::ErrorKind::InvalidData,
203            format!("chunk map magic mismatch: {magic:#x}"),
204        ));
205    }
206    let version = cursor.read_u32::<LittleEndian>()?;
207    if version != VERSION {
208        return Err(io::Error::new(
209            io::ErrorKind::InvalidData,
210            format!("unsupported chunk map version {version} (expected {VERSION})"),
211        ));
212    }
213    let num_fields = cursor.read_u32::<LittleEndian>()? as usize;
214    if data.len() < HEADER_SIZE + TOC_ENTRY_SIZE * num_fields {
215        return Err(io::Error::new(
216            io::ErrorKind::InvalidData,
217            "chunk map table of contents truncated",
218        ));
219    }
220    let mut maps = FxHashMap::default();
221    for _ in 0..num_fields {
222        let field_id = cursor.read_u32::<LittleEndian>()?;
223        let num_chunks = cursor.read_u32::<LittleEndian>()?;
224        let total_tokens = cursor.read_u64::<LittleEndian>()?;
225        let offset = cursor.read_u64::<LittleEndian>()? as usize;
226        let n = num_chunks as usize;
227        let end = offset
228            .checked_add(n.checked_mul(8).ok_or_else(|| {
229                io::Error::new(io::ErrorKind::InvalidData, "chunk map size overflow")
230            })?)
231            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "chunk map size overflow"))?;
232        if end > data.len() {
233            return Err(io::Error::new(
234                io::ErrorKind::InvalidData,
235                format!("chunk map section of field {field_id} exceeds file length"),
236            ));
237        }
238        let doc_ids = bytes.slice(offset..offset + n * 4);
239        let ordinals = bytes.slice(offset + n * 4..offset + n * 6);
240        let lengths = bytes.slice(offset + n * 6..end);
241        maps.insert(
242            field_id,
243            ChunkMap {
244                doc_ids,
245                ordinals,
246                lengths,
247                num_chunks,
248                total_tokens,
249            },
250        );
251    }
252    Ok(maps)
253}
254
255/// One source section of a merged chunk map.
256pub struct ChunkMapSource<'a> {
257    pub map: &'a ChunkMap,
258    /// Added to every document id of the source.
259    pub doc_offset: u32,
260}
261
262/// Write the merged `.chunks` file: per field, the sources' sections are
263/// concatenated in order (virtual ids of a later source are offset by the
264/// chunk counts of the earlier ones, matching the posting merge).
265///
266/// `fields` must be sorted by field id; a field with zero total chunks is
267/// skipped.
268pub fn write_merged_chunk_maps<W: Write + ?Sized>(
269    writer: &mut W,
270    fields: &[(u32, Vec<ChunkMapSource<'_>>)],
271) -> io::Result<u64> {
272    let live: Vec<&(u32, Vec<ChunkMapSource<'_>>)> = fields
273        .iter()
274        .filter(|(_, sources)| sources.iter().any(|s| s.map.num_chunks() > 0))
275        .collect();
276    let mut offset = (HEADER_SIZE + TOC_ENTRY_SIZE * live.len()) as u64;
277    writer.write_u32::<LittleEndian>(MAGIC)?;
278    writer.write_u32::<LittleEndian>(VERSION)?;
279    writer.write_u32::<LittleEndian>(live.len() as u32)?;
280    for (field_id, sources) in &live {
281        let mut num_chunks = 0u64;
282        let mut total_tokens = 0u64;
283        for source in sources {
284            num_chunks += u64::from(source.map.num_chunks());
285            total_tokens += source.map.total_tokens();
286        }
287        let num_chunks = u32::try_from(num_chunks).map_err(|_| {
288            io::Error::new(
289                io::ErrorKind::InvalidData,
290                format!("chunked field {field_id} exceeds u32::MAX chunks after merge"),
291            )
292        })?;
293        writer.write_u32::<LittleEndian>(*field_id)?;
294        writer.write_u32::<LittleEndian>(num_chunks)?;
295        writer.write_u64::<LittleEndian>(total_tokens)?;
296        writer.write_u64::<LittleEndian>(offset)?;
297        offset += u64::from(num_chunks) * 8;
298    }
299    let mut patched: Vec<u8> = Vec::new();
300    for (_, sources) in &live {
301        for source in sources {
302            if source.doc_offset == 0 {
303                writer.write_all(source.map.doc_id_bytes())?;
304                continue;
305            }
306            patched.clear();
307            patched.reserve(source.map.doc_id_bytes().len());
308            for chunk in source.map.doc_id_bytes().chunks_exact(4) {
309                let doc = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
310                let remapped = doc.checked_add(source.doc_offset).ok_or_else(|| {
311                    io::Error::new(
312                        io::ErrorKind::InvalidData,
313                        "document id overflow while merging chunk maps",
314                    )
315                })?;
316                patched.extend_from_slice(&remapped.to_le_bytes());
317            }
318            writer.write_all(&patched)?;
319        }
320        for source in sources {
321            writer.write_all(source.map.ordinal_bytes())?;
322        }
323        for source in sources {
324            writer.write_all(source.map.length_bytes())?;
325        }
326    }
327    Ok(offset)
328}
329
330#[cfg(test)]
331mod tests {
332    use super::*;
333
334    fn build(entries: &[(u32, u16, u32)]) -> ChunkMapBuilder {
335        let mut builder = ChunkMapBuilder::default();
336        for &(doc, ord, len) in entries {
337            builder.push(doc, ord, len).unwrap();
338        }
339        builder
340    }
341
342    #[test]
343    fn round_trips_two_fields() {
344        let a = build(&[(0, 0, 10), (0, 1, 20), (3, 0, 70_000)]);
345        let b = build(&[(1, 0, 5)]);
346        let mut out = Vec::new();
347        write_chunk_maps(&mut out, &[(2, &a), (7, &b)]).unwrap();
348        let maps = read_chunk_maps(OwnedBytes::new(out)).unwrap();
349        let a = &maps[&2];
350        assert_eq!(a.num_chunks(), 3);
351        assert_eq!(a.resolve(0), (0, 0));
352        assert_eq!(a.resolve(1), (0, 1));
353        assert_eq!(a.resolve(2), (3, 0));
354        assert_eq!(a.length(1), 20);
355        assert_eq!(a.length(2), MAX_CHUNK_LENGTH, "lengths saturate at u16");
356        assert_eq!(a.total_tokens(), 70_030);
357        assert_eq!(maps[&7].resolve(0), (1, 0));
358        assert_eq!(maps[&7].avg_len(), 5.0);
359    }
360
361    #[test]
362    fn merged_maps_offset_doc_ids_and_keep_ordinals() {
363        let first = build(&[(0, 0, 10), (1, 0, 11), (1, 1, 12)]);
364        let second = build(&[(0, 0, 20), (0, 1, 21)]);
365        let mut raw_first = Vec::new();
366        write_chunk_maps(&mut raw_first, &[(4, &first)]).unwrap();
367        let mut raw_second = Vec::new();
368        write_chunk_maps(&mut raw_second, &[(4, &second)]).unwrap();
369        let first = read_chunk_maps(OwnedBytes::new(raw_first)).unwrap();
370        let second = read_chunk_maps(OwnedBytes::new(raw_second)).unwrap();
371
372        let mut merged = Vec::new();
373        write_merged_chunk_maps(
374            &mut merged,
375            &[(
376                4,
377                vec![
378                    ChunkMapSource {
379                        map: &first[&4],
380                        doc_offset: 0,
381                    },
382                    ChunkMapSource {
383                        map: &second[&4],
384                        doc_offset: 2,
385                    },
386                ],
387            )],
388        )
389        .unwrap();
390        let merged = read_chunk_maps(OwnedBytes::new(merged)).unwrap();
391        let map = &merged[&4];
392        assert_eq!(map.num_chunks(), 5);
393        assert_eq!(map.total_tokens(), 74);
394        assert_eq!(
395            (0..5).map(|v| map.resolve(v)).collect::<Vec<_>>(),
396            vec![(0, 0), (1, 0), (1, 1), (2, 0), (2, 1)]
397        );
398        assert_eq!(
399            (0..5).map(|v| map.length(v)).collect::<Vec<_>>(),
400            vec![10, 11, 12, 20, 21]
401        );
402    }
403
404    #[test]
405    fn rejects_foreign_or_truncated_files() {
406        assert!(read_chunk_maps(OwnedBytes::new(vec![0u8; 4])).is_err());
407        let mut bad_magic = Vec::new();
408        bad_magic.write_u32::<LittleEndian>(0xDEAD_BEEF).unwrap();
409        bad_magic.write_u32::<LittleEndian>(VERSION).unwrap();
410        bad_magic.write_u32::<LittleEndian>(0).unwrap();
411        assert!(read_chunk_maps(OwnedBytes::new(bad_magic)).is_err());
412
413        let a = build(&[(0, 0, 10)]);
414        let mut out = Vec::new();
415        write_chunk_maps(&mut out, &[(1, &a)]).unwrap();
416        out.truncate(out.len() - 1);
417        assert!(read_chunk_maps(OwnedBytes::new(out)).is_err());
418    }
419}
hermes_core/segment/chunk_map.rs

hermes_core/segment/
chunk_map.rs
