Skip to main content

hermes_core/segment/
store.rs

1//! Document store with Zstd compression and lazy loading
2//!
3//! Optimized for static indexes:
4//! - Maximum compression level (22) for best compression ratio
5//! - Larger block sizes (256KB) for better compression efficiency
6//! - Optional trained dictionary support for even better compression
7//! - Parallel compression support for faster indexing
8//!
9//! Writer stores documents in compressed blocks.
10//! Reader only loads index into memory, blocks are loaded on-demand.
11
12use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
13use lru::LruCache;
14use parking_lot::RwLock;
15use rustc_hash::FxHashMap;
16use std::io::{self, Write};
17use std::sync::Arc;
18
19use crate::DocId;
20use crate::compression::CompressionDict;
21#[cfg(feature = "native")]
22use crate::compression::CompressionLevel;
23use crate::directories::FileHandle;
24use crate::dsl::{Document, Schema};
25
26const STORE_MAGIC: u32 = 0x53544F52; // "STOR"
27const STORE_VERSION: u32 = 2; // Version 2 supports dictionaries
28
29/// Block size for document store (16KB).
30/// Smaller blocks reduce read amplification for single-doc fetches at the
31/// cost of slightly worse compression ratio. Zstd dictionary training
32/// recovers most of the compression loss.
33pub const STORE_BLOCK_SIZE: usize = 16 * 1024;
34
35/// Default dictionary size (4KB is a good balance)
36pub const DEFAULT_DICT_SIZE: usize = 4 * 1024;
37
38/// Hard safety bounds for individual on-disk store objects. Writers normally
39/// emit ~16 KiB blocks and 4 KiB dictionaries; these generous limits preserve
40/// unusually large stored documents while bounding corrupt compressed frames.
41const MAX_STORE_BLOCK_BYTES: usize = 64 * 1024 * 1024;
42const MAX_STORE_DICTIONARY_BYTES: u64 = 16 * 1024 * 1024;
43
44/// Default compression level for document store
45#[cfg(feature = "native")]
46const DEFAULT_COMPRESSION_LEVEL: CompressionLevel = CompressionLevel(3);
47
48/// Write block index + footer to a store file.
49///
50/// Shared by `EagerParallelStoreWriter::finish` and `StoreMerger::finish`.
51fn write_store_index_and_footer(
52    writer: &mut (impl Write + ?Sized),
53    index: &[StoreBlockIndex],
54    data_end_offset: u64,
55    dict_offset: u64,
56    num_docs: u32,
57    has_dict: bool,
58) -> io::Result<()> {
59    writer.write_u32::<LittleEndian>(u32::try_from(index.len()).map_err(|_| {
60        io::Error::new(
61            io::ErrorKind::InvalidInput,
62            "too many document store blocks",
63        )
64    })?)?;
65    for entry in index {
66        writer.write_u32::<LittleEndian>(entry.first_doc_id)?;
67        writer.write_u64::<LittleEndian>(entry.offset)?;
68        writer.write_u32::<LittleEndian>(entry.length)?;
69        writer.write_u32::<LittleEndian>(entry.num_docs)?;
70    }
71    writer.write_u64::<LittleEndian>(data_end_offset)?;
72    writer.write_u64::<LittleEndian>(dict_offset)?;
73    writer.write_u32::<LittleEndian>(num_docs)?;
74    writer.write_u32::<LittleEndian>(if has_dict { 1 } else { 0 })?;
75    writer.write_u32::<LittleEndian>(STORE_VERSION)?;
76    writer.write_u32::<LittleEndian>(STORE_MAGIC)?;
77    Ok(())
78}
79
80/// Binary document format:
81///   num_fields: u16
82///   per field: field_id: u16, type_tag: u8, value data
83///     0=Text:         len:u32 + utf8
84///     1=U64:          u64 LE
85///     2=I64:          i64 LE
86///     3=F64:          f64 LE
87///     4=Bytes:        len:u32 + raw
88///     5=SparseVector: count:u32 + count*(u32+f32)
89///     6=DenseVector:  count:u32 + count*f32
90///     7=Json:         len:u32 + json utf8
91pub fn serialize_document(doc: &Document, schema: &Schema) -> io::Result<Vec<u8>> {
92    let mut buf = Vec::with_capacity(256);
93    serialize_document_into(doc, schema, &mut buf)?;
94    Ok(buf)
95}
96
97/// Serialize a document into a reusable buffer (clears it first).
98/// Avoids per-document allocation when called in a loop.
99pub fn serialize_document_into(
100    doc: &Document,
101    schema: &Schema,
102    buf: &mut Vec<u8>,
103) -> io::Result<()> {
104    use crate::dsl::FieldValue;
105
106    buf.clear();
107
108    // Two-pass approach avoids allocating a Vec just to count + iterate stored fields.
109    let is_stored = |field: &crate::dsl::Field, value: &FieldValue| -> bool {
110        // Dense/binary vectors live in .vectors (LazyFlatVectorData), not in .store
111        if matches!(
112            value,
113            FieldValue::DenseVector(_) | FieldValue::BinaryDenseVector(_)
114        ) {
115            return false;
116        }
117        schema.get_field_entry(*field).is_some_and(|e| e.stored)
118    };
119
120    let stored_count = doc
121        .field_values()
122        .iter()
123        .filter(|(field, value)| is_stored(field, value))
124        .count();
125
126    buf.write_u16::<LittleEndian>(
127        u16::try_from(stored_count)
128            .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "too many stored fields"))?,
129    )?;
130
131    for (field, value) in doc.field_values().iter().filter(|(f, v)| is_stored(f, v)) {
132        buf.write_u16::<LittleEndian>(u16::try_from(field.0).map_err(|_| {
133            io::Error::new(io::ErrorKind::InvalidInput, "stored field id exceeds u16")
134        })?)?;
135        match value {
136            FieldValue::Text(s) => {
137                buf.push(0);
138                let bytes = s.as_bytes();
139                buf.write_u32::<LittleEndian>(u32::try_from(bytes.len()).map_err(|_| {
140                    io::Error::new(io::ErrorKind::InvalidInput, "stored text is too large")
141                })?)?;
142                buf.extend_from_slice(bytes);
143            }
144            FieldValue::U64(v) => {
145                buf.push(1);
146                buf.write_u64::<LittleEndian>(*v)?;
147            }
148            FieldValue::I64(v) => {
149                buf.push(2);
150                buf.write_i64::<LittleEndian>(*v)?;
151            }
152            FieldValue::F64(v) => {
153                buf.push(3);
154                buf.write_f64::<LittleEndian>(*v)?;
155            }
156            FieldValue::Bytes(b) => {
157                buf.push(4);
158                buf.write_u32::<LittleEndian>(u32::try_from(b.len()).map_err(|_| {
159                    io::Error::new(
160                        io::ErrorKind::InvalidInput,
161                        "stored byte field is too large",
162                    )
163                })?)?;
164                buf.extend_from_slice(b);
165            }
166            FieldValue::SparseVector(entries) => {
167                buf.push(5);
168                buf.write_u32::<LittleEndian>(u32::try_from(entries.len()).map_err(|_| {
169                    io::Error::new(
170                        io::ErrorKind::InvalidInput,
171                        "stored sparse vector is too large",
172                    )
173                })?)?;
174                for (idx, val) in entries {
175                    buf.write_u32::<LittleEndian>(*idx)?;
176                    buf.write_f32::<LittleEndian>(*val)?;
177                }
178            }
179            FieldValue::DenseVector(values) => {
180                buf.push(6);
181                buf.write_u32::<LittleEndian>(u32::try_from(values.len()).map_err(|_| {
182                    io::Error::new(
183                        io::ErrorKind::InvalidInput,
184                        "stored dense vector is too large",
185                    )
186                })?)?;
187                // Write raw f32 bytes directly
188                let byte_slice = unsafe {
189                    std::slice::from_raw_parts(values.as_ptr() as *const u8, values.len() * 4)
190                };
191                buf.extend_from_slice(byte_slice);
192            }
193            FieldValue::Json(v) => {
194                buf.push(7);
195                let json_bytes = serde_json::to_vec(v)
196                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
197                buf.write_u32::<LittleEndian>(u32::try_from(json_bytes.len()).map_err(|_| {
198                    io::Error::new(io::ErrorKind::InvalidInput, "stored JSON is too large")
199                })?)?;
200                buf.extend_from_slice(&json_bytes);
201            }
202            FieldValue::BinaryDenseVector(b) => {
203                buf.push(8);
204                buf.write_u32::<LittleEndian>(u32::try_from(b.len()).map_err(|_| {
205                    io::Error::new(
206                        io::ErrorKind::InvalidInput,
207                        "stored binary dense vector is too large",
208                    )
209                })?)?;
210                buf.extend_from_slice(b);
211            }
212        }
213    }
214
215    Ok(())
216}
217
218/// Compressed block result
219#[cfg(feature = "native")]
220struct CompressedBlock {
221    seq: usize,
222    first_doc_id: DocId,
223    num_docs: u32,
224    compressed: Vec<u8>,
225}
226
227/// Parallel document store writer - compresses blocks immediately when queued
228///
229/// Spawns compression tasks as soon as blocks are ready, overlapping document
230/// ingestion with compression to reduce total indexing time.
231///
232/// Uses background threads to compress blocks while the main thread continues
233/// accepting documents.
234#[cfg(feature = "native")]
235pub struct EagerParallelStoreWriter<'a> {
236    writer: &'a mut dyn Write,
237    block_buffer: Vec<u8>,
238    /// Reusable buffer for document serialization (avoids per-doc allocation)
239    serialize_buf: Vec<u8>,
240    /// Compressed blocks ready to be written (may arrive out of order)
241    compressed_blocks: Vec<CompressedBlock>,
242    /// Handles for in-flight compression tasks
243    pending_handles: Vec<std::thread::JoinHandle<CompressedBlock>>,
244    next_seq: usize,
245    next_doc_id: DocId,
246    block_first_doc: DocId,
247    dict: Option<Arc<CompressionDict>>,
248    compression_level: CompressionLevel,
249}
250
251#[cfg(feature = "native")]
252impl<'a> EagerParallelStoreWriter<'a> {
253    /// Create a new eager parallel store writer
254    pub fn new(writer: &'a mut dyn Write, _num_threads: usize) -> Self {
255        Self::with_compression_level(writer, _num_threads, DEFAULT_COMPRESSION_LEVEL)
256    }
257
258    /// Create with specific compression level
259    pub fn with_compression_level(
260        writer: &'a mut dyn Write,
261        _num_threads: usize,
262        compression_level: CompressionLevel,
263    ) -> Self {
264        Self {
265            writer,
266            block_buffer: Vec::with_capacity(STORE_BLOCK_SIZE),
267            serialize_buf: Vec::with_capacity(512),
268            compressed_blocks: Vec::new(),
269            pending_handles: Vec::new(),
270            next_seq: 0,
271            next_doc_id: 0,
272            block_first_doc: 0,
273            dict: None,
274            compression_level,
275        }
276    }
277
278    /// Create with dictionary
279    pub fn with_dict(
280        writer: &'a mut dyn Write,
281        dict: CompressionDict,
282        _num_threads: usize,
283    ) -> Self {
284        Self::with_dict_and_level(writer, dict, _num_threads, DEFAULT_COMPRESSION_LEVEL)
285    }
286
287    /// Create with dictionary and specific compression level
288    pub fn with_dict_and_level(
289        writer: &'a mut dyn Write,
290        dict: CompressionDict,
291        _num_threads: usize,
292        compression_level: CompressionLevel,
293    ) -> Self {
294        Self {
295            writer,
296            block_buffer: Vec::with_capacity(STORE_BLOCK_SIZE),
297            serialize_buf: Vec::with_capacity(512),
298            compressed_blocks: Vec::new(),
299            pending_handles: Vec::new(),
300            next_seq: 0,
301            next_doc_id: 0,
302            block_first_doc: 0,
303            dict: Some(Arc::new(dict)),
304            compression_level,
305        }
306    }
307
308    pub fn store(&mut self, doc: &Document, schema: &Schema) -> io::Result<DocId> {
309        serialize_document_into(doc, schema, &mut self.serialize_buf)?;
310        if self.serialize_buf.len() > MAX_STORE_BLOCK_BYTES.saturating_sub(4) {
311            return Err(io::Error::new(
312                io::ErrorKind::InvalidInput,
313                "serialized document exceeds store block limit",
314            ));
315        }
316        let doc_id = self.next_doc_id;
317        self.next_doc_id = self
318            .next_doc_id
319            .checked_add(1)
320            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "document id overflow"))?;
321        self.block_buffer
322            .write_u32::<LittleEndian>(self.serialize_buf.len() as u32)?;
323        self.block_buffer.extend_from_slice(&self.serialize_buf);
324        if self.block_buffer.len() >= STORE_BLOCK_SIZE {
325            self.spawn_compression();
326        }
327        Ok(doc_id)
328    }
329
330    /// Store pre-serialized document bytes directly (avoids deserialize+reserialize roundtrip).
331    pub fn store_raw(&mut self, doc_bytes: &[u8]) -> io::Result<DocId> {
332        if doc_bytes.len() > MAX_STORE_BLOCK_BYTES.saturating_sub(4) {
333            return Err(io::Error::new(
334                io::ErrorKind::InvalidInput,
335                "serialized document exceeds store block limit",
336            ));
337        }
338        let doc_id = self.next_doc_id;
339        self.next_doc_id = self
340            .next_doc_id
341            .checked_add(1)
342            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "document id overflow"))?;
343
344        self.block_buffer
345            .write_u32::<LittleEndian>(doc_bytes.len() as u32)?;
346        self.block_buffer.extend_from_slice(doc_bytes);
347
348        if self.block_buffer.len() >= STORE_BLOCK_SIZE {
349            self.spawn_compression();
350        }
351
352        Ok(doc_id)
353    }
354
355    /// Spawn compression for the current block immediately
356    fn spawn_compression(&mut self) {
357        if self.block_buffer.is_empty() {
358            return;
359        }
360
361        let num_docs = self.next_doc_id - self.block_first_doc;
362        let data = std::mem::replace(&mut self.block_buffer, Vec::with_capacity(STORE_BLOCK_SIZE));
363        let seq = self.next_seq;
364        let first_doc_id = self.block_first_doc;
365        let dict = self.dict.clone();
366
367        self.next_seq += 1;
368        self.block_first_doc = self.next_doc_id;
369
370        // Spawn compression task using thread
371        let level = self.compression_level;
372        let handle = std::thread::spawn(move || {
373            let compressed = if let Some(ref d) = dict {
374                crate::compression::compress_with_dict(&data, level, d).expect("compression failed")
375            } else {
376                crate::compression::compress(&data, level).expect("compression failed")
377            };
378
379            CompressedBlock {
380                seq,
381                first_doc_id,
382                num_docs,
383                compressed,
384            }
385        });
386
387        self.pending_handles.push(handle);
388    }
389
390    /// Collect any completed compression tasks
391    fn collect_completed(&mut self) {
392        let mut remaining = Vec::new();
393        for handle in self.pending_handles.drain(..) {
394            if handle.is_finished() {
395                match handle.join() {
396                    Ok(block) => self.compressed_blocks.push(block),
397                    Err(payload) => std::panic::resume_unwind(payload),
398                }
399            } else {
400                remaining.push(handle);
401            }
402        }
403        self.pending_handles = remaining;
404    }
405
406    pub fn finish(mut self) -> io::Result<u32> {
407        // Spawn compression for any remaining data
408        self.spawn_compression();
409
410        // Collect any already-completed tasks
411        self.collect_completed();
412
413        // Wait for all remaining compression tasks
414        for handle in self.pending_handles.drain(..) {
415            match handle.join() {
416                Ok(block) => self.compressed_blocks.push(block),
417                Err(payload) => std::panic::resume_unwind(payload),
418            }
419        }
420
421        if self.compressed_blocks.is_empty() {
422            write_store_index_and_footer(&mut self.writer, &[], 0, 0, 0, false)?;
423            return Ok(0);
424        }
425
426        // Sort by sequence to maintain order
427        self.compressed_blocks.sort_by_key(|b| b.seq);
428
429        // Write blocks in order and build index
430        let mut index = Vec::with_capacity(self.compressed_blocks.len());
431        let mut current_offset = 0u64;
432
433        for block in &self.compressed_blocks {
434            index.push(StoreBlockIndex {
435                first_doc_id: block.first_doc_id,
436                offset: current_offset,
437                length: u32::try_from(block.compressed.len()).map_err(|_| {
438                    io::Error::new(
439                        io::ErrorKind::InvalidData,
440                        "compressed store block too large",
441                    )
442                })?,
443                num_docs: block.num_docs,
444            });
445
446            self.writer.write_all(&block.compressed)?;
447            current_offset = current_offset
448                .checked_add(block.compressed.len() as u64)
449                .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "store size overflow"))?;
450        }
451
452        let data_end_offset = current_offset;
453
454        // Write dictionary if present
455        let dict_offset = if let Some(ref dict) = self.dict {
456            let offset = current_offset;
457            let dict_bytes = dict.as_bytes();
458            self.writer
459                .write_u32::<LittleEndian>(dict_bytes.len() as u32)?;
460            self.writer.write_all(dict_bytes)?;
461            Some(offset)
462        } else {
463            None
464        };
465
466        // Write index + footer
467        write_store_index_and_footer(
468            &mut self.writer,
469            &index,
470            data_end_offset,
471            dict_offset.unwrap_or(0),
472            self.next_doc_id,
473            self.dict.is_some(),
474        )?;
475
476        Ok(self.next_doc_id)
477    }
478}
479
480/// Block index entry for document store
481#[derive(Debug, Clone)]
482pub(crate) struct StoreBlockIndex {
483    pub(crate) first_doc_id: DocId,
484    pub(crate) offset: u64,
485    pub(crate) length: u32,
486    pub(crate) num_docs: u32,
487}
488
489/// Async document store reader - loads blocks on demand
490pub struct AsyncStoreReader {
491    /// FileHandle for the data portion - fetches ranges on demand
492    data_slice: FileHandle,
493    /// Block index
494    index: Vec<StoreBlockIndex>,
495    num_docs: u32,
496    /// Optional compression dictionary
497    dict: Option<CompressionDict>,
498    /// Process-wide byte-bounded block cache.
499    cache: Arc<SharedStoreCache>,
500    /// Stable directory + segment namespace for shared-cache keys.
501    cache_namespace: StoreCacheNamespace,
502}
503
504/// Decompressed block with pre-built doc offset table.
505///
506/// The offset table is built once on decompression: `offsets[i]` is the byte
507/// position in `data` where doc `i`'s length prefix starts. This turns the
508/// O(n) linear scan per `get()` into O(1) direct indexing.
509struct CachedBlock {
510    data: Vec<u8>,
511    /// Byte offset of each doc's length prefix within `data`.
512    /// `offsets.len()` == number of docs in the block.
513    offsets: Vec<u32>,
514}
515
516impl CachedBlock {
517    fn build(data: Vec<u8>, num_docs: u32) -> io::Result<Self> {
518        if num_docs as usize > data.len() / 4 {
519            return Err(io::Error::new(
520                io::ErrorKind::InvalidData,
521                "store block document count exceeds block length",
522            ));
523        }
524        let mut offsets = Vec::new();
525        offsets.try_reserve_exact(num_docs as usize).map_err(|_| {
526            io::Error::new(
527                io::ErrorKind::InvalidData,
528                "store block has too many documents",
529            )
530        })?;
531        let mut pos = 0usize;
532        for _ in 0..num_docs {
533            let length_end = pos.checked_add(4).ok_or_else(|| {
534                io::Error::new(io::ErrorKind::InvalidData, "store block offset overflow")
535            })?;
536            if length_end > data.len() {
537                return Err(io::Error::new(
538                    io::ErrorKind::InvalidData,
539                    "truncated block while building offset table",
540                ));
541            }
542            offsets.push(u32::try_from(pos).map_err(|_| {
543                io::Error::new(io::ErrorKind::InvalidData, "store block offset exceeds u32")
544            })?);
545            let doc_len =
546                u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]])
547                    as usize;
548            pos = length_end.checked_add(doc_len).ok_or_else(|| {
549                io::Error::new(io::ErrorKind::InvalidData, "store document length overflow")
550            })?;
551            if pos > data.len() {
552                return Err(io::Error::new(
553                    io::ErrorKind::UnexpectedEof,
554                    "store document is truncated",
555                ));
556            }
557        }
558        if pos != data.len() {
559            return Err(io::Error::new(
560                io::ErrorKind::InvalidData,
561                "store block contains trailing data",
562            ));
563        }
564        Ok(Self { data, offsets })
565    }
566
567    /// Get doc bytes by index within the block (O(1))
568    fn doc_bytes(&self, doc_offset_in_block: u32) -> io::Result<&[u8]> {
569        let idx = doc_offset_in_block as usize;
570        if idx >= self.offsets.len() {
571            return Err(io::Error::new(
572                io::ErrorKind::InvalidData,
573                "doc offset out of range",
574            ));
575        }
576        let start = self.offsets[idx] as usize;
577        let data_start = start.checked_add(4).ok_or_else(|| {
578            io::Error::new(io::ErrorKind::InvalidData, "store document offset overflow")
579        })?;
580        if data_start > self.data.len() {
581            return Err(io::Error::new(
582                io::ErrorKind::InvalidData,
583                "truncated doc length",
584            ));
585        }
586        let doc_len = u32::from_le_bytes([
587            self.data[start],
588            self.data[start + 1],
589            self.data[start + 2],
590            self.data[start + 3],
591        ]) as usize;
592        let data_end = data_start.checked_add(doc_len).ok_or_else(|| {
593            io::Error::new(io::ErrorKind::InvalidData, "store document length overflow")
594        })?;
595        if data_end > self.data.len() {
596            return Err(io::Error::new(
597                io::ErrorKind::InvalidData,
598                "doc data overflow",
599            ));
600        }
601        Ok(&self.data[data_start..data_end])
602    }
603
604    #[inline]
605    fn retained_bytes(&self) -> usize {
606        self.data.capacity() + self.offsets.capacity() * std::mem::size_of::<u32>()
607    }
608}
609
610#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
611struct StoreCacheNamespace {
612    directory: usize,
613    segment: u128,
614}
615
616#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
617struct StoreCacheKey {
618    namespace: StoreCacheNamespace,
619    first_doc_id: DocId,
620}
621
622struct SharedStoreCacheState {
623    blocks: LruCache<StoreCacheKey, Arc<CachedBlock>>,
624    retained_bytes: usize,
625    namespace_bytes: FxHashMap<StoreCacheNamespace, usize>,
626    namespace_readers: FxHashMap<StoreCacheNamespace, usize>,
627}
628
629/// Process-wide byte-bounded cache for decompressed document-store blocks.
630///
631/// The old cache bounded each segment by an entry count. One large stored
632/// body can make a block close to `MAX_STORE_BLOCK_BYTES`, so 32 entries per
633/// segment retained up to 2 GiB and multiplied that by segment fan-out. This
634/// cache has one hard byte ceiling across all indexes using the same policy.
635/// Hits take a shared read lock; eviction is insertion-ordered rather than
636/// serializing every hit solely for exact LRU promotion.
637pub(crate) struct SharedStoreCache {
638    state: RwLock<SharedStoreCacheState>,
639    max_bytes: usize,
640    /// Very large one-document blocks have almost no spatial reuse and can
641    /// evict thousands of ordinary result blocks. The OS compressed-file page
642    /// cache remains available when decompressed admission is bypassed.
643    max_entry_bytes: usize,
644}
645
646impl std::fmt::Debug for SharedStoreCache {
647    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
648        formatter
649            .debug_struct("SharedStoreCache")
650            .field("max_bytes", &self.max_bytes)
651            .field("max_entry_bytes", &self.max_entry_bytes)
652            .field("retained_bytes", &self.total_bytes())
653            .finish()
654    }
655}
656
657impl SharedStoreCache {
658    const MAX_ADMITTED_ENTRY_BYTES: usize = 8 * 1024 * 1024;
659
660    pub(crate) fn new(max_bytes: usize) -> Self {
661        Self::with_limits(max_bytes, max_bytes.min(Self::MAX_ADMITTED_ENTRY_BYTES))
662    }
663
664    fn with_limits(max_bytes: usize, max_entry_bytes: usize) -> Self {
665        Self {
666            state: RwLock::new(SharedStoreCacheState {
667                blocks: LruCache::unbounded(),
668                retained_bytes: 0,
669                namespace_bytes: FxHashMap::default(),
670                namespace_readers: FxHashMap::default(),
671            }),
672            max_bytes,
673            max_entry_bytes: max_entry_bytes.min(max_bytes),
674        }
675    }
676
677    fn register(&self, namespace: StoreCacheNamespace) {
678        if self.max_bytes == 0 {
679            return;
680        }
681        let mut state = self.state.write();
682        *state.namespace_readers.entry(namespace).or_default() += 1;
683    }
684
685    fn unregister(&self, namespace: StoreCacheNamespace) {
686        if self.max_bytes == 0 {
687            return;
688        }
689        let mut state = self.state.write();
690        let Some(readers) = state.namespace_readers.get_mut(&namespace) else {
691            return;
692        };
693        *readers -= 1;
694        if *readers > 0 {
695            return;
696        }
697        state.namespace_readers.remove(&namespace);
698
699        // A merged-away segment will never hit these entries again. Remove
700        // them immediately instead of waiting for unrelated searches to
701        // create enough pressure for ordinary LRU eviction.
702        let keys: Vec<_> = state
703            .blocks
704            .iter()
705            .filter_map(|(key, _)| (key.namespace == namespace).then_some(*key))
706            .collect();
707        for key in keys {
708            if let Some(block) = state.blocks.pop(&key) {
709                state.retained_bytes = state.retained_bytes.saturating_sub(block.retained_bytes());
710            }
711        }
712        state.namespace_bytes.remove(&namespace);
713    }
714
715    fn get(&self, key: StoreCacheKey) -> Option<Arc<CachedBlock>> {
716        // Do not serialize all process-wide hits merely to update exact LRU
717        // order. Concurrent readers use a shared lock; insertion and
718        // decompression-race resolution still promote entries.
719        self.state.read().blocks.peek(&key).map(Arc::clone)
720    }
721
722    /// Admit one block and return the canonical cached allocation if another
723    /// request won the decompression race.
724    fn insert(&self, key: StoreCacheKey, block: Arc<CachedBlock>) -> Arc<CachedBlock> {
725        let bytes = block.retained_bytes();
726        if self.max_bytes == 0 || bytes == 0 || bytes > self.max_entry_bytes {
727            return block;
728        }
729
730        let mut state = self.state.write();
731        if let Some(existing) = state.blocks.get(&key) {
732            return Arc::clone(existing);
733        }
734
735        state.retained_bytes = state.retained_bytes.saturating_add(bytes);
736        *state.namespace_bytes.entry(key.namespace).or_default() = state
737            .namespace_bytes
738            .get(&key.namespace)
739            .copied()
740            .unwrap_or(0)
741            .saturating_add(bytes);
742        state.blocks.put(key, Arc::clone(&block));
743
744        while state.retained_bytes > self.max_bytes {
745            let Some((evicted_key, evicted)) = state.blocks.pop_lru() else {
746                state.retained_bytes = 0;
747                state.namespace_bytes.clear();
748                break;
749            };
750            let evicted_bytes = evicted.retained_bytes();
751            state.retained_bytes = state.retained_bytes.saturating_sub(evicted_bytes);
752            if let Some(namespace_bytes) = state.namespace_bytes.get_mut(&evicted_key.namespace) {
753                *namespace_bytes = namespace_bytes.saturating_sub(evicted_bytes);
754                if *namespace_bytes == 0 {
755                    state.namespace_bytes.remove(&evicted_key.namespace);
756                }
757            }
758        }
759        block
760    }
761
762    pub(crate) fn total_bytes(&self) -> usize {
763        self.state.read().retained_bytes
764    }
765
766    pub(crate) fn total_blocks(&self) -> usize {
767        self.state.read().blocks.len()
768    }
769
770    fn namespace_bytes(&self, namespace: StoreCacheNamespace) -> usize {
771        self.state
772            .read()
773            .namespace_bytes
774            .get(&namespace)
775            .copied()
776            .unwrap_or(0)
777    }
778
779    fn namespace_blocks(&self, namespace: StoreCacheNamespace) -> usize {
780        self.state
781            .read()
782            .blocks
783            .iter()
784            .filter(|(key, _)| key.namespace == namespace)
785            .count()
786    }
787}
788
789impl Drop for AsyncStoreReader {
790    fn drop(&mut self) {
791        self.cache.unregister(self.cache_namespace);
792    }
793}
794
795impl AsyncStoreReader {
796    /// Open a document store from FileHandle
797    /// Only loads footer and index into memory, data blocks are fetched on-demand
798    pub(crate) async fn open(
799        file_handle: FileHandle,
800        directory_namespace: usize,
801        segment_namespace: u128,
802        cache: Arc<SharedStoreCache>,
803    ) -> io::Result<Self> {
804        let file_len = file_handle.len();
805        // Footer: data_end(8) + dict_offset(8) + num_docs(4) + has_dict(4) + version(4) + magic(4) = 32 bytes
806        if file_len < 32 {
807            return Err(io::Error::new(
808                io::ErrorKind::InvalidData,
809                "Store too small",
810            ));
811        }
812
813        // Read footer (32 bytes)
814        let footer = file_handle
815            .read_bytes_range(file_len - 32..file_len)
816            .await?;
817        let mut reader = footer.as_slice();
818        let data_end_offset = reader.read_u64::<LittleEndian>()?;
819        let dict_offset = reader.read_u64::<LittleEndian>()?;
820        let num_docs = reader.read_u32::<LittleEndian>()?;
821        let has_dict = reader.read_u32::<LittleEndian>()? != 0;
822        let version = reader.read_u32::<LittleEndian>()?;
823        let magic = reader.read_u32::<LittleEndian>()?;
824
825        if magic != STORE_MAGIC {
826            return Err(io::Error::new(
827                io::ErrorKind::InvalidData,
828                "Invalid store magic",
829            ));
830        }
831        if version != STORE_VERSION {
832            return Err(io::Error::new(
833                io::ErrorKind::InvalidData,
834                format!("Unsupported store version: {}", version),
835            ));
836        }
837
838        let index_end = file_len - 32;
839        if data_end_offset > index_end {
840            return Err(io::Error::new(
841                io::ErrorKind::InvalidData,
842                "store data section extends past its footer",
843            ));
844        }
845
846        // Load dictionary if present, and compute index_start in one pass
847        let (dict, index_start) = if has_dict {
848            if dict_offset < data_end_offset || dict_offset >= index_end {
849                return Err(io::Error::new(
850                    io::ErrorKind::InvalidData,
851                    "store dictionary offset is out of bounds",
852                ));
853            }
854            let dict_start = dict_offset;
855            let dict_header_end = dict_start.checked_add(4).ok_or_else(|| {
856                io::Error::new(
857                    io::ErrorKind::InvalidData,
858                    "store dictionary range overflow",
859                )
860            })?;
861            if dict_header_end > index_end {
862                return Err(io::Error::new(
863                    io::ErrorKind::UnexpectedEof,
864                    "store dictionary length is truncated",
865                ));
866            }
867            let dict_len_bytes = file_handle
868                .read_bytes_range(dict_start..dict_header_end)
869                .await?;
870            let dict_len = (&dict_len_bytes[..]).read_u32::<LittleEndian>()? as u64;
871            if dict_len > MAX_STORE_DICTIONARY_BYTES {
872                return Err(io::Error::new(
873                    io::ErrorKind::InvalidData,
874                    "store dictionary exceeds safety limit",
875                ));
876            }
877            let dict_end = dict_header_end.checked_add(dict_len).ok_or_else(|| {
878                io::Error::new(
879                    io::ErrorKind::InvalidData,
880                    "store dictionary range overflow",
881                )
882            })?;
883            if dict_end > index_end {
884                return Err(io::Error::new(
885                    io::ErrorKind::UnexpectedEof,
886                    "store dictionary is truncated",
887                ));
888            }
889            let dict_bytes = file_handle
890                .read_bytes_range(dict_header_end..dict_end)
891                .await?;
892            (
893                Some(CompressionDict::from_owned_bytes(dict_bytes)),
894                dict_end,
895            )
896        } else {
897            if dict_offset != 0 {
898                return Err(io::Error::new(
899                    io::ErrorKind::InvalidData,
900                    "store without a dictionary has a dictionary offset",
901                ));
902            }
903            (None, data_end_offset)
904        };
905
906        if index_start > index_end {
907            return Err(io::Error::new(
908                io::ErrorKind::InvalidData,
909                "store index offset is out of bounds",
910            ));
911        }
912
913        let index_bytes = file_handle.read_bytes_range(index_start..index_end).await?;
914        let mut reader = index_bytes.as_slice();
915
916        let num_blocks = reader.read_u32::<LittleEndian>()? as usize;
917        let required_index_bytes = num_blocks.checked_mul(20).ok_or_else(|| {
918            io::Error::new(io::ErrorKind::InvalidData, "store index size overflow")
919        })?;
920        if reader.len() != required_index_bytes {
921            return Err(io::Error::new(
922                io::ErrorKind::InvalidData,
923                "store index length is inconsistent",
924            ));
925        }
926        let mut index = Vec::new();
927        index
928            .try_reserve_exact(num_blocks)
929            .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "too many store blocks"))?;
930
931        let mut expected_doc = 0u32;
932        let mut expected_offset = 0u64;
933
934        for _ in 0..num_blocks {
935            let first_doc_id = reader.read_u32::<LittleEndian>()?;
936            let offset = reader.read_u64::<LittleEndian>()?;
937            let length = reader.read_u32::<LittleEndian>()?;
938            let num_docs_in_block = reader.read_u32::<LittleEndian>()?;
939
940            let end = offset.checked_add(length as u64).ok_or_else(|| {
941                io::Error::new(io::ErrorKind::InvalidData, "store block range overflow")
942            })?;
943            if first_doc_id != expected_doc
944                || num_docs_in_block == 0
945                || offset != expected_offset
946                || end > data_end_offset
947            {
948                return Err(io::Error::new(
949                    io::ErrorKind::InvalidData,
950                    "store block index is inconsistent",
951                ));
952            }
953            expected_doc = expected_doc.checked_add(num_docs_in_block).ok_or_else(|| {
954                io::Error::new(io::ErrorKind::InvalidData, "store document count overflow")
955            })?;
956            expected_offset = end;
957
958            index.push(StoreBlockIndex {
959                first_doc_id,
960                offset,
961                length,
962                num_docs: num_docs_in_block,
963            });
964        }
965
966        if expected_doc != num_docs || expected_offset != data_end_offset {
967            return Err(io::Error::new(
968                io::ErrorKind::InvalidData,
969                "store footer totals do not match its block index",
970            ));
971        }
972
973        // Create lazy slice for data portion only
974        let data_slice = file_handle.slice(0..data_end_offset);
975
976        let cache_namespace = StoreCacheNamespace {
977            directory: directory_namespace,
978            segment: segment_namespace,
979        };
980        cache.register(cache_namespace);
981        Ok(Self {
982            data_slice,
983            index,
984            num_docs,
985            dict,
986            cache,
987            cache_namespace,
988        })
989    }
990
991    /// Number of documents
992    pub fn num_docs(&self) -> u32 {
993        self.num_docs
994    }
995
996    /// Number of blocks currently in the cache
997    pub fn cached_blocks(&self) -> usize {
998        self.cache.namespace_blocks(self.cache_namespace)
999    }
1000
1001    /// Heap bytes retained by decompressed blocks and their document offsets.
1002    pub fn cached_bytes(&self) -> usize {
1003        self.cache.namespace_bytes(self.cache_namespace)
1004    }
1005
1006    /// Get a document by doc_id (async - may load block)
1007    pub async fn get(&self, doc_id: DocId, schema: &Schema) -> io::Result<Option<Document>> {
1008        if doc_id >= self.num_docs {
1009            return Ok(None);
1010        }
1011
1012        let t = crate::observe::Timer::start();
1013        let (entry, block) = self.find_and_load_block(doc_id).await?;
1014        let doc_bytes = block.doc_bytes(doc_id - entry.first_doc_id)?;
1015        let result = deserialize_document(doc_bytes, schema).map(Some);
1016        crate::observe::store_get(schema.index_label(), t.secs());
1017        result
1018    }
1019
1020    /// Get specific fields of a document by doc_id (async - may load block)
1021    ///
1022    /// Only deserializes the requested fields, skipping over unwanted data.
1023    /// Much faster than `get()` when documents have large fields (text bodies,
1024    /// vectors) that aren't needed for the response.
1025    pub async fn get_fields(
1026        &self,
1027        doc_id: DocId,
1028        schema: &Schema,
1029        field_ids: &[u32],
1030    ) -> io::Result<Option<Document>> {
1031        if doc_id >= self.num_docs {
1032            return Ok(None);
1033        }
1034
1035        let t = crate::observe::Timer::start();
1036        let (entry, block) = self.find_and_load_block(doc_id).await?;
1037        let doc_bytes = block.doc_bytes(doc_id - entry.first_doc_id)?;
1038        let result = deserialize_document_fields(doc_bytes, schema, field_ids).map(Some);
1039        crate::observe::store_get(schema.index_label(), t.secs());
1040        result
1041    }
1042
1043    /// Find the block index entry and load/cache the block for a given doc_id
1044    async fn find_and_load_block(
1045        &self,
1046        doc_id: DocId,
1047    ) -> io::Result<(&StoreBlockIndex, Arc<CachedBlock>)> {
1048        let block_idx = self
1049            .index
1050            .binary_search_by(|entry| {
1051                if doc_id < entry.first_doc_id {
1052                    std::cmp::Ordering::Greater
1053                } else if doc_id >= entry.first_doc_id + entry.num_docs {
1054                    std::cmp::Ordering::Less
1055                } else {
1056                    std::cmp::Ordering::Equal
1057                }
1058            })
1059            .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Doc not found in index"))?;
1060
1061        let entry = &self.index[block_idx];
1062        let block = self.load_block(entry).await?;
1063        Ok((entry, block))
1064    }
1065
1066    async fn load_block(&self, entry: &StoreBlockIndex) -> io::Result<Arc<CachedBlock>> {
1067        let key = StoreCacheKey {
1068            namespace: self.cache_namespace,
1069            first_doc_id: entry.first_doc_id,
1070        };
1071        if let Some(block) = self.cache.get(key) {
1072            return Ok(block);
1073        }
1074
1075        // Load from FileSlice
1076        let start = entry.offset;
1077        let end = start.checked_add(entry.length as u64).ok_or_else(|| {
1078            io::Error::new(io::ErrorKind::InvalidData, "store block range overflow")
1079        })?;
1080        let compressed = self.data_slice.read_bytes_range(start..end).await?;
1081
1082        // Use dictionary decompression if available
1083        let decompressed = if let Some(ref dict) = self.dict {
1084            crate::compression::decompress_with_dict_limited(
1085                compressed.as_slice(),
1086                dict,
1087                MAX_STORE_BLOCK_BYTES,
1088            )?
1089        } else {
1090            crate::compression::decompress_limited(compressed.as_slice(), MAX_STORE_BLOCK_BYTES)?
1091        };
1092
1093        // Build offset table for O(1) doc lookup within the block
1094        let cached = CachedBlock::build(decompressed, entry.num_docs)?;
1095        Ok(self.cache.insert(key, Arc::new(cached)))
1096    }
1097}
1098
1099/// Deserialize only specific fields from document bytes.
1100///
1101/// Skips over unwanted fields without allocating their values — just advances
1102/// the reader past their length-prefixed data. For large documents with many
1103/// fields (e.g., full text body), this avoids allocating/copying data that
1104/// the caller doesn't need.
1105pub fn deserialize_document_fields(
1106    data: &[u8],
1107    schema: &Schema,
1108    field_ids: &[u32],
1109) -> io::Result<Document> {
1110    deserialize_document_inner(data, schema, Some(field_ids))
1111}
1112
1113/// Deserialize all fields from document bytes.
1114///
1115/// Delegates to the shared field-parsing core with no field filter.
1116pub fn deserialize_document(data: &[u8], schema: &Schema) -> io::Result<Document> {
1117    deserialize_document_inner(data, schema, None)
1118}
1119
1120/// Shared deserialization core. `field_filter = None` means all fields wanted.
1121fn deserialize_document_inner(
1122    data: &[u8],
1123    _schema: &Schema,
1124    field_filter: Option<&[u32]>,
1125) -> io::Result<Document> {
1126    use crate::dsl::Field;
1127
1128    let mut reader = data;
1129    let num_fields = reader.read_u16::<LittleEndian>()? as usize;
1130    let mut doc = Document::new();
1131
1132    for _ in 0..num_fields {
1133        let field_id = reader.read_u16::<LittleEndian>()?;
1134        let type_tag = reader.read_u8()?;
1135
1136        let wanted = field_filter.is_none_or(|ids| ids.contains(&(field_id as u32)));
1137
1138        match type_tag {
1139            0 => {
1140                // Text
1141                let len = reader.read_u32::<LittleEndian>()? as usize;
1142                let bytes = take_document_bytes(&mut reader, len, "text field")?;
1143                if wanted {
1144                    let s = std::str::from_utf8(bytes)
1145                        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
1146                    doc.add_text(Field(field_id as u32), s);
1147                }
1148            }
1149            1 => {
1150                // U64
1151                let v = reader.read_u64::<LittleEndian>()?;
1152                if wanted {
1153                    doc.add_u64(Field(field_id as u32), v);
1154                }
1155            }
1156            2 => {
1157                // I64
1158                let v = reader.read_i64::<LittleEndian>()?;
1159                if wanted {
1160                    doc.add_i64(Field(field_id as u32), v);
1161                }
1162            }
1163            3 => {
1164                // F64
1165                let v = reader.read_f64::<LittleEndian>()?;
1166                if wanted {
1167                    doc.add_f64(Field(field_id as u32), v);
1168                }
1169            }
1170            4 => {
1171                // Bytes
1172                let len = reader.read_u32::<LittleEndian>()? as usize;
1173                let bytes = take_document_bytes(&mut reader, len, "byte field")?;
1174                if wanted {
1175                    doc.add_bytes(Field(field_id as u32), bytes.to_vec());
1176                }
1177            }
1178            5 => {
1179                // SparseVector
1180                let count = reader.read_u32::<LittleEndian>()? as usize;
1181                let byte_len = count.checked_mul(8).ok_or_else(|| {
1182                    io::Error::new(io::ErrorKind::InvalidData, "sparse vector size overflow")
1183                })?;
1184                let bytes = take_document_bytes(&mut reader, byte_len, "sparse vector")?;
1185                if wanted {
1186                    let mut entries = Vec::new();
1187                    entries.try_reserve_exact(count).map_err(|_| {
1188                        io::Error::new(io::ErrorKind::InvalidData, "sparse vector is too large")
1189                    })?;
1190                    let mut vector_reader = bytes;
1191                    for _ in 0..count {
1192                        let idx = vector_reader.read_u32::<LittleEndian>()?;
1193                        let val = vector_reader.read_f32::<LittleEndian>()?;
1194                        entries.push((idx, val));
1195                    }
1196                    doc.add_sparse_vector(Field(field_id as u32), entries);
1197                }
1198            }
1199            6 => {
1200                // DenseVector
1201                let count = reader.read_u32::<LittleEndian>()? as usize;
1202                let byte_len = count.checked_mul(4).ok_or_else(|| {
1203                    io::Error::new(io::ErrorKind::InvalidData, "dense vector size overflow")
1204                })?;
1205                let bytes = take_document_bytes(&mut reader, byte_len, "dense vector")?;
1206                if wanted {
1207                    let mut values = vec![0.0f32; count];
1208                    unsafe {
1209                        std::ptr::copy_nonoverlapping(
1210                            bytes.as_ptr(),
1211                            values.as_mut_ptr() as *mut u8,
1212                            byte_len,
1213                        );
1214                    }
1215                    doc.add_dense_vector(Field(field_id as u32), values);
1216                }
1217            }
1218            7 => {
1219                // Json
1220                let len = reader.read_u32::<LittleEndian>()? as usize;
1221                let bytes = take_document_bytes(&mut reader, len, "JSON field")?;
1222                if wanted {
1223                    let v: serde_json::Value = serde_json::from_slice(bytes)
1224                        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
1225                    doc.add_json(Field(field_id as u32), v);
1226                }
1227            }
1228            8 => {
1229                // BinaryDenseVector
1230                let len = reader.read_u32::<LittleEndian>()? as usize;
1231                let bytes = take_document_bytes(&mut reader, len, "binary dense vector")?;
1232                if wanted {
1233                    doc.add_binary_dense_vector(Field(field_id as u32), bytes.to_vec());
1234                }
1235            }
1236            _ => {
1237                return Err(io::Error::new(
1238                    io::ErrorKind::InvalidData,
1239                    format!("Unknown field type tag: {}", type_tag),
1240                ));
1241            }
1242        }
1243    }
1244
1245    Ok(doc)
1246}
1247
1248fn take_document_bytes<'a>(reader: &mut &'a [u8], len: usize, field: &str) -> io::Result<&'a [u8]> {
1249    if len > reader.len() {
1250        return Err(io::Error::new(
1251            io::ErrorKind::UnexpectedEof,
1252            format!("{field} is truncated"),
1253        ));
1254    }
1255    let (value, remaining) = reader.split_at(len);
1256    *reader = remaining;
1257    Ok(value)
1258}
1259
1260/// Raw block info for store merging (without decompression)
1261#[derive(Debug, Clone)]
1262pub struct RawStoreBlock {
1263    pub first_doc_id: DocId,
1264    pub num_docs: u32,
1265    pub offset: u64,
1266    pub length: u32,
1267}
1268
1269/// Store merger - concatenates compressed blocks from multiple stores without recompression
1270///
1271/// This is much faster than rebuilding stores since it avoids:
1272/// - Decompressing blocks from source stores
1273/// - Re-serializing documents
1274/// - Re-compressing blocks at level 22
1275///
1276/// Limitations:
1277/// - All source stores must NOT use dictionaries (or use the same dictionary)
1278/// - Doc IDs are remapped sequentially
1279pub struct StoreMerger<'a, W: Write> {
1280    writer: &'a mut W,
1281    index: Vec<StoreBlockIndex>,
1282    current_offset: u64,
1283    next_doc_id: DocId,
1284}
1285
1286impl<'a, W: Write> StoreMerger<'a, W> {
1287    pub fn new(writer: &'a mut W) -> Self {
1288        Self {
1289            writer,
1290            index: Vec::new(),
1291            current_offset: 0,
1292            next_doc_id: 0,
1293        }
1294    }
1295
1296    /// Append raw compressed blocks from a store file
1297    ///
1298    /// `data_slice` should be the data portion of the store (before index/footer)
1299    /// `blocks` contains the block metadata from the source store
1300    pub async fn append_store(
1301        &mut self,
1302        data_slice: &FileHandle,
1303        blocks: &[RawStoreBlock],
1304    ) -> io::Result<()> {
1305        for block in blocks {
1306            // Read raw compressed block data
1307            let start = block.offset;
1308            let end = start.checked_add(block.length as u64).ok_or_else(|| {
1309                io::Error::new(io::ErrorKind::InvalidData, "store block range overflow")
1310            })?;
1311            if end > data_slice.len() {
1312                return Err(io::Error::new(
1313                    io::ErrorKind::UnexpectedEof,
1314                    "store block range is out of bounds",
1315                ));
1316            }
1317            let compressed_data = data_slice.read_bytes_range(start..end).await?;
1318
1319            // Write to output
1320            self.writer.write_all(compressed_data.as_slice())?;
1321
1322            // Add to index with remapped doc IDs
1323            self.index.push(StoreBlockIndex {
1324                first_doc_id: self.next_doc_id,
1325                offset: self.current_offset,
1326                length: block.length,
1327                num_docs: block.num_docs,
1328            });
1329
1330            self.current_offset = self
1331                .current_offset
1332                .checked_add(block.length as u64)
1333                .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "store size overflow"))?;
1334            self.next_doc_id = self
1335                .next_doc_id
1336                .checked_add(block.num_docs)
1337                .ok_or_else(|| {
1338                    io::Error::new(io::ErrorKind::InvalidData, "store document count overflow")
1339                })?;
1340        }
1341
1342        Ok(())
1343    }
1344
1345    /// Append blocks from a dict-compressed store by decompressing and recompressing.
1346    ///
1347    /// For stores that use dictionary compression, raw blocks can't be stacked
1348    /// directly because the decompressor needs the original dictionary.
1349    /// This method decompresses each block with the source dict, then
1350    /// recompresses without a dictionary so the merged output is self-contained.
1351    pub async fn append_store_recompressing(&mut self, store: &AsyncStoreReader) -> io::Result<()> {
1352        let dict = store.dict();
1353        let data_slice = store.data_slice();
1354        let blocks = store.block_index();
1355
1356        for block in blocks {
1357            let start = block.offset;
1358            let end = start.checked_add(block.length as u64).ok_or_else(|| {
1359                io::Error::new(io::ErrorKind::InvalidData, "store block range overflow")
1360            })?;
1361            if end > data_slice.len() {
1362                return Err(io::Error::new(
1363                    io::ErrorKind::UnexpectedEof,
1364                    "store block range is out of bounds",
1365                ));
1366            }
1367            let compressed = data_slice.read_bytes_range(start..end).await?;
1368
1369            // Decompress with source dict (or without if no dict)
1370            let decompressed = if let Some(d) = dict {
1371                crate::compression::decompress_with_dict_limited(
1372                    compressed.as_slice(),
1373                    d,
1374                    MAX_STORE_BLOCK_BYTES,
1375                )?
1376            } else {
1377                crate::compression::decompress_limited(
1378                    compressed.as_slice(),
1379                    MAX_STORE_BLOCK_BYTES,
1380                )?
1381            };
1382
1383            // Recompress without dictionary
1384            let recompressed = crate::compression::compress(
1385                &decompressed,
1386                crate::compression::CompressionLevel::default(),
1387            )?;
1388
1389            self.writer.write_all(&recompressed)?;
1390
1391            self.index.push(StoreBlockIndex {
1392                first_doc_id: self.next_doc_id,
1393                offset: self.current_offset,
1394                length: u32::try_from(recompressed.len()).map_err(|_| {
1395                    io::Error::new(
1396                        io::ErrorKind::InvalidData,
1397                        "compressed store block too large",
1398                    )
1399                })?,
1400                num_docs: block.num_docs,
1401            });
1402
1403            self.current_offset = self
1404                .current_offset
1405                .checked_add(recompressed.len() as u64)
1406                .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "store size overflow"))?;
1407            self.next_doc_id = self
1408                .next_doc_id
1409                .checked_add(block.num_docs)
1410                .ok_or_else(|| {
1411                    io::Error::new(io::ErrorKind::InvalidData, "store document count overflow")
1412                })?;
1413        }
1414
1415        Ok(())
1416    }
1417
1418    /// Finish writing the merged store
1419    pub fn finish(self) -> io::Result<u32> {
1420        let data_end_offset = self.current_offset;
1421
1422        // No dictionary support for merged stores (would need same dict across all sources)
1423        let dict_offset = 0u64;
1424
1425        // Write index + footer
1426        write_store_index_and_footer(
1427            self.writer,
1428            &self.index,
1429            data_end_offset,
1430            dict_offset,
1431            self.next_doc_id,
1432            false,
1433        )?;
1434
1435        Ok(self.next_doc_id)
1436    }
1437}
1438
1439impl AsyncStoreReader {
1440    /// Get raw block metadata for merging (without loading block data)
1441    pub fn raw_blocks(&self) -> Vec<RawStoreBlock> {
1442        self.index
1443            .iter()
1444            .map(|entry| RawStoreBlock {
1445                first_doc_id: entry.first_doc_id,
1446                num_docs: entry.num_docs,
1447                offset: entry.offset,
1448                length: entry.length,
1449            })
1450            .collect()
1451    }
1452
1453    /// Get the data slice for raw block access
1454    pub fn data_slice(&self) -> &FileHandle {
1455        &self.data_slice
1456    }
1457
1458    /// Check if this store uses a dictionary (incompatible with raw merging)
1459    pub fn has_dict(&self) -> bool {
1460        self.dict.is_some()
1461    }
1462
1463    /// Get the decompression dictionary (if any)
1464    pub fn dict(&self) -> Option<&CompressionDict> {
1465        self.dict.as_ref()
1466    }
1467
1468    /// Get block index for iteration
1469    pub(crate) fn block_index(&self) -> &[StoreBlockIndex] {
1470        &self.index
1471    }
1472}
1473
1474#[cfg(test)]
1475mod tests {
1476    use super::*;
1477
1478    fn cached_test_block(byte: u8) -> Arc<CachedBlock> {
1479        Arc::new(CachedBlock::build(vec![4, 0, 0, 0, byte, byte, byte, byte], 1).unwrap())
1480    }
1481
1482    #[test]
1483    fn cached_block_rejects_truncated_and_trailing_documents() {
1484        assert!(CachedBlock::build(vec![8, 0, 0, 0, 1], 1).is_err());
1485        assert!(CachedBlock::build(vec![0, 0, 0, 0, 1], 1).is_err());
1486    }
1487
1488    #[test]
1489    fn document_deserializer_rejects_length_prefixed_slice_overrun() {
1490        let schema = Schema::builder().build();
1491        let truncated_text = [1, 0, 0, 0, 0, 5, 0, 0, 0, b'x'];
1492        assert!(deserialize_document(&truncated_text, &schema).is_err());
1493
1494        let truncated_sparse = [1, 0, 0, 0, 5, 2, 0, 0, 0, 1, 0, 0, 0];
1495        assert!(deserialize_document(&truncated_sparse, &schema).is_err());
1496    }
1497
1498    #[test]
1499    fn shared_store_cache_is_byte_bounded_and_read_concurrent() {
1500        let block_bytes = cached_test_block(1).retained_bytes();
1501        let cache = SharedStoreCache::with_limits(block_bytes * 2, block_bytes);
1502        let key = |first_doc_id| StoreCacheKey {
1503            namespace: StoreCacheNamespace {
1504                directory: 1,
1505                segment: 7,
1506            },
1507            first_doc_id,
1508        };
1509
1510        cache.insert(key(1), cached_test_block(1));
1511        cache.insert(key(2), cached_test_block(2));
1512        assert!(cache.get(key(1)).is_some());
1513        cache.insert(key(3), cached_test_block(3));
1514
1515        assert!(cache.get(key(1)).is_none());
1516        assert!(cache.get(key(2)).is_some());
1517        assert!(cache.get(key(3)).is_some());
1518        assert!(cache.total_bytes() <= block_bytes * 2);
1519    }
1520
1521    #[test]
1522    fn shared_store_cache_bypasses_oversized_entries() {
1523        let block = cached_test_block(1);
1524        let cache = SharedStoreCache::with_limits(1024, block.retained_bytes() - 1);
1525        let key = StoreCacheKey {
1526            namespace: StoreCacheNamespace {
1527                directory: 1,
1528                segment: 9,
1529            },
1530            first_doc_id: 0,
1531        };
1532        cache.insert(key, block);
1533        assert_eq!(cache.total_bytes(), 0);
1534        assert!(cache.get(key).is_none());
1535    }
1536
1537    #[test]
1538    fn shared_store_cache_purges_closed_segment_namespace() {
1539        let block = cached_test_block(1);
1540        let cache = SharedStoreCache::with_limits(1024, 1024);
1541        let key = StoreCacheKey {
1542            namespace: StoreCacheNamespace {
1543                directory: 1,
1544                segment: 11,
1545            },
1546            first_doc_id: 0,
1547        };
1548        cache.register(key.namespace);
1549        cache.insert(key, block);
1550        assert!(cache.total_bytes() > 0);
1551        cache.unregister(key.namespace);
1552        assert_eq!(cache.total_bytes(), 0);
1553        assert!(cache.get(key).is_none());
1554    }
1555
1556    #[test]
1557    fn shared_store_cache_isolates_equal_segment_ids_across_directories() {
1558        let cache = SharedStoreCache::with_limits(1024, 1024);
1559        let key = |directory| StoreCacheKey {
1560            namespace: StoreCacheNamespace {
1561                directory,
1562                segment: 42,
1563            },
1564            first_doc_id: 0,
1565        };
1566        let left = cache.insert(key(1), cached_test_block(1));
1567        let right = cache.insert(key(2), cached_test_block(2));
1568
1569        assert!(!Arc::ptr_eq(&left, &right));
1570        assert_eq!(cache.total_blocks(), 2);
1571    }
1572}