1use 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; const 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
46pub const MAX_CHUNK_LENGTH: u32 = u16::MAX as u32;
48
49#[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 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 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 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 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
100pub struct DocLengthsColumn<'a> {
102 pub field_id: u32,
103 pub lengths: &'a [u16],
105 pub total_tokens: u64,
107}
108
109pub 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#[derive(Debug, Clone)]
162pub struct DocLengths {
163 lengths: OwnedBytes,
164 num_docs: u32,
165 total_tokens: u64,
166}
167
168impl DocLengths {
169 #[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 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 #[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#[derive(Debug, Default)]
224pub struct ChunkMapFile {
225 pub chunk_maps: FxHashMap<u32, ChunkMap>,
226 pub doc_lengths: FxHashMap<u32, DocLengths>,
227}
228
229#[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 #[inline]
242 pub fn num_chunks(&self) -> u32 {
243 self.num_chunks
244 }
245
246 pub fn total_tokens(&self) -> u64 {
248 self.total_tokens
249 }
250
251 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 #[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 #[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 #[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 #[inline]
286 pub fn resolve(&self, vid: u32) -> (DocId, u16) {
287 (self.doc_id(vid), self.ordinal(vid))
288 }
289
290 pub(crate) fn doc_id_bytes(&self) -> &[u8] {
292 self.doc_ids.as_slice()
293 }
294
295 pub(crate) fn ordinal_bytes(&self) -> &[u8] {
297 self.ordinals.as_slice()
298 }
299
300 pub(crate) fn length_bytes(&self) -> &[u8] {
302 self.lengths.as_slice()
303 }
304}
305
306pub 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
404pub struct ChunkMapSource<'a> {
406 pub map: &'a ChunkMap,
407 pub doc_offset: u32,
409}
410
411pub struct DocLengthsSource<'a> {
414 pub lengths: Option<&'a DocLengths>,
415 pub num_docs: u32,
416}
417
418pub 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 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}