Skip to main content

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 = 2][num_sections u32]
11//! TOC × num_sections: [field_id u32][kind u32][count u32][total_tokens u64][data_offset u64]
12//! kind 0 (chunk map):   doc_ids u32 × n | ordinals u16 × n | lengths u16 × n
13//! kind 1 (doc lengths): lengths u16 × num_docs        (norms of a plain text field)
14//! ```
15//!
16//! Version 1 files have 24-byte entries without `kind` and hold chunk maps
17//! only; they are still read.
18//!
19//! Virtual ids are assigned in indexing order, and documents are indexed in
20//! doc-id order, so `doc_ids` starts out non-decreasing. A reorder pass on a
21//! field with the `reorder` attribute permutes the virtual ids (BP over the
22//! field's postings, `segment/text_reorder.rs`); no query path depends on the
23//! order. Merges concatenate sections and add the document offset to
24//! `doc_ids`; ordinals and lengths are copied verbatim.
25//!
26//! A doc-length section stores the token count of the field in every
27//! document of the segment (0 when the document has no value), so BM25 can
28//! normalise plain fields by their real length instead of `tf`.
29
30use std::io::{self, Write};
31
32use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
33use rustc_hash::FxHashMap;
34
35use crate::DocId;
36use crate::directories::OwnedBytes;
37
38const MAGIC: u32 = 0x4B4E_4843; // "CHNK"
39const VERSION: u32 = 2;
40const HEADER_SIZE: usize = 12;
41const TOC_ENTRY_SIZE_V1: usize = 24;
42const TOC_ENTRY_SIZE: usize = 28;
43const KIND_CHUNK_MAP: u32 = 0;
44const KIND_DOC_LENGTHS: u32 = 1;
45
46/// Token count stored per chunk; longer chunks saturate.
47pub const MAX_CHUNK_LENGTH: u32 = u16::MAX as u32;
48
49/// In-memory map of one chunked field while a segment is being built.
50#[derive(Debug, Default, Clone)]
51pub struct ChunkMapBuilder {
52    doc_ids: Vec<DocId>,
53    ordinals: Vec<u16>,
54    lengths: Vec<u16>,
55    total_tokens: u64,
56}
57
58impl ChunkMapBuilder {
59    /// Number of chunks so far (the next virtual id).
60    pub fn len(&self) -> usize {
61        self.doc_ids.len()
62    }
63
64    pub fn is_empty(&self) -> bool {
65        self.doc_ids.is_empty()
66    }
67
68    /// Register the next chunk. Returns its virtual id.
69    pub fn push(&mut self, doc_id: DocId, ordinal: u16, token_count: u32) -> io::Result<u32> {
70        let vid = u32::try_from(self.doc_ids.len()).map_err(|_| {
71            io::Error::new(
72                io::ErrorKind::InvalidData,
73                "chunked text field exceeds u32::MAX chunks in one segment",
74            )
75        })?;
76        self.doc_ids.push(doc_id);
77        self.ordinals.push(ordinal);
78        self.lengths.push(token_count.min(MAX_CHUNK_LENGTH) as u16);
79        self.total_tokens += u64::from(token_count);
80        Ok(vid)
81    }
82
83    /// Heap bytes held by this builder (memory-budget accounting).
84    pub fn estimated_bytes(&self) -> usize {
85        self.doc_ids.capacity() * 4 + self.ordinals.capacity() * 2 + self.lengths.capacity() * 2
86    }
87
88    fn section_bytes(&self) -> u64 {
89        self.doc_ids.len() as u64 * 8
90    }
91
92    /// Token count of virtual id `vid` (saturated at `MAX_CHUNK_LENGTH`).
93    pub fn length(&self, vid: u32) -> u32 {
94        self.lengths
95            .get(vid as usize)
96            .map_or(0, |len| u32::from(*len))
97    }
98}
99
100/// Per-document token counts of one plain text field, ready to be written.
101pub struct DocLengthsColumn<'a> {
102    pub field_id: u32,
103    /// One entry per document of the segment (0 = no value).
104    pub lengths: &'a [u16],
105    /// Sum of the unsaturated token counts.
106    pub total_tokens: u64,
107}
108
109/// Write every chunked field's map and every plain field's length column as
110/// one `.chunks` file.
111///
112/// `fields` must be sorted by field id and contain only non-empty builders;
113/// `norms` likewise sorted, one column per field.
114pub fn write_chunk_maps<W: Write + ?Sized>(
115    writer: &mut W,
116    fields: &[(u32, &ChunkMapBuilder)],
117    norms: &[DocLengthsColumn<'_>],
118) -> io::Result<u64> {
119    let sections = fields.len() + norms.len();
120    let mut offset = (HEADER_SIZE + TOC_ENTRY_SIZE * sections) as u64;
121    writer.write_u32::<LittleEndian>(MAGIC)?;
122    writer.write_u32::<LittleEndian>(VERSION)?;
123    writer.write_u32::<LittleEndian>(sections as u32)?;
124    for (field_id, map) in fields {
125        writer.write_u32::<LittleEndian>(*field_id)?;
126        writer.write_u32::<LittleEndian>(KIND_CHUNK_MAP)?;
127        writer.write_u32::<LittleEndian>(map.len() as u32)?;
128        writer.write_u64::<LittleEndian>(map.total_tokens)?;
129        writer.write_u64::<LittleEndian>(offset)?;
130        offset += map.section_bytes();
131    }
132    for column in norms {
133        writer.write_u32::<LittleEndian>(column.field_id)?;
134        writer.write_u32::<LittleEndian>(KIND_DOC_LENGTHS)?;
135        writer.write_u32::<LittleEndian>(column.lengths.len() as u32)?;
136        writer.write_u64::<LittleEndian>(column.total_tokens)?;
137        writer.write_u64::<LittleEndian>(offset)?;
138        offset += column.lengths.len() as u64 * 2;
139    }
140    for (_, map) in fields {
141        for doc_id in &map.doc_ids {
142            writer.write_u32::<LittleEndian>(*doc_id)?;
143        }
144        for ordinal in &map.ordinals {
145            writer.write_u16::<LittleEndian>(*ordinal)?;
146        }
147        for length in &map.lengths {
148            writer.write_u16::<LittleEndian>(*length)?;
149        }
150    }
151    for column in norms {
152        for length in column.lengths {
153            writer.write_u16::<LittleEndian>(*length)?;
154        }
155    }
156    Ok(offset)
157}
158
159/// Read-only per-document lengths of one plain text field, backed by the
160/// mapped `.chunks` file.
161#[derive(Debug, Clone)]
162pub struct DocLengths {
163    lengths: OwnedBytes,
164    num_docs: u32,
165    total_tokens: u64,
166}
167
168impl DocLengths {
169    /// In-memory lengths column (tests).
170    #[cfg(test)]
171    pub(crate) fn from_lengths(lengths: &[u16]) -> Self {
172        let mut bytes = Vec::with_capacity(lengths.len() * 2);
173        for len in lengths {
174            bytes.extend_from_slice(&len.to_le_bytes());
175        }
176        Self {
177            lengths: OwnedBytes::new(bytes),
178            num_docs: lengths.len() as u32,
179            total_tokens: lengths.iter().map(|&l| u64::from(l)).sum(),
180        }
181    }
182
183    pub fn num_docs(&self) -> u32 {
184        self.num_docs
185    }
186
187    pub fn total_tokens(&self) -> u64 {
188        self.total_tokens
189    }
190
191    /// Average length over documents that have the field (1.0 when none).
192    pub fn avg_len(&self) -> f32 {
193        let with_value = self
194            .lengths
195            .as_slice()
196            .chunks_exact(2)
197            .filter(|b| b[0] != 0 || b[1] != 0)
198            .count();
199        if with_value == 0 {
200            1.0
201        } else {
202            (self.total_tokens as f64 / with_value as f64) as f32
203        }
204    }
205
206    /// Token count of the field in `doc_id` (0 when absent or out of range,
207    /// saturated at `MAX_CHUNK_LENGTH`).
208    #[inline]
209    pub fn length(&self, doc_id: DocId) -> u32 {
210        let at = doc_id as usize * 2;
211        self.lengths
212            .as_slice()
213            .get(at..at + 2)
214            .map_or(0, |b| u32::from(u16::from_le_bytes([b[0], b[1]])))
215    }
216
217    pub(crate) fn length_bytes(&self) -> &[u8] {
218        self.lengths.as_slice()
219    }
220}
221
222/// Everything a `.chunks` file holds.
223#[derive(Debug, Default)]
224pub struct ChunkMapFile {
225    pub chunk_maps: FxHashMap<u32, ChunkMap>,
226    pub doc_lengths: FxHashMap<u32, DocLengths>,
227}
228
229/// Read-only chunk map of one field, backed by the mapped `.chunks` file.
230#[derive(Debug, Clone)]
231pub struct ChunkMap {
232    doc_ids: OwnedBytes,
233    ordinals: OwnedBytes,
234    lengths: OwnedBytes,
235    num_chunks: u32,
236    total_tokens: u64,
237}
238
239impl ChunkMap {
240    /// Number of chunks (virtual ids) in this segment.
241    #[inline]
242    pub fn num_chunks(&self) -> u32 {
243        self.num_chunks
244    }
245
246    /// Sum of all chunk token counts.
247    pub fn total_tokens(&self) -> u64 {
248        self.total_tokens
249    }
250
251    /// Average chunk length in tokens (1.0 when empty).
252    pub fn avg_len(&self) -> f32 {
253        if self.num_chunks == 0 {
254            1.0
255        } else {
256            (self.total_tokens as f64 / f64::from(self.num_chunks)) as f32
257        }
258    }
259
260    /// Document owning virtual id `vid`.
261    #[inline]
262    pub fn doc_id(&self, vid: u32) -> DocId {
263        let at = vid as usize * 4;
264        let b = &self.doc_ids.as_slice()[at..at + 4];
265        u32::from_le_bytes([b[0], b[1], b[2], b[3]])
266    }
267
268    /// Ordinal (value index within the document) of virtual id `vid`.
269    #[inline]
270    pub fn ordinal(&self, vid: u32) -> u16 {
271        let at = vid as usize * 2;
272        let b = &self.ordinals.as_slice()[at..at + 2];
273        u16::from_le_bytes([b[0], b[1]])
274    }
275
276    /// Token count of virtual id `vid` (saturated at `MAX_CHUNK_LENGTH`).
277    #[inline]
278    pub fn length(&self, vid: u32) -> u32 {
279        let at = vid as usize * 2;
280        let b = &self.lengths.as_slice()[at..at + 2];
281        u32::from(u16::from_le_bytes([b[0], b[1]]))
282    }
283
284    /// `(doc_id, ordinal)` of virtual id `vid`.
285    #[inline]
286    pub fn resolve(&self, vid: u32) -> (DocId, u16) {
287        (self.doc_id(vid), self.ordinal(vid))
288    }
289
290    /// Raw little-endian document-id column (merge copy).
291    pub(crate) fn doc_id_bytes(&self) -> &[u8] {
292        self.doc_ids.as_slice()
293    }
294
295    /// Raw little-endian ordinal column (merge copy).
296    pub(crate) fn ordinal_bytes(&self) -> &[u8] {
297        self.ordinals.as_slice()
298    }
299
300    /// Raw little-endian length column (merge copy).
301    pub(crate) fn length_bytes(&self) -> &[u8] {
302        self.lengths.as_slice()
303    }
304}
305
306/// Parse a `.chunks` file into per-field chunk maps and length columns.
307pub fn read_chunk_maps(bytes: OwnedBytes) -> io::Result<ChunkMapFile> {
308    let data = bytes.as_slice();
309    if data.len() < HEADER_SIZE {
310        return Err(io::Error::new(
311            io::ErrorKind::InvalidData,
312            "chunk map file shorter than its header",
313        ));
314    }
315    let mut cursor = io::Cursor::new(data);
316    let magic = cursor.read_u32::<LittleEndian>()?;
317    if magic != MAGIC {
318        return Err(io::Error::new(
319            io::ErrorKind::InvalidData,
320            format!("chunk map magic mismatch: {magic:#x}"),
321        ));
322    }
323    let version = cursor.read_u32::<LittleEndian>()?;
324    let entry_size = match version {
325        1 => TOC_ENTRY_SIZE_V1,
326        VERSION => TOC_ENTRY_SIZE,
327        other => {
328            return Err(io::Error::new(
329                io::ErrorKind::InvalidData,
330                format!("unsupported chunk map version {other} (expected {VERSION})"),
331            ));
332        }
333    };
334    let num_sections = cursor.read_u32::<LittleEndian>()? as usize;
335    if data.len() < HEADER_SIZE + entry_size * num_sections {
336        return Err(io::Error::new(
337            io::ErrorKind::InvalidData,
338            "chunk map table of contents truncated",
339        ));
340    }
341    let overflow = || io::Error::new(io::ErrorKind::InvalidData, "chunk map size overflow");
342    let mut file = ChunkMapFile::default();
343    for _ in 0..num_sections {
344        let field_id = cursor.read_u32::<LittleEndian>()?;
345        let kind = if version == 1 {
346            KIND_CHUNK_MAP
347        } else {
348            cursor.read_u32::<LittleEndian>()?
349        };
350        let count = cursor.read_u32::<LittleEndian>()?;
351        let total_tokens = cursor.read_u64::<LittleEndian>()?;
352        let offset = cursor.read_u64::<LittleEndian>()? as usize;
353        let n = count as usize;
354        let bytes_per_entry = match kind {
355            KIND_CHUNK_MAP => 8,
356            KIND_DOC_LENGTHS => 2,
357            other => {
358                return Err(io::Error::new(
359                    io::ErrorKind::InvalidData,
360                    format!("unknown chunk map section kind {other} for field {field_id}"),
361                ));
362            }
363        };
364        let end = offset
365            .checked_add(n.checked_mul(bytes_per_entry).ok_or_else(overflow)?)
366            .ok_or_else(overflow)?;
367        if end > data.len() {
368            return Err(io::Error::new(
369                io::ErrorKind::InvalidData,
370                format!("chunk map section of field {field_id} exceeds file length"),
371            ));
372        }
373        match kind {
374            KIND_CHUNK_MAP => {
375                let doc_ids = bytes.slice(offset..offset + n * 4);
376                let ordinals = bytes.slice(offset + n * 4..offset + n * 6);
377                let lengths = bytes.slice(offset + n * 6..end);
378                file.chunk_maps.insert(
379                    field_id,
380                    ChunkMap {
381                        doc_ids,
382                        ordinals,
383                        lengths,
384                        num_chunks: count,
385                        total_tokens,
386                    },
387                );
388            }
389            _ => {
390                file.doc_lengths.insert(
391                    field_id,
392                    DocLengths {
393                        lengths: bytes.slice(offset..end),
394                        num_docs: count,
395                        total_tokens,
396                    },
397                );
398            }
399        }
400    }
401    Ok(file)
402}
403
404/// One source section of a merged chunk map.
405pub struct ChunkMapSource<'a> {
406    pub map: &'a ChunkMap,
407    /// Added to every document id of the source.
408    pub doc_offset: u32,
409}
410
411/// One source of a merged length column: the source segment's column when it
412/// has one, and its document count (zeros are written for a missing column).
413pub struct DocLengthsSource<'a> {
414    pub lengths: Option<&'a DocLengths>,
415    pub num_docs: u32,
416}
417
418/// Write the merged `.chunks` file: per field, the sources' sections are
419/// concatenated in order (virtual ids of a later source are offset by the
420/// chunk counts of the earlier ones, matching the posting merge; length
421/// columns follow the document order of the merge).
422///
423/// `fields` and `norms` must be sorted by field id; a field with zero total
424/// chunks is skipped.
425pub fn write_merged_chunk_maps<W: Write + ?Sized>(
426    writer: &mut W,
427    fields: &[(u32, Vec<ChunkMapSource<'_>>)],
428    norms: &[(u32, Vec<DocLengthsSource<'_>>)],
429) -> io::Result<u64> {
430    let live: Vec<&(u32, Vec<ChunkMapSource<'_>>)> = fields
431        .iter()
432        .filter(|(_, sources)| sources.iter().any(|s| s.map.num_chunks() > 0))
433        .collect();
434    let sections = live.len() + norms.len();
435    let mut offset = (HEADER_SIZE + TOC_ENTRY_SIZE * sections) as u64;
436    writer.write_u32::<LittleEndian>(MAGIC)?;
437    writer.write_u32::<LittleEndian>(VERSION)?;
438    writer.write_u32::<LittleEndian>(sections as u32)?;
439    for (field_id, sources) in &live {
440        let mut num_chunks = 0u64;
441        let mut total_tokens = 0u64;
442        for source in sources {
443            num_chunks += u64::from(source.map.num_chunks());
444            total_tokens += source.map.total_tokens();
445        }
446        let num_chunks = u32::try_from(num_chunks).map_err(|_| {
447            io::Error::new(
448                io::ErrorKind::InvalidData,
449                format!("chunked field {field_id} exceeds u32::MAX chunks after merge"),
450            )
451        })?;
452        writer.write_u32::<LittleEndian>(*field_id)?;
453        writer.write_u32::<LittleEndian>(KIND_CHUNK_MAP)?;
454        writer.write_u32::<LittleEndian>(num_chunks)?;
455        writer.write_u64::<LittleEndian>(total_tokens)?;
456        writer.write_u64::<LittleEndian>(offset)?;
457        offset += u64::from(num_chunks) * 8;
458    }
459    for (field_id, sources) in norms {
460        let num_docs: u64 = sources.iter().map(|s| u64::from(s.num_docs)).sum();
461        let num_docs = u32::try_from(num_docs).map_err(|_| {
462            io::Error::new(
463                io::ErrorKind::InvalidData,
464                format!("field {field_id} exceeds u32::MAX documents after merge"),
465            )
466        })?;
467        let total_tokens: u64 = sources
468            .iter()
469            .filter_map(|s| s.lengths.map(DocLengths::total_tokens))
470            .sum();
471        writer.write_u32::<LittleEndian>(*field_id)?;
472        writer.write_u32::<LittleEndian>(KIND_DOC_LENGTHS)?;
473        writer.write_u32::<LittleEndian>(num_docs)?;
474        writer.write_u64::<LittleEndian>(total_tokens)?;
475        writer.write_u64::<LittleEndian>(offset)?;
476        offset += u64::from(num_docs) * 2;
477    }
478    let mut patched: Vec<u8> = Vec::new();
479    for (_, sources) in &live {
480        for source in sources {
481            if source.doc_offset == 0 {
482                writer.write_all(source.map.doc_id_bytes())?;
483                continue;
484            }
485            patched.clear();
486            patched.reserve(source.map.doc_id_bytes().len());
487            for chunk in source.map.doc_id_bytes().chunks_exact(4) {
488                let doc = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
489                let remapped = doc.checked_add(source.doc_offset).ok_or_else(|| {
490                    io::Error::new(
491                        io::ErrorKind::InvalidData,
492                        "document id overflow while merging chunk maps",
493                    )
494                })?;
495                patched.extend_from_slice(&remapped.to_le_bytes());
496            }
497            writer.write_all(&patched)?;
498        }
499        for source in sources {
500            writer.write_all(source.map.ordinal_bytes())?;
501        }
502        for source in sources {
503            writer.write_all(source.map.length_bytes())?;
504        }
505    }
506    let zeros = [0u8; 2 * 1024];
507    for (_, sources) in norms {
508        for source in sources {
509            match source.lengths {
510                Some(lengths) if lengths.num_docs() == source.num_docs => {
511                    writer.write_all(lengths.length_bytes())?;
512                }
513                Some(lengths) => {
514                    return Err(io::Error::new(
515                        io::ErrorKind::InvalidData,
516                        format!(
517                            "length column covers {} documents, segment has {}",
518                            lengths.num_docs(),
519                            source.num_docs
520                        ),
521                    ));
522                }
523                None => {
524                    let mut remaining = source.num_docs as usize * 2;
525                    while remaining > 0 {
526                        let take = remaining.min(zeros.len());
527                        writer.write_all(&zeros[..take])?;
528                        remaining -= take;
529                    }
530                }
531            }
532        }
533    }
534    Ok(offset)
535}
536
537#[cfg(test)]
538mod tests {
539    use super::*;
540
541    fn build(entries: &[(u32, u16, u32)]) -> ChunkMapBuilder {
542        let mut builder = ChunkMapBuilder::default();
543        for &(doc, ord, len) in entries {
544            builder.push(doc, ord, len).unwrap();
545        }
546        builder
547    }
548
549    #[test]
550    fn round_trips_two_fields() {
551        let a = build(&[(0, 0, 10), (0, 1, 20), (3, 0, 70_000)]);
552        let b = build(&[(1, 0, 5)]);
553        let mut out = Vec::new();
554        write_chunk_maps(&mut out, &[(2, &a), (7, &b)], &[]).unwrap();
555        let maps = read_chunk_maps(OwnedBytes::new(out)).unwrap().chunk_maps;
556        let a = &maps[&2];
557        assert_eq!(a.num_chunks(), 3);
558        assert_eq!(a.resolve(0), (0, 0));
559        assert_eq!(a.resolve(1), (0, 1));
560        assert_eq!(a.resolve(2), (3, 0));
561        assert_eq!(a.length(1), 20);
562        assert_eq!(a.length(2), MAX_CHUNK_LENGTH, "lengths saturate at u16");
563        assert_eq!(a.total_tokens(), 70_030);
564        assert_eq!(maps[&7].resolve(0), (1, 0));
565        assert_eq!(maps[&7].avg_len(), 5.0);
566    }
567
568    #[test]
569    fn merged_maps_offset_doc_ids_and_keep_ordinals() {
570        let first = build(&[(0, 0, 10), (1, 0, 11), (1, 1, 12)]);
571        let second = build(&[(0, 0, 20), (0, 1, 21)]);
572        let mut raw_first = Vec::new();
573        write_chunk_maps(&mut raw_first, &[(4, &first)], &[]).unwrap();
574        let mut raw_second = Vec::new();
575        write_chunk_maps(&mut raw_second, &[(4, &second)], &[]).unwrap();
576        let first = read_chunk_maps(OwnedBytes::new(raw_first))
577            .unwrap()
578            .chunk_maps;
579        let second = read_chunk_maps(OwnedBytes::new(raw_second))
580            .unwrap()
581            .chunk_maps;
582
583        let mut merged = Vec::new();
584        write_merged_chunk_maps(
585            &mut merged,
586            &[(
587                4,
588                vec![
589                    ChunkMapSource {
590                        map: &first[&4],
591                        doc_offset: 0,
592                    },
593                    ChunkMapSource {
594                        map: &second[&4],
595                        doc_offset: 2,
596                    },
597                ],
598            )],
599            &[],
600        )
601        .unwrap();
602        let merged = read_chunk_maps(OwnedBytes::new(merged)).unwrap().chunk_maps;
603        let map = &merged[&4];
604        assert_eq!(map.num_chunks(), 5);
605        assert_eq!(map.total_tokens(), 74);
606        assert_eq!(
607            (0..5).map(|v| map.resolve(v)).collect::<Vec<_>>(),
608            vec![(0, 0), (1, 0), (1, 1), (2, 0), (2, 1)]
609        );
610        assert_eq!(
611            (0..5).map(|v| map.length(v)).collect::<Vec<_>>(),
612            vec![10, 11, 12, 20, 21]
613        );
614    }
615
616    #[test]
617    fn rejects_foreign_or_truncated_files() {
618        assert!(read_chunk_maps(OwnedBytes::new(vec![0u8; 4])).is_err());
619        let mut bad_magic = Vec::new();
620        bad_magic.write_u32::<LittleEndian>(0xDEAD_BEEF).unwrap();
621        bad_magic.write_u32::<LittleEndian>(VERSION).unwrap();
622        bad_magic.write_u32::<LittleEndian>(0).unwrap();
623        assert!(read_chunk_maps(OwnedBytes::new(bad_magic)).is_err());
624
625        let a = build(&[(0, 0, 10)]);
626        let mut out = Vec::new();
627        write_chunk_maps(&mut out, &[(1, &a)], &[]).unwrap();
628        out.truncate(out.len() - 1);
629        assert!(read_chunk_maps(OwnedBytes::new(out)).is_err());
630    }
631
632    #[test]
633    fn doc_length_columns_round_trip_and_merge_with_zero_fill() {
634        let a = build(&[(0, 0, 10)]);
635        let column = [7u16, 0, 300];
636        let mut out = Vec::new();
637        write_chunk_maps(
638            &mut out,
639            &[(1, &a)],
640            &[DocLengthsColumn {
641                field_id: 5,
642                lengths: &column,
643                total_tokens: 307,
644            }],
645        )
646        .unwrap();
647        let file = read_chunk_maps(OwnedBytes::new(out)).unwrap();
648        assert_eq!(file.chunk_maps[&1].num_chunks(), 1);
649        let norms = &file.doc_lengths[&5];
650        assert_eq!(norms.num_docs(), 3);
651        assert_eq!(
652            (0..4).map(|d| norms.length(d)).collect::<Vec<_>>(),
653            vec![7, 0, 300, 0]
654        );
655        assert_eq!(norms.total_tokens(), 307);
656        assert!(
657            (norms.avg_len() - 153.5).abs() < 1e-3,
658            "{}",
659            norms.avg_len()
660        );
661
662        // Merge: a source without the column contributes zeros for its docs.
663        let mut merged = Vec::new();
664        write_merged_chunk_maps(
665            &mut merged,
666            &[],
667            &[(
668                5,
669                vec![
670                    DocLengthsSource {
671                        lengths: None,
672                        num_docs: 2,
673                    },
674                    DocLengthsSource {
675                        lengths: Some(norms),
676                        num_docs: 3,
677                    },
678                ],
679            )],
680        )
681        .unwrap();
682        let merged = read_chunk_maps(OwnedBytes::new(merged)).unwrap();
683        assert!(merged.chunk_maps.is_empty());
684        let norms = &merged.doc_lengths[&5];
685        assert_eq!(norms.num_docs(), 5);
686        assert_eq!(
687            (0..5).map(|d| norms.length(d)).collect::<Vec<_>>(),
688            vec![0, 0, 7, 0, 300]
689        );
690        assert_eq!(norms.total_tokens(), 307);
691    }
692
693    #[test]
694    fn version_one_files_still_read() {
695        let a = build(&[(0, 0, 10), (2, 0, 4)]);
696        let mut out = Vec::new();
697        out.write_u32::<LittleEndian>(MAGIC).unwrap();
698        out.write_u32::<LittleEndian>(1).unwrap();
699        out.write_u32::<LittleEndian>(1).unwrap();
700        out.write_u32::<LittleEndian>(9).unwrap();
701        out.write_u32::<LittleEndian>(2).unwrap();
702        out.write_u64::<LittleEndian>(14).unwrap();
703        out.write_u64::<LittleEndian>((HEADER_SIZE + TOC_ENTRY_SIZE_V1) as u64)
704            .unwrap();
705        for doc in &a.doc_ids {
706            out.write_u32::<LittleEndian>(*doc).unwrap();
707        }
708        for ord in &a.ordinals {
709            out.write_u16::<LittleEndian>(*ord).unwrap();
710        }
711        for len in &a.lengths {
712            out.write_u16::<LittleEndian>(*len).unwrap();
713        }
714        let file = read_chunk_maps(OwnedBytes::new(out)).unwrap();
715        assert!(file.doc_lengths.is_empty());
716        let map = &file.chunk_maps[&9];
717        assert_eq!(map.resolve(1), (2, 0));
718        assert_eq!(map.length(1), 4);
719    }
720}