diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 3434376..fc0204e 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -89,7 +89,11 @@ Key properties: - **Bounded memory.** Documents are buffered up to a configurable cap and flushed to a new segment file before the cap is reached, so indexing a large batch does not require - holding the whole inverted index in RAM. + holding the whole inverted index in RAM. `optimize()` is bounded too: it merges as a + streaming, k-way pass (see below), keeping only a per-term working set plus small + O(document-count) bookkeeping resident, so its peak heap does not scale with the corpus's + total text volume. (One caveat: a single extremely frequent term still materializes its + own full posting list transiently during the merge.) - **Segments are invisible until commit.** A flushed segment file exists on disk but is not referenced by any generation record until `commit()` writes the new `segments_N` listing it. If the process dies mid-batch, the on-disk generation still @@ -102,7 +106,13 @@ Key properties: Recomputing anything from the original text (re-tokenizing, re-scoring norms) would reintroduce drift the merge is specifically designed to avoid; the merged segment must be indistinguishable, byte-for-byte in content, from what the legacy engine's own - merge would have produced. + merge would have produced. The merge runs as a single streaming pass: a k-way merge over + the input segments' term cursors emits the merged term dictionary in order while the large + postings/positions/stored blocks are streamed to temp files + (`.fdt.tmp`/`.frq.tmp`/`.prx.tmp` in the index dir) and assembled into the final + compound file — so the whole merged index is never held in RAM. The older in-RAM merge + (`merge_segments`) is retained only as the differential-test oracle that pins this + streaming path byte-for-byte. - **Ordered, crash-safe durability.** `durability.rs` writes new segment/generation data and only then atomically flips the pointer that makes it visible (`segments.gen`), fsyncing before the flip on the merge path; unreferenced leftovers from an interrupted diff --git a/README.md b/README.md index 4fc5c6d..ffef394 100644 --- a/README.md +++ b/README.md @@ -99,6 +99,10 @@ Order-of-magnitude numbers from local benchmarking against the legacy PHP engine - **Indexing throughput:** the native streaming writer indexed the same document batches **~95–150× faster** than ZSL's PHP writer, with peak RSS bounded by a configurable buffered-docs cap rather than growing with the batch size. +- **Optimize (merge) memory:** `optimize()` uses a streaming, k-way merge whose peak heap is + bounded by a per-term working set plus small per-document bookkeeping, *not* by the corpus + size. On the same ~135k-document index this cut optimize's peak heap from **~3.2 GB to + ~0.1 GB** (and it ran faster), so compacting a large index no longer risks an OOM. These are single-machine, order-of-magnitude measurements, not precise benchmarks — treat them as a shape-of-the-win indicator, not a guarantee. diff --git a/docs/API.md b/docs/API.md index a98f80f..a137423 100644 --- a/docs/API.md +++ b/docs/API.md @@ -52,6 +52,20 @@ echo sdsearch_version(); // "0.1.0" — also a smoke test that the extension loa `commit()` and `optimize()` consume the writer: after either, the object is closed and any further method throws "writer not open". Create a fresh `Writer` for the next batch. +**When to call `optimize()`.** Each `commit()` adds one or more segments; the segment count +only shrinks when you `optimize()` (there is no automatic merge policy). A read (search) opens +every live segment, so an index that is committed many times without optimizing gets slower +and heavier to open. The recommended pattern for a bulk feed is **open once → `add_document` +per doc → `optimize()` once at the end** (rather than open/commit per document), which keeps +the index compacted to a single segment. + +**`optimize()` resource profile.** The merge is streaming and bounded-memory: peak heap is a +per-term working set plus small per-document bookkeeping, independent of the corpus text +volume (on a ~135k-doc index, ~0.1 GB peak heap). While it runs it writes temporary files +(`.fdt.tmp` / `.frq.tmp` / `.prx.tmp`) into the index directory sized in aggregate +close to the final segment, so provision disk accordingly. Stale generation manifests +(`segments_N`) are pruned automatically after each commit/optimize. + ## Indexing (write path) ```php diff --git a/sdsearch-core/examples/optimize_bench.rs b/sdsearch-core/examples/optimize_bench.rs new file mode 100644 index 0000000..8924f3c --- /dev/null +++ b/sdsearch-core/examples/optimize_bench.rs @@ -0,0 +1,263 @@ +//! OPTIMIZE baseline: measures the RAM and time cost of `IndexWriter::optimize()` (the +//! full-index merge) as a function of index size, to evaluate a future streaming merge. +//! +//! Two memory numbers are reported, because they answer different questions: +//! - `heap_peak_kb`: peak HEAP bytes attributed to the optimize alone (a tracking global +//! allocator; the counter is reset to steady-state right before optimize). This is the +//! number a streaming merge is meant to bound — it ignores the mmap of the source segments. +//! - `vmhwm_kb`: process peak RSS (VmHWM). Includes the mmap of the source segments, which a +//! streaming merge does NOT reduce. Reported for context, not as the target metric. +//! +//! Flow per run: copy the KB fixture as base → stream N synthetic docs with `cap` (→ a +//! multi-segment index) → reset the heap counter → optimize() → report. +//! +//! Usage: +//! cargo run -p sdsearch-core --release --example optimize_bench -- [cap] +//! N number of docs to add on top of the base (default 2000) +//! cap max_buffered_docs while building (default 1000 → ceil(N/cap) flushed segments) +//! +//! Prints ONE JSON line: {"n":..,"cap":..,"segments_before":..,"doc_count":.., +//! "optimize_ms":..,"heap_peak_kb":..,"vmhwm_kb":..} + +use sdsearch_core::zsl::writer::{FieldKind, IndexWriter, WriterDoc, WriterField, WriterOpts}; +use std::alloc::{GlobalAlloc, Layout, System}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicUsize, Ordering}; + +// ---- tracking allocator: measures live + peak heap bytes ---- +static LIVE: AtomicUsize = AtomicUsize::new(0); +static PEAK: AtomicUsize = AtomicUsize::new(0); + +struct TrackingAlloc; + +unsafe impl GlobalAlloc for TrackingAlloc { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let p = System.alloc(layout); + if !p.is_null() { + let live = LIVE.fetch_add(layout.size(), Ordering::Relaxed) + layout.size(); + PEAK.fetch_max(live, Ordering::Relaxed); + } + p + } + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + System.dealloc(ptr, layout); + LIVE.fetch_sub(layout.size(), Ordering::Relaxed); + } + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let p = System.realloc(ptr, layout, new_size); + if !p.is_null() { + let old = layout.size(); + if new_size >= old { + let live = LIVE.fetch_add(new_size - old, Ordering::Relaxed) + (new_size - old); + PEAK.fetch_max(live, Ordering::Relaxed); + } else { + LIVE.fetch_sub(old - new_size, Ordering::Relaxed); + } + } + p + } +} + +#[global_allocator] +static GLOBAL: TrackingAlloc = TrackingAlloc; + +/// Anchors PEAK to the current live bytes, so the next PEAK read reflects only what is +/// allocated from here on (i.e. the optimize's own heap high-water mark above steady state). +fn reset_peak() -> usize { + let live = LIVE.load(Ordering::Relaxed); + PEAK.store(live, Ordering::Relaxed); + live +} + +/// process peak RSS (VmHWM from /proc/self/status), in KB; 0 if unavailable. +fn vmhwm_kb() -> u64 { + std::fs::read_to_string("/proc/self/status") + .ok() + .and_then(|s| { + s.lines() + .find(|l| l.starts_with("VmHWM:")) + .and_then(|l| l.split_whitespace().nth(1)) + .and_then(|kb| kb.parse().ok()) + }) + .unwrap_or(0) +} + +/// copies the committed KB fixture to a fresh temp dir (skips locks and `.sti`), returns it. +fn copy_kb_base(n: usize, cap: usize) -> PathBuf { + let src = PathBuf::from(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/zsl_index_kb" + )); + let dst = std::env::temp_dir().join(format!( + "sdsearch_optbench_{}_{}_{}", + std::process::id(), + n, + cap + )); + if dst.is_dir() { + std::fs::remove_dir_all(&dst).ok(); + } + std::fs::create_dir_all(&dst).expect("create temp dir"); + for entry in std::fs::read_dir(&src).expect("read KB fixture") { + let p = entry.unwrap().path(); + let name = p.file_name().unwrap().to_string_lossy().to_string(); + if name.contains("lock") || name.ends_with(".sti") { + continue; + } + std::fs::copy(&p, dst.join(&name)).expect("copy fixture file"); + } + dst +} + +/// N deterministic docs: title ~5 tokens, body ~40 tokens, id keyword (== perf_writer.php). +fn gen_docs(n: usize) -> Vec { + const POOL: &[&str] = &[ + "printer", "network", "vpn", "login", "email", "server", "crash", "slow", "reset", + "password", "access", "error", "update", "install", "config", "backup", "restore", + "timeout", "license", "upgrade", "firewall", "router", "disk", "memory", "cpu", + ]; + let np = POOL.len(); + (0..n) + .map(|i| { + let title = format!( + "ticket {i} {} {} issue{i}", + POOL[i % np], + POOL[(i * 3) % np] + ); + let body: String = (0..40) + .map(|j| POOL[(i * 7 + j * 5) % np]) + .collect::>() + .join(" "); + let body = format!("{body} ref{i}"); + WriterDoc { + fields: vec![ + WriterField { + name: "title".into(), + value: title, + kind: FieldKind::Text, + stored: true, + }, + WriterField { + name: "body".into(), + value: body, + kind: FieldKind::Text, + stored: true, + }, + WriterField { + name: "id".into(), + value: format!("REC-{i}"), + kind: FieldKind::Keyword, + stored: true, + }, + ], + } + }) + .collect() +} + +/// copies an arbitrary index dir to `dst` (skips lock files and `.sti`). Overwrites `dst`. +fn copy_index(src: &Path, dst: &Path) { + if dst.is_dir() { + std::fs::remove_dir_all(dst).ok(); + } + std::fs::create_dir_all(dst).expect("create scratch dir"); + for entry in std::fs::read_dir(src).expect("read source index") { + let p = entry.unwrap().path(); + let name = p.file_name().unwrap().to_string_lossy().to_string(); + if name.contains("lock") || name.ends_with(".sti") { + continue; + } + std::fs::copy(&p, dst.join(&name)).expect("copy index file"); + } +} + +/// `optimize_bench existing `: copies a REAL index to a +/// disk-backed scratch, adds `n_extra` tiny docs (cap=1 → one segment each) to force a +/// multi-segment optimize, then measures optimize() over the whole corpus. +fn run_existing(args: &[String]) { + let src = Path::new( + args.get(2) + .expect("usage: optimize_bench existing "), + ); + let n_extra: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(2); + let scratch = Path::new( + args.get(4) + .expect("scratch dir required (use a disk-backed path)"), + ); + + copy_index(src, scratch); + + { + let opts = WriterOpts { + max_buffered_docs: 1, + ..WriterOpts::default() + }; + let mut w = IndexWriter::open(scratch, opts).expect("open (build) failed"); + for d in gen_docs(n_extra) { + w.add_document(d).expect("add_document failed"); + } + w.commit().expect("commit failed"); + } + + let segments_before = sdsearch_core::zsl::segments::read_segment_infos(scratch) + .expect("read segment infos") + .len(); + + reset_peak(); + let t0 = std::time::Instant::now(); + let w = IndexWriter::open(scratch, WriterOpts::default()).expect("open (optimize) failed"); + let report = w.optimize().expect("optimize failed"); + let optimize_ms = t0.elapsed().as_secs_f64() * 1000.0; + let heap_peak_kb = PEAK.load(Ordering::Relaxed) as u64 / 1024; + + println!( + "{{\"source\":\"{}\",\"n_extra\":{},\"segments_before\":{},\"doc_count\":{},\"optimize_ms\":{:.2},\"heap_peak_kb\":{},\"vmhwm_kb\":{}}}", + src.display(), n_extra, segments_before, report.doc_count, optimize_ms, heap_peak_kb, vmhwm_kb() + ); + + std::fs::remove_dir_all(scratch).ok(); +} + +fn main() { + let args: Vec = std::env::args().collect(); + if args.get(1).map(String::as_str) == Some("existing") { + run_existing(&args); + return; + } + let n: usize = args.get(1).and_then(|s| s.parse().ok()).unwrap_or(2000); + let cap: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(1000); + + let dir = copy_kb_base(n, cap); + + // ---- build: stream N docs into a multi-segment index (bounded-memory add path) ---- + { + let opts = WriterOpts { + max_buffered_docs: cap, + ..WriterOpts::default() + }; + let mut w = IndexWriter::open(&dir, opts).expect("open (build) failed"); + for d in gen_docs(n) { + w.add_document(d).expect("add_document failed"); + } + w.commit().expect("commit failed"); + } + + let segments_before = sdsearch_core::zsl::segments::read_segment_infos(&dir) + .expect("read segment infos") + .len(); + + // ---- measure: optimize() only ---- + reset_peak(); + let t0 = std::time::Instant::now(); + let w = IndexWriter::open(&dir, WriterOpts::default()).expect("open (optimize) failed"); + let report = w.optimize().expect("optimize failed"); + let optimize_ms = t0.elapsed().as_secs_f64() * 1000.0; + let heap_peak_kb = PEAK.load(Ordering::Relaxed) as u64 / 1024; + + println!( + "{{\"n\":{},\"cap\":{},\"segments_before\":{},\"doc_count\":{},\"optimize_ms\":{:.2},\"heap_peak_kb\":{},\"vmhwm_kb\":{}}}", + n, cap, segments_before, report.doc_count, optimize_ms, heap_peak_kb, vmhwm_kb() + ); + + std::fs::remove_dir_all(&dir).ok(); +} diff --git a/sdsearch-core/src/zsl/segment.rs b/sdsearch-core/src/zsl/segment.rs index 81be422..96a1fa9 100644 --- a/sdsearch-core/src/zsl/segment.rs +++ b/sdsearch-core/src/zsl/segment.rs @@ -6,7 +6,7 @@ use crate::zsl::fields::{read_field_infos, FieldInfo}; use crate::zsl::norms::{approx_field_len, read_norms}; use crate::zsl::postings::{read_all_positions, read_freqs, read_positions}; use crate::zsl::stored::{read_stored_fields, read_stored_raw, StoredRaw}; -use crate::zsl::terms::TermDict; +use crate::zsl::terms::{TermCursor, TermDict}; use std::collections::HashMap; use std::path::Path; @@ -58,6 +58,14 @@ impl ZslSegment { self.dict.iter_terms() } + /// lazy cursor over every `(field, term)` pair in ZSL canonical order + /// (field names ascending, terms ascending within each field), without + /// materializing a `Vec` of all terms like `all_terms` does. Used by the + /// bounded-memory k-way streaming merge. + pub fn term_cursor(&self) -> TermCursor<'_> { + self.dict.cursor() + } + /// stored fields of a doc in write order (LOCAL field_num + tokenized flag). pub fn stored_raw(&self, local_doc: usize) -> std::io::Result> { read_stored_raw( @@ -329,6 +337,40 @@ mod tests { let _ = std::fs::remove_dir_all(&dir); } + fn seg_kb() -> ZslSegment { + let dir = PathBuf::from(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/zsl_index_kb" + )); + // KB has a single segment "_2" with no deletes (del_gen -1); it spans + // several stored/indexed fields, unlike the tiny "zsl_index" fixture. + ZslSegment::open_named(&dir, "_2", -1).unwrap() + } + + #[test] + fn term_cursor_yields_all_terms_in_zsl_canonical_order() { + let s = seg_kb(); + let mut expected = s.all_terms(); + expected.sort(); + + // sanity: the fixture must actually exercise field-name ordering, not + // just within-field term ordering. + let distinct_fields: std::collections::HashSet<&String> = + expected.iter().map(|(f, _)| f).collect(); + assert!( + distinct_fields.len() >= 2, + "fixture has too few fields to exercise field ordering: {distinct_fields:?}" + ); + + let mut got: Vec<(String, String)> = Vec::new(); + let mut cur = s.term_cursor(); + while let Some((field, term)) = cur.peek() { + got.push((field.to_string(), term.to_string())); + cur.advance(); + } + assert_eq!(got, expected); + } + #[test] fn merge_accessors_expose_fields_deletes_norms_terms_stored() { let s = seg(); // incidents fixture, 4 docs, no deletes diff --git a/sdsearch-core/src/zsl/terms.rs b/sdsearch-core/src/zsl/terms.rs index 1b97bbb..acb3377 100644 --- a/sdsearch-core/src/zsl/terms.rs +++ b/sdsearch-core/src/zsl/terms.rs @@ -30,6 +30,15 @@ impl FieldTerms { fn term(&self, i: usize) -> &[u8] { &self.text[self.offsets[i] as usize..self.offsets[i + 1] as usize] } + /// zero-copy `&str` view of term `i`. `text` is built exclusively from + /// `String`s (`read_modified_utf8` decodes into a real `String`, never raw + /// bytes) and `offsets` only ever land on whole-term boundaries, so every + /// slice is guaranteed valid UTF-8 — no lossy/alloc path needed here, + /// unlike `iter_terms`'s owned `String::from_utf8_lossy`. + fn term_str(&self, i: usize) -> &str { + std::str::from_utf8(self.term(i)) + .expect("FieldTerms invariant violated: term bytes are valid UTF-8") + } } pub struct TermDict { @@ -138,6 +147,14 @@ impl TermDict { out } + /// lazy cursor over every `(field, term)` pair in ZSL canonical order + /// (`fieldName · \0 · text`, i.e. field names ascending, terms ascending + /// within each field). Unlike `iter_terms`, never materializes a `Vec` of + /// all terms up front — the k-way streaming merge walks this instead. + pub fn cursor(&self) -> TermCursor<'_> { + TermCursor::new(self) + } + /// Enumerates ALL terms as `(field, text)`. Order not guaranteed (grouped by /// field, each field ascending). Used by the merge to walk each source segment's /// terms and copy their postings via `positions_all(field, text)`. @@ -155,6 +172,60 @@ impl TermDict { } } +/// Lazy cursor over all `(field, term)` pairs of a `TermDict`, in ZSL +/// canonical order (field names sorted ascending, terms ascending within each +/// field — equivalent to sorting `(fieldName, text)` tuples since `\0` is the +/// minimum byte and never appears inside a field name or term). Backed +/// directly by each field's already-sorted `FieldTerms` buffer: no `Vec` of +/// all terms is built up front, so memory stays bounded regardless of how +/// many terms the segment holds. +pub struct TermCursor<'a> { + dict: &'a TermDict, + fields: Vec<&'a str>, + field_idx: usize, + term_idx: usize, +} + +impl<'a> TermCursor<'a> { + fn new(dict: &'a TermDict) -> TermCursor<'a> { + let mut fields: Vec<&str> = dict.by_field.keys().map(String::as_str).collect(); + fields.sort_unstable(); + TermCursor { + dict, + fields, + field_idx: 0, + term_idx: 0, + } + } + + /// current `(field, term)` pair, or `None` once every field is exhausted. + pub fn peek(&self) -> Option<(&str, &str)> { + let field = *self.fields.get(self.field_idx)?; + let ft = self + .dict + .by_field + .get(field) + .expect("field came from this dict's own key set"); + Some((field, ft.term_str(self.term_idx))) + } + + /// moves to the next pair in canonical order. No-op once exhausted. + pub fn advance(&mut self) { + if self.field_idx >= self.fields.len() { + return; + } + self.term_idx += 1; + // `by_field` only ever holds fields with >=1 term (a field is inserted + // lazily, on its first term, while reading `.tis`), so the next field + // (if any) is guaranteed non-empty — no need to loop-skip empties. + let current_len = self.dict.by_field[self.fields[self.field_idx]].len(); + if self.term_idx >= current_len { + self.field_idx += 1; + self.term_idx = 0; + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/sdsearch-core/src/zsl/writer/cfs.rs b/sdsearch-core/src/zsl/writer/cfs.rs index c312891..a2e4a12 100644 --- a/sdsearch-core/src/zsl/writer/cfs.rs +++ b/sdsearch-core/src/zsl/writer/cfs.rs @@ -5,6 +5,8 @@ //! `.fdx .fdt .fnm .nrm .tis .tii .frq .prx` (stored first, then fnm/nrm, then dict). use crate::zsl::bytes::{write_i64_be, write_modified_utf8, write_vint}; +use std::io::{self, Write}; +use std::path::Path; /// Packs the `(ext, data)` sub-files into a `.cfs`. `segment_name` is the prefix /// (e.g. `"_0"`); the name in the directory is `segment_name + ext`. @@ -30,6 +32,99 @@ pub fn write_cfs(segment_name: &str, files: &[(&str, &[u8])]) -> Vec { out } +/// Source for one CFS sub-file fed to `write_cfs_streaming`: either an in-memory buffer or an +/// on-disk temp file. `Path` sources are streamed with a fixed-size buffer (`std::io::copy`), +/// never loaded fully into RAM. +pub enum CfsSource<'a> { + Mem(&'a [u8]), + Path(&'a Path), +} + +impl<'a> CfsSource<'a> { + fn len(&self) -> io::Result { + match self { + CfsSource::Mem(data) => Ok(data.len() as u64), + CfsSource::Path(path) => Ok(std::fs::metadata(path)?.len()), + } + } + + /// Writes this source's bytes to `out`. `expected_len` must be the exact value that `len()` + /// returned for this same source and that was used to compute its CFS directory offset. If + /// the number of bytes actually copied differs (e.g. the underlying file was truncated or + /// replaced between the `len()` call and this write — a TOCTOU race), this returns an + /// `io::Error` instead of silently emitting a `.cfs` whose directory offsets don't match the + /// data actually written. + fn write_to(&self, out: &mut W, expected_len: u64) -> io::Result<()> { + match self { + CfsSource::Mem(data) => out.write_all(data), + CfsSource::Path(path) => { + let mut f = std::fs::File::open(path)?; + let copied = io::copy(&mut f, out)?; + if copied != expected_len { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!( + "cfs source {}: expected {expected_len} bytes (length used to \ + compute the directory offset) but copied {copied}; file changed \ + size between measurement and write", + path.display() + ), + )); + } + Ok(()) + } + } + } +} + +/// Streaming counterpart to `write_cfs`: identical byte layout (`VInt(fileCount)` + directory +/// `[Long(offset) + String(fullName)]` + data blocks). Every source length is known up front +/// (`Mem`: slice len; `Path`: file metadata), so all directory offsets are computed directly and +/// the directory is written to `out` once — no back-patch into `out` is needed (the small +/// in-RAM directory buffer is patched before anything is written). Each data block is then +/// streamed into `out` in turn; `Path` sources are copied via `std::io::copy` and never loaded +/// fully into memory. Each `Path` source's length is measured once and that same value is used +/// both for its directory offset and to verify the number of bytes actually copied, so a source +/// that changes size between the two (e.g. a truncated temp file) is caught as an error rather +/// than silently producing a `.cfs` with offsets that don't match its data. +/// +/// On `Err`, `out` may already contain a partial write (directory and/or some data blocks +/// written before the failing source was reached); callers must discard it and must not treat it +/// as a valid, committable `.cfs`. +pub fn write_cfs_streaming( + out: &mut W, + segment_name: &str, + files: &[(&str, CfsSource)], +) -> io::Result<()> { + // Build the directory in a small in-RAM buffer first: its length depends only on the file + // count and names (not on the — possibly huge — data), so once built we know every data + // block's final offset and can patch them here, still in RAM, before writing anything to + // `out`. + let mut dir = Vec::new(); + write_vint(&mut dir, files.len() as u64); + let mut ptr_positions = Vec::with_capacity(files.len()); + for (ext, _) in files { + ptr_positions.push(dir.len()); + write_i64_be(&mut dir, 0); // placeholder, patched below (in RAM, not yet in `out`) + write_modified_utf8(&mut dir, &format!("{segment_name}{ext}")); + } + + let mut offset = dir.len() as u64; + let mut lengths = Vec::with_capacity(files.len()); + for (i, (_, src)) in files.iter().enumerate() { + let len = src.len()?; + dir[ptr_positions[i]..ptr_positions[i] + 8].copy_from_slice(&offset.to_be_bytes()); + offset += len; + lengths.push(len); + } + + out.write_all(&dir)?; + for ((_, src), len) in files.iter().zip(lengths.iter()) { + src.write_to(out, *len)?; + } + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -43,6 +138,101 @@ mod tests { std::env::temp_dir().join(format!("sdsearch_cfs_{}_{}.cfs", std::process::id(), n)) } + fn temp_path_named(tag: &str) -> std::path::PathBuf { + let n = COUNTER.fetch_add(1, Ordering::Relaxed); + std::env::temp_dir().join(format!( + "sdsearch_cfs_{}_{}_{}.tmp", + std::process::id(), + n, + tag + )) + } + + #[test] + fn write_cfs_streaming_matches_write_cfs_for_mem_sources() { + let fnm = b"field-infos-bytes".to_vec(); + let tis = b"term-dict-bytes-longer".to_vec(); + let frq = vec![0u8, 1, 2, 3, 4]; + let files: Vec<(&str, &[u8])> = vec![(".fnm", &fnm), (".tis", &tis), (".frq", &frq)]; + let expected = write_cfs("_7", &files); + + let streaming_files: Vec<(&str, CfsSource)> = vec![ + (".fnm", CfsSource::Mem(&fnm)), + (".tis", CfsSource::Mem(&tis)), + (".frq", CfsSource::Mem(&frq)), + ]; + let mut actual = Vec::new(); + write_cfs_streaming(&mut actual, "_7", &streaming_files).unwrap(); + + assert_eq!(actual, expected); + } + + #[test] + fn write_cfs_streaming_matches_write_cfs_with_mixed_path_and_mem_sources() { + let fnm = b"field-infos-bytes".to_vec(); + let tis = b"term-dict-bytes-longer-for-the-path-source-case".to_vec(); + let frq = vec![0u8, 1, 2, 3, 4]; + + // All-Mem oracle: same sub-files, same order, produced by the existing `write_cfs`. + let files: Vec<(&str, &[u8])> = vec![(".fnm", &fnm), (".tis", &tis), (".frq", &frq)]; + let expected = write_cfs("_9", &files); + + // Write the ".tis" data to a temp file and pass it as a `Path` source; the rest stay + // `Mem`. Output must still be byte-identical to the all-Mem oracle above. + let tis_path = temp_path_named("tis"); + std::fs::write(&tis_path, &tis).unwrap(); + + let streaming_files: Vec<(&str, CfsSource)> = vec![ + (".fnm", CfsSource::Mem(&fnm)), + (".tis", CfsSource::Path(&tis_path)), + (".frq", CfsSource::Mem(&frq)), + ]; + let mut actual = Vec::new(); + write_cfs_streaming(&mut actual, "_9", &streaming_files).unwrap(); + + std::fs::remove_file(&tis_path).ok(); + + assert_eq!(actual, expected); + } + + #[test] + fn write_to_errors_when_path_source_shrinks_after_len_was_measured() { + let path = temp_path_named("toctou"); + std::fs::write(&path, b"this file starts out long enough to matter").unwrap(); + + let src = CfsSource::Path(&path); + // Same call `write_cfs_streaming` makes to compute the CFS directory offset. + let measured_len = src.len().unwrap(); + + // Simulate a TOCTOU race: the file is truncated/replaced by something else between the + // `len()` call used for the directory offset and the actual data write. + std::fs::write(&path, b"short").unwrap(); + + let mut out = Vec::new(); + let err = src.write_to(&mut out, measured_len).unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof); + + std::fs::remove_file(&path).ok(); + } + + #[test] + fn write_to_succeeds_when_expected_len_matches_actual_file_size() { + // Guard against the fix spuriously erroring on a legitimate full copy. + let path = temp_path_named("full_copy"); + let data = b"a normal file that is not touched after measurement".to_vec(); + std::fs::write(&path, &data).unwrap(); + + let src = CfsSource::Path(&path); + let measured_len = src.len().unwrap(); + assert_eq!(measured_len, data.len() as u64); + + let mut out = Vec::new(); + src.write_to(&mut out, measured_len).unwrap(); + assert_eq!(out, data); + + std::fs::remove_file(&path).ok(); + } + #[test] fn cfs_roundtrips_sub_files_through_reader() { let fnm = b"field-infos-bytes".to_vec(); diff --git a/sdsearch-core/src/zsl/writer/index_writer.rs b/sdsearch-core/src/zsl/writer/index_writer.rs index 6bca196..0d9f66a 100644 --- a/sdsearch-core/src/zsl/writer/index_writer.rs +++ b/sdsearch-core/src/zsl/writer/index_writer.rs @@ -6,8 +6,8 @@ //! open (excludes another native writer and ZSL). use super::lock::WriteLock; +use super::merge; use super::segments::{self, Generation, NewSegment}; -use super::{durability, merge}; use super::{write_segment_cfs, WriterDoc, WriterOpts}; use crate::index::IndexReader; use crate::zsl::deletes::DeletedDocs; @@ -154,26 +154,18 @@ impl IndexWriter { }); } - // 3) merge -> bytes of the merged .cfs. Name = next from name_counter. + // 3+4) stream the merge straight to a DURABLE {merged_name}.cfs (fsync happens inside). + // Peak heap is bounded — postings/positions/stored are streamed through temp files — and + // the bytes are identical to merge_segments. mmaps are dropped on return, so on Windows + // the old .cfs can be unlinked afterwards. Name = next from name_counter. let gen = segments::read_generation(&self.dir)?; let merged_name = segments::segment_name(gen.name_counter); let refs: Vec<(String, i64)> = infos.iter().map(|s| (s.name.clone(), s.del_gen)).collect(); - let result = merge::merge_segments(&self.dir, &merged_name, &refs)?; - // (merge_segments dropped its mmaps on return => on Windows the old ones can be unlinked) - - // 4) write the merged .cfs DURABLE (fsync). Not yet referenced by any generation. - durability::write_durable( - &self.dir.join(format!("{merged_name}.cfs")), - &result.cfs_bytes, - )?; + let doc_count = merge::merge_segments_streaming(&self.dir, &merged_name, &refs)?; // 5) write segments_{N+1} (fsync) + atomic flip of segments.gen. - let new_gen = segments::write_optimized_generation( - &self.dir, - &gen, - &merged_name, - result.doc_count as u32, - )?; + let new_gen = + segments::write_optimized_generation(&self.dir, &gen, &merged_name, doc_count as u32)?; // 6) orphan cleanup POST-flip (best-effort): .cfs/.del/.sti of the old segments. // Never before the flip (crash-safety invariant: the merged segment is referenced only @@ -198,7 +190,7 @@ impl IndexWriter { Ok(CommitReport { generation: new_gen, segments: vec![merged_name], - doc_count: result.doc_count, + doc_count, }) } @@ -602,6 +594,105 @@ mod tests { std::fs::remove_dir_all(&dir).ok(); } + #[test] + fn commit_prunes_old_generation_manifests_keeping_current_and_previous() { + let dir = temp_kb_full(); // KB: gen 6 (segments_6) + assert!(dir.join("segments_6").exists()); + + for i in 0..5 { + let mut w = IndexWriter::open(&dir, WriterOpts::default()).unwrap(); + w.add_document(doc_mark(i)).unwrap(); + w.commit().unwrap(); + } + // 5 commits: gen 6 -> 11 (7,8,9,10,11). + let gen = segments::read_generation(&dir).unwrap(); + assert_eq!(gen.generation, 11); + + // strictly-older-than-previous manifests are gone. + assert!(!dir.join("segments_6").exists()); + assert!(!dir.join("segments_7").exists()); + assert!(!dir.join("segments_8").exists()); + assert!(!dir.join("segments_9").exists()); + // current (11 == "b") and immediately-previous (10 == "a") survive: grace window for + // lock-free concurrent readers that read segments.gen just before the last flip. + assert!(dir + .join(format!("segments_{}", crate::zsl::segments::to_base36(10))) + .exists()); + assert!(dir + .join(format!("segments_{}", crate::zsl::segments::to_base36(11))) + .exists()); + + // segments.gen and the actual segment data (.cfs) are untouched by the pruning. + assert!(dir.join("segments.gen").exists()); + let idx = ZslIndex::open(&dir).unwrap(); + assert_eq!(idx.num_docs(), 20 + 5); + + std::fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn commit_prunes_using_numeric_not_lexical_order_across_base36_rollover() { + let dir = temp_kb_full(); // KB: gen 6 + // enough commits to cross the base36 rollover ("z" = 35 -> "10" = 36), plus margin. + for i in 0..35 { + let mut w = IndexWriter::open(&dir, WriterOpts::default()).unwrap(); + w.add_document(doc_mark(100 + i)).unwrap(); + w.commit().unwrap(); + } + let gen = segments::read_generation(&dir).unwrap(); + assert_eq!(gen.generation, 41); // 6 + 35 + + // a lexical-comparison bug would keep "segments_z" (35) forever (the string "z" + // sorts after two-digit base36 names); numeric comparison must have pruned it. + assert!(!dir.join("segments_z").exists()); + assert!(!dir.join("segments_10").exists()); // 36, also strictly older than 40 + + // current (41) and immediately-previous (40) survive. + let prev_name = format!("segments_{}", crate::zsl::segments::to_base36(40)); + let cur_name = format!("segments_{}", crate::zsl::segments::to_base36(41)); + assert!(dir.join(&prev_name).exists()); + assert!(dir.join(&cur_name).exists()); + + assert!(dir.join("segments.gen").exists()); + let idx = ZslIndex::open(&dir).unwrap(); + assert_eq!(idx.num_docs(), 20 + 35); + + std::fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn optimize_prunes_old_generation_manifests_too() { + let dir = temp_kb_full(); // KB: gen 6 + for i in 0..3 { + let mut w = IndexWriter::open(&dir, WriterOpts::default()).unwrap(); + w.add_document(doc_mark(i)).unwrap(); + w.commit().unwrap(); + } + let before = segments::read_generation(&dir).unwrap(); + assert_eq!(before.generation, 9); // 6 + 3 + assert!(dir.join("segments_9").exists()); + + let w = IndexWriter::open(&dir, WriterOpts::default()).unwrap(); + let rep = w.optimize().unwrap(); + assert_eq!(rep.generation, 10); // one flip: merge only (commit_inner was a no-op) + + // everything strictly older than the previous gen (9) is gone. + assert!(!dir.join("segments_6").exists()); + assert!(!dir.join("segments_7").exists()); + assert!(!dir.join("segments_8").exists()); + // current (10 == "a") and immediately-previous (9) survive. + assert!(dir.join("segments_9").exists()); + assert!(dir + .join(format!("segments_{}", crate::zsl::segments::to_base36(10))) + .exists()); + + assert!(dir.join("segments.gen").exists()); + let idx = ZslIndex::open(&dir).unwrap(); + assert_eq!(idx.num_docs(), 20 + 3); + + std::fs::remove_dir_all(&dir).ok(); + } + #[test] fn empty_commit_is_a_noop() { let dir = temp_kb_full(); diff --git a/sdsearch-core/src/zsl/writer/merge.rs b/sdsearch-core/src/zsl/writer/merge.rs index 898bec4..4f5badc 100644 --- a/sdsearch-core/src/zsl/writer/merge.rs +++ b/sdsearch-core/src/zsl/writer/merge.rs @@ -11,11 +11,14 @@ //! sorts by new doc-id; drops terms with no live docs. //! 4. serializes reusing the `write_segment_cfs` primitives (norms COPIED verbatim via write_norms_raw). +use super::cfs::{write_cfs_streaming, CfsSource}; use super::invert::{FieldMeta, StoredField, TermPostings}; use super::{assemble_cfs, fnm, norms, stored, terms}; use crate::index::IndexReader; use crate::zsl::segment::ZslSegment; -use std::collections::{BTreeMap, HashMap}; +use std::cmp::Reverse; +use std::collections::{BTreeMap, BinaryHeap, HashMap}; +use std::fs::File; use std::io; use std::path::Path; @@ -27,6 +30,10 @@ pub struct MergeResult { /// Merges `segments` (name, del_gen) — in `segments_N` order — into a `.cfs` named /// `merged_name`. See the module doc. +/// +/// Retained as the differential-test oracle (see `assert_streaming_byte_identical` below) and +/// as a potential future small-index fallback; it is NOT on the production `optimize()` path, +/// which uses [`merge_segments_streaming`]. pub fn merge_segments( index_dir: &Path, merged_name: &str, @@ -180,6 +187,255 @@ pub fn merge_segments( }) } +/// Streaming, bounded-memory counterpart of [`merge_segments`]. Reproduces its 4 phases EXACTLY +/// (byte-identical `.cfs`) but streams the big blocks (`.fdt`/`.frq`/`.prx`) through temp files +/// under `index_dir` instead of building them fully in RAM, and writes the merged `.cfs` durably +/// (fsync) to `index_dir/{merged_name}.cfs` itself before returning. Returns the live `doc_count`. +/// +/// Only the small per-field/index blocks stay in RAM (`.fnm`/`.fdx`/`.nrm`/`.tis`/`.tii` + +/// `doc_maps`/`norm_cols`), so peak heap is independent of total text volume. The term merge is a +/// k-way merge across each segment's [`ZslSegment::term_cursor`] (already in ZSL canonical order), +/// so terms are fed to the streaming dict writer in the SAME order the batch path sorts them into. +/// +/// Temp files (`{merged_name}.fdt.tmp` / `.frq.tmp` / `.prx.tmp`) are removed on BOTH the success +/// and error paths. +pub fn merge_segments_streaming( + index_dir: &Path, + merged_name: &str, + segments: &[(String, i64)], +) -> io::Result { + let fdt_tmp = index_dir.join(format!("{merged_name}.fdt.tmp")); + let frq_tmp = index_dir.join(format!("{merged_name}.frq.tmp")); + let prx_tmp = index_dir.join(format!("{merged_name}.prx.tmp")); + + let result = merge_streaming_inner( + index_dir, + merged_name, + segments, + &fdt_tmp, + &frq_tmp, + &prx_tmp, + ); + + // Clean up temp files on BOTH the success and error paths (best-effort). + let _ = std::fs::remove_file(&fdt_tmp); + let _ = std::fs::remove_file(&frq_tmp); + let _ = std::fs::remove_file(&prx_tmp); + + result +} + +fn merge_streaming_inner( + index_dir: &Path, + merged_name: &str, + segments: &[(String, i64)], + fdt_tmp: &Path, + frq_tmp: &Path, + prx_tmp: &Path, +) -> io::Result { + let segs: Vec = segments + .iter() + .map(|(name, dg)| ZslSegment::open_named(index_dir, name, *dg)) + .collect::>()?; + + // ---- phase 1: field union by first-seen (identical to merge_segments) ---- + let mut field_index: HashMap = HashMap::new(); + let mut fields: Vec = Vec::new(); + for seg in &segs { + for fi in seg.field_infos() { + let idx = *field_index.entry(fi.name.clone()).or_insert_with(|| { + fields.push(FieldMeta { + name: fi.name.clone(), + indexed: false, + }); + fields.len() - 1 + }); + fields[idx].indexed |= fi.is_indexed; + } + } + + // ---- phase 2: dense renumbering; stream stored to a temp .fdt (fdx kept in RAM), norms in RAM ---- + let mut doc_maps: Vec>> = Vec::with_capacity(segs.len()); + let mut norm_cols: Vec> = vec![Vec::new(); fields.len()]; + let mut next_id = 0usize; + + let mut fdx_buf: Vec = Vec::new(); + let mut stored_writer = stored::StoredStreamWriter::new(File::create(fdt_tmp)?, &mut fdx_buf); + + for seg in &segs { + let local_fields = seg.field_infos(); + let mut map = vec![None; seg.max_doc()]; + // `local` is a semantic doc-id (indexes is_deleted/stored_raw/norm_bytes and feeds + // doc_maps), not a mere iteration index: hence the range loop, not an iter. + #[allow(clippy::needless_range_loop)] + for local in 0..seg.max_doc() { + if seg.is_deleted(local) { + continue; + } + map[local] = Some(next_id); + + // stored: stream in order, remapping field_num local -> merged. + let raw = seg.stored_raw(local)?; + let mut remapped: Vec = Vec::with_capacity(raw.len()); + for r in raw { + if r.is_binary { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "merge: binary stored field not supported (local field_num {})", + r.field_num + ), + )); + } + let name = &local_fields + .get(r.field_num) + .ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("merge: stored field_num {} out of range", r.field_num), + ) + })? + .name; + remapped.push(StoredField { + field_num: field_index[name], + value: r.value, + tokenized: r.tokenized, + }); + } + stored_writer.add_doc(&remapped)?; + + // norms: for each indexed merged field, copy the segment's raw byte (255 if absent). + for (mf, field) in fields.iter().enumerate() { + if !field.indexed { + continue; + } + let byte = seg + .norm_bytes(&field.name) + .and_then(|col| col.get(local).copied()) + .unwrap_or(255); + norm_cols[mf].push(byte); + } + next_id += 1; + } + doc_maps.push(map); + } + let doc_count = next_id; + stored_writer.finish()?; // flush + close the temp .fdt; fdx_buf is now complete + + // ---- phase 3 guard: a segment with indexed fields MUST have .prx (identical to merge_segments) ---- + for seg in &segs { + let has_indexed = seg.field_infos().iter().any(|fi| fi.is_indexed); + if has_indexed && !seg.has_prx() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "merge: segment with indexed fields but no .prx — positions unavailable", + )); + } + } + + // ---- phase 3: k-way term merge → streamed to temp .frq/.prx, .tis/.tii kept in RAM ---- + // .tis/.tii are Write+Seek (Cursor over a Vec, so the header term-counts can be back-patched); + // .frq/.prx are the big blocks, streamed to temp files. + let mut tis_buf = io::Cursor::new(Vec::new()); + let mut tii_buf = io::Cursor::new(Vec::new()); + let mut dict_writer = terms::TermDictStreamWriter::new( + &mut tis_buf, + &mut tii_buf, + File::create(frq_tmp)?, + File::create(prx_tmp)?, + )?; + + // one cursor per segment; each yields (field, term) ascending. A min-heap (BinaryHeap is a + // max-heap, so keys are wrapped in Reverse) drives the k-way merge. Keys are cloned into the + // heap (bounded: at most one small (field, term) per segment at a time) because the cursors' + // borrows end at each `peek`, so we cannot hold `&str` across an `advance`. + let mut cursors: Vec<_> = segs.iter().map(|s| s.term_cursor()).collect(); + let mut heap: BinaryHeap> = BinaryHeap::new(); + for (si, cur) in cursors.iter().enumerate() { + if let Some((f, t)) = cur.peek() { + heap.push(Reverse((f.to_string(), t.to_string(), si))); + } + } + + while let Some(Reverse((field, term, first_si))) = heap.pop() { + // gather EVERY segment currently positioned at this exact (field, term). Advancing a + // cursor and re-pushing its next (strictly greater) term keeps the heap in sync. + let mut contributing: Vec = vec![first_si]; + cursors[first_si].advance(); + if let Some((f, t)) = cursors[first_si].peek() { + heap.push(Reverse((f.to_string(), t.to_string(), first_si))); + } + loop { + let same = matches!(heap.peek(), Some(Reverse((f, t, _))) if *f == field && *t == term); + if !same { + break; + } + let Reverse((_, _, si)) = heap.pop().unwrap(); + contributing.push(si); + cursors[si].advance(); + if let Some((f, t)) = cursors[si].peek() { + heap.push(Reverse((f.to_string(), t.to_string(), si))); + } + } + + // Contributing segments in ASCENDING index (+ locals ascending within each) => new + // doc-ids come out ascending, matching merge_segments' BTreeMap ordering + // (dense renumbering assigns lower ids to earlier segments). + contributing.sort_unstable(); + let mf = field_index[&field]; + let mut docs: Vec<(usize, Vec)> = Vec::new(); + for &si in &contributing { + // positions_all ALREADY filters deletes. It may return an unordered HashMap, so sort + // locals ascending before remapping so the new ids come out ascending. + let mut locals: Vec<(usize, Vec)> = + segs[si].positions_all(&field, &term).into_iter().collect(); + locals.sort_by_key(|(local, _)| *local); + for (local, positions) in locals { + if let Some(new_id) = doc_maps[si][local] { + docs.push((new_id, positions)); + } + } + } + // terms with no live docs are dropped (== ZSL / merge_segments). + if docs.is_empty() { + continue; + } + dict_writer.add_term(&TermPostings { + field_num: mf, + text: term, + docs, + })?; + } + dict_writer.finish()?; // back-patch term counts + close the temp .frq/.prx + + // ---- phase 4: assemble the .cfs (small blocks from RAM, big ones streamed from temp files) ---- + let fnm_bytes = fnm::write_fnm(&fields); + let nrm = norms::write_norms_raw(&norm_cols); + let tis_bytes = tis_buf.into_inner(); + let tii_bytes = tii_buf.into_inner(); + + // Same file order as `assemble_cfs`: .fdx .fdt .fnm .nrm .tis .tii .frq .prx. + let files: Vec<(&str, CfsSource)> = vec![ + (".fdx", CfsSource::Mem(&fdx_buf)), + (".fdt", CfsSource::Path(fdt_tmp)), + (".fnm", CfsSource::Mem(&fnm_bytes)), + (".nrm", CfsSource::Mem(&nrm)), + (".tis", CfsSource::Mem(&tis_bytes)), + (".tii", CfsSource::Mem(&tii_bytes)), + (".frq", CfsSource::Path(frq_tmp)), + (".prx", CfsSource::Path(prx_tmp)), + ]; + + // Write the merged .cfs durably: stream it into the file, then fsync BEFORE returning so the + // generation flip in optimize() only references a durable .cfs. + let cfs_path = index_dir.join(format!("{merged_name}.cfs")); + let mut cfs_file = File::create(&cfs_path)?; + write_cfs_streaming(&mut cfs_file, merged_name, &files)?; + cfs_file.sync_all()?; + + Ok(doc_count) +} + #[cfg(test)] mod tests { use super::*; @@ -378,6 +634,202 @@ mod tests { merge_segments(&dir, "_m", &[("_x".to_string(), -1)]).unwrap(); } + // ---- differential gate: streaming merge must be byte-identical to the batch oracle ---- + + /// Runs BOTH the batch oracle (`merge_segments`, bytes captured in RAM) and + /// `merge_segments_streaming` (writes `_m.cfs`) over the SAME `merged_name` and asserts the + /// resulting `.cfs` bytes are identical. The name must match: the CFS sub-file directory + /// embeds `{name}{ext}`, so different names would differ trivially. Also asserts the temp + /// files were cleaned up. + fn assert_streaming_byte_identical(dir: &std::path::Path, refs: &[(String, i64)]) { + let oracle = merge_segments(dir, "_m", refs).unwrap(); // RAM only; does NOT write + let doc_count = merge_segments_streaming(dir, "_m", refs).unwrap(); // writes _m.cfs (fsync) + assert_eq!(doc_count, oracle.doc_count, "doc_count mismatch"); + let streamed = std::fs::read(dir.join("_m.cfs")).unwrap(); + assert_eq!( + streamed, oracle.cfs_bytes, + "merged .cfs bytes differ (streaming vs batch)" + ); + // temp files removed on the success path + assert!(!dir.join("_m.fdt.tmp").exists()); + assert!(!dir.join("_m.frq.tmp").exists()); + assert!(!dir.join("_m.prx.tmp").exists()); + } + + #[test] + fn streaming_matches_batch_multi_segment_no_deletes() { + let dir = temp_kb_full(); + let opts = WriterOpts { + max_buffered_docs: 2, + ..WriterOpts::default() + }; + let mut w = IndexWriter::open(&dir, opts).unwrap(); + for i in 0..4 { + w.add_document(doc_mark(i)).unwrap(); + } + w.commit().unwrap(); + + let infos = read_segment_infos(&dir).unwrap(); + assert!(infos.len() >= 2, "expected multi-segment base"); + assert!( + infos.iter().all(|s| s.del_gen == -1), + "scenario (a) must have NO deletes" + ); + let refs: Vec<(String, i64)> = infos.iter().map(|s| (s.name.clone(), s.del_gen)).collect(); + assert_streaming_byte_identical(&dir, &refs); + std::fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn streaming_matches_batch_deletes_across_two_base_segments() { + let dir = temp_kb_full(); + let opts = WriterOpts { + max_buffered_docs: 2, + ..WriterOpts::default() + }; + let mut w = IndexWriter::open(&dir, opts).unwrap(); + for i in 0..4 { + w.add_document(doc_mark(i)).unwrap(); + } + w.commit().unwrap(); + + // delete a doc in _2 (gid 0) and the first doc of _3 (gid 20), in one commit. + let mut w2 = IndexWriter::open(&dir, WriterOpts::default()).unwrap(); + w2.delete_document(0); + w2.delete_document(20); + w2.commit().unwrap(); + + let infos = read_segment_infos(&dir).unwrap(); + assert!( + infos.iter().any(|s| s.del_gen != -1), + "scenario (b) needs deletes" + ); + let refs: Vec<(String, i64)> = infos.iter().map(|s| (s.name.clone(), s.del_gen)).collect(); + assert_streaming_byte_identical(&dir, &refs); + std::fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn streaming_matches_batch_single_segment_with_deletes() { + let dir = temp_kb_full(); + let mut w = IndexWriter::open(&dir, WriterOpts::default()).unwrap(); + w.delete_document(5); + w.commit().unwrap(); + + let infos = read_segment_infos(&dir).unwrap(); + assert_eq!(infos.len(), 1, "scenario (c) is single-segment"); + assert_ne!(infos[0].del_gen, -1, "must have a .del"); + let refs: Vec<(String, i64)> = infos.iter().map(|s| (s.name.clone(), s.del_gen)).collect(); + assert_streaming_byte_identical(&dir, &refs); + std::fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn streaming_matches_batch_all_docs_deleted_zero_live_docs() { + // Empty-merge scenario: delete EVERY doc in the single-segment KB fixture (all 20, + // global ids 0..20 — the KB fixture is exactly one segment, `_2`, with `max_doc() == + // 20`, see `sdsearch-core/src/zsl/segment.rs`), so the merge input has ZERO live docs. + // This exercises phase 2/3 with empty `stored`/`term_map` (empty `.fdt`/`.fdx`, no + // terms survive the "drop terms with no live docs" filter) — the smallest possible + // merge input, not covered by the other differential scenarios (which all retain >=1 + // live doc). + let dir = temp_kb_full(); + let mut w = IndexWriter::open(&dir, WriterOpts::default()).unwrap(); + for gid in 0..20 { + w.delete_document(gid); + } + w.commit().unwrap(); + + let infos = read_segment_infos(&dir).unwrap(); + assert_eq!(infos.len(), 1, "scenario is single-segment, fully deleted"); + assert_eq!( + ZslIndex::open(&dir).unwrap().num_docs(), + 0, + "all docs must be deleted" + ); + + let refs: Vec<(String, i64)> = infos.iter().map(|s| (s.name.clone(), s.del_gen)).collect(); + let oracle = merge_segments(&dir, "_m", &refs).unwrap(); + assert_eq!( + oracle.doc_count, 0, + "merge of a fully-deleted segment has 0 live docs" + ); + + assert_streaming_byte_identical(&dir, &refs); + std::fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn streaming_errors_on_binary_stored_field() { + // same hand-crafted binary-flagged segment as the batch guard test, but exercised via the + // streaming path (which returns an Err rather than panicking through an unwrap). + let n = COUNTER.fetch_add(1, Ordering::Relaxed); + let dir = + std::env::temp_dir().join(format!("sdsearch_smerge_bin_{}_{}", std::process::id(), n)); + std::fs::create_dir_all(&dir).unwrap(); + + let fields = vec![FieldMeta { + name: "bin_field".to_string(), + indexed: false, + }]; + let fnm_bytes = fnm::write_fnm(&fields); + let payload = b"binpayload"; + let mut fdt = vec![0x01u8, 0x00u8, 0x02u8]; + crate::zsl::bytes::write_vint(&mut fdt, payload.len() as u64); + fdt.extend_from_slice(payload); + let fdx = vec![0u8; 8]; + let nrm = norms::write_norms_raw(&[Vec::new()]); + let dict = terms::write_term_dict(&[]); + let cfs_bytes = assemble_cfs("_x", &fnm_bytes, &fdt, &fdx, &nrm, &dict); + std::fs::write(dir.join("_x.cfs"), &cfs_bytes).unwrap(); + + let err = merge_segments_streaming(&dir, "_m", &[("_x".to_string(), -1)]).unwrap_err(); + assert!( + err.to_string() + .contains("binary stored field not supported"), + "unexpected error: {err}" + ); + // partial temp files (if any) are cleaned up even on the error path + assert!(!dir.join("_m.fdt.tmp").exists()); + std::fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn streaming_errors_on_indexed_field_without_prx() { + // same hand-crafted no-.prx segment as the batch guard test, exercised via streaming. + let n = COUNTER.fetch_add(1, Ordering::Relaxed); + let dir = std::env::temp_dir().join(format!( + "sdsearch_smerge_noprx_{}_{}", + std::process::id(), + n + )); + std::fs::create_dir_all(&dir).unwrap(); + + let fields = vec![FieldMeta { + name: "body".to_string(), + indexed: true, + }]; + let fnm_bytes = fnm::write_fnm(&fields); + let (fdt, fdx) = stored::write_stored(&[Vec::new()]); + let dict = terms::write_term_dict(&[]); + let files: Vec<(&str, &[u8])> = vec![ + (".fdx", &fdx), + (".fdt", &fdt), + (".fnm", &fnm_bytes), + (".tis", &dict.tis), + (".frq", &dict.frq), + ]; + let cfs_bytes = crate::zsl::writer::cfs::write_cfs("_x", &files); + std::fs::write(dir.join("_x.cfs"), &cfs_bytes).unwrap(); + + let err = merge_segments_streaming(&dir, "_m", &[("_x".to_string(), -1)]).unwrap_err(); + assert!( + err.to_string().contains("indexed fields but no .prx"), + "unexpected error: {err}" + ); + std::fs::remove_dir_all(&dir).ok(); + } + #[test] fn merge_segments_collapses_multiseg_with_delete() { let dir = temp_kb_full(); diff --git a/sdsearch-core/src/zsl/writer/segments.rs b/sdsearch-core/src/zsl/writer/segments.rs index 909ab18..0a9efaf 100644 --- a/sdsearch-core/src/zsl/writer/segments.rs +++ b/sdsearch-core/src/zsl/writer/segments.rs @@ -53,6 +53,59 @@ pub fn segment_name(counter: u32) -> String { format!("_{}", to_base36(counter as u64)) } +/// inverse of `to_base36`: parses a lowercase base-36 string to a number. `None` if any +/// byte isn't a valid base-36 digit or the string is empty — defensive, so a file we don't +/// recognize is never touched by the generation-manifest pruning below. +fn from_base36(s: &str) -> Option { + if s.is_empty() { + return None; + } + let mut n: u64 = 0; + for b in s.bytes() { + let digit = match b { + b'0'..=b'9' => b - b'0', + b'a'..=b'z' => b - b'a' + 10, + _ => return None, + }; + n = n.checked_mul(36)?.checked_add(digit as u64)?; + } + Some(n) +} + +/// Best-effort deletion of stale generation manifests: any `segments_` whose parsed +/// number is strictly less than `keep_from_gen`. Only ever touches files matching that exact +/// pattern — never `segments.gen`, never `.cfs`/`.del`/`.sti`/`.tmp`/lock files, never the +/// legacy no-suffix `segments` (generation 0) file. Comparison is NUMERIC (parsed), not +/// lexical, so e.g. `segments_10` (36) is correctly treated as newer than `segments_z` (35). +/// Individual `read_dir`/`remove_file` errors are ignored (mirrors the orphan `.cfs` cleanup +/// in `IndexWriter::optimize`): this is housekeeping, never load-bearing for correctness. +/// +/// MUST be called only AFTER `segments.gen` has been durably flipped to the new generation. +/// Callers pass `keep_from_gen = new_gen - 1` so both the current generation and the +/// immediately-previous one survive — a grace window for lock-free concurrent readers that +/// may have read the old `segments.gen` value an instant before the flip and are about to +/// open the generation manifest it pointed to. +fn prune_old_generations(index_dir: &Path, keep_from_gen: u64) { + let entries = match std::fs::read_dir(index_dir) { + Ok(e) => e, + Err(_) => return, + }; + for entry in entries.flatten() { + let Some(name) = entry.file_name().to_str().map(str::to_string) else { + continue; + }; + let Some(suffix) = name.strip_prefix("segments_") else { + continue; + }; + let Some(n) = from_base36(suffix) else { + continue; + }; + if n < keep_from_gen { + let _ = std::fs::remove_file(entry.path()); + } + } +} + fn read_gen_number(index_dir: &Path) -> std::io::Result { let data = std::fs::read(index_dir.join("segments.gen"))?; let mut pos = 0usize; @@ -219,6 +272,9 @@ pub fn write_generation_with_delgens( write_i64_be(&mut g, new_gen as i64); super::durability::write_atomic(&index_dir.join("segments.gen"), &g)?; + // best-effort prune of stale generation manifests, only now that the flip is durable. + prune_old_generations(index_dir, new_gen.saturating_sub(1)); + Ok(new_gen) } @@ -263,6 +319,9 @@ pub fn write_optimized_generation( write_i64_be(&mut g, new_gen as i64); super::durability::write_atomic(&index_dir.join("segments.gen"), &g)?; + // best-effort prune of stale generation manifests, only now that the flip is durable. + prune_old_generations(index_dir, new_gen.saturating_sub(1)); + Ok(new_gen) } @@ -432,4 +491,57 @@ mod tests { assert_eq!(reread.seg_count, 1); std::fs::remove_dir_all(&dir).ok(); } + + #[test] + fn write_generation_with_delgens_prunes_old_manifests_keeps_current_and_previous() { + use std::collections::HashMap; + let dir = temp_kb_gen(); // KB: segments_6 (gen 6) + segments.gen + let mut gen = read_generation(&dir).unwrap(); + for _ in 0..5 { + let new_gen = write_generation_with_delgens(&dir, &gen, &HashMap::new(), &[]).unwrap(); + gen = read_generation(&dir).unwrap(); + assert_eq!(gen.generation, new_gen); + } + // 5 flips from gen 6 -> gen 11. + assert_eq!(gen.generation, 11); + + // strictly-older-than-previous manifests are gone (base36: 6,7,8,9 == same digits). + assert!(!dir.join("segments_6").exists()); + assert!(!dir.join("segments_7").exists()); + assert!(!dir.join("segments_8").exists()); + assert!(!dir.join("segments_9").exists()); + // current (11 == "b") and immediately-previous (10 == "a") survive. + assert!(dir.join(format!("segments_{}", to_base36(10))).exists()); + assert!(dir.join(format!("segments_{}", to_base36(11))).exists()); + // never touch segments.gen. + assert!(dir.join("segments.gen").exists()); + + std::fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn prune_uses_numeric_not_lexical_comparison_across_base36_rollover() { + use std::collections::HashMap; + let dir = temp_kb_gen(); // KB: gen 6 + let mut gen = read_generation(&dir).unwrap(); + // advance past the base36 rollover ("z" = 35 -> "10" = 36) with margin. + while gen.generation < 40 { + write_generation_with_delgens(&dir, &gen, &HashMap::new(), &[]).unwrap(); + gen = read_generation(&dir).unwrap(); + } + assert_eq!(gen.generation, 40); + + // a lexical-comparison bug would keep "segments_z" forever (the string "z" sorts + // after "10"/"13"/etc.); numeric comparison must have pruned it (35 < 40-1). + assert!(!dir.join("segments_z").exists()); + assert!(!dir.join("segments_10").exists()); // 36, also strictly older than 39 + + // current (40) and immediately-previous (39) survive. + let prev_name = format!("segments_{}", to_base36(39)); + let cur_name = format!("segments_{}", to_base36(40)); + assert!(dir.join(&prev_name).exists()); + assert!(dir.join(&cur_name).exists()); + + std::fs::remove_dir_all(&dir).ok(); + } } diff --git a/sdsearch-core/src/zsl/writer/stored.rs b/sdsearch-core/src/zsl/writer/stored.rs index 74aa624..531353d 100644 --- a/sdsearch-core/src/zsl/writer/stored.rs +++ b/sdsearch-core/src/zsl/writer/stored.rs @@ -5,20 +5,70 @@ //! The trailing `\n` that compactText adds is preserved verbatim (it is part of the value). use super::invert::StoredField; -use crate::zsl::bytes::{write_i64_be, write_modified_utf8, write_vint}; +use crate::zsl::bytes::{write_modified_utf8, write_vint}; +use std::io::{self, Write}; + +/// Writes `.fdt`/`.fdx` incrementally, one doc at a time. `.fdx` gets the current `.fdt` +/// length (i64 BE) at the start of each doc; `.fdt` gets `VInt(count)` + per-field +/// `VInt(field_num)` + flags + `String(value)`. No back-patch: same byte layout as +/// `write_stored`, just emitted doc-by-doc instead of built in one pass. +pub struct StoredStreamWriter { + fdt: Fdt, + fdx: Fdx, + fdt_len: u64, +} + +impl StoredStreamWriter { + pub fn new(fdt: Fdt, fdx: Fdx) -> Self { + Self { + fdt, + fdx, + fdt_len: 0, + } + } + + /// Appends one doc's stored fields: an `.fdx` pointer to the current `.fdt` length, + /// then the doc's `.fdt` block. + pub fn add_doc(&mut self, fields: &[StoredField]) -> io::Result<()> { + self.fdx.write_all(&(self.fdt_len as i64).to_be_bytes())?; + + // Build the block in RAM first (same layout as `write_stored`) so we can measure its + // length without requiring `Fdt: Seek`, then stream it out in one write. + let mut block = Vec::new(); + write_vint(&mut block, fields.len() as u64); + for sf in fields { + write_vint(&mut block, sf.field_num as u64); + block.push(if sf.tokenized { 0x01 } else { 0x00 }); // never binary here + write_modified_utf8(&mut block, &sf.value); + } + + self.fdt.write_all(&block)?; + self.fdt_len += block.len() as u64; + Ok(()) + } + + /// Flushes both sinks. No back-patch is needed for this format. + pub fn finish(mut self) -> io::Result<()> { + self.fdt.flush()?; + self.fdx.flush()?; + Ok(()) + } +} /// Writes (fdt, fdx) for the per-doc stored fields. pub fn write_stored(docs_stored: &[Vec]) -> (Vec, Vec) { let mut fdt = Vec::new(); let mut fdx = Vec::new(); - for fields in docs_stored { - write_i64_be(&mut fdx, fdt.len() as i64); // offset of this doc's block in .fdt - write_vint(&mut fdt, fields.len() as u64); - for sf in fields { - write_vint(&mut fdt, sf.field_num as u64); - fdt.push(if sf.tokenized { 0x01 } else { 0x00 }); // never binary here - write_modified_utf8(&mut fdt, &sf.value); + { + let mut writer = StoredStreamWriter::new(&mut fdt, &mut fdx); + for fields in docs_stored { + writer + .add_doc(fields) + .expect("writing to an in-memory Vec cannot fail"); } + writer + .finish() + .expect("flushing an in-memory Vec cannot fail"); } (fdt, fdx) } @@ -29,6 +79,86 @@ mod tests { use crate::zsl::fields::FieldInfo; use crate::zsl::stored::read_stored_fields; + fn sample_docs() -> Vec> { + vec![ + // multiple fields, tokenized flag on and off + vec![ + StoredField { + field_num: 0, + value: "New workflow\n".into(), + tokenized: true, + }, + StoredField { + field_num: 1, + value: "42".into(), + tokenized: false, + }, + ], + // empty doc (no stored fields at all) + vec![], + // unicode value (multi-byte chars + NUL) + vec![StoredField { + field_num: 2, + value: "über\u{1F680}\u{0}end".into(), + tokenized: true, + }], + // single field, not tokenized + vec![StoredField { + field_num: 3, + value: "TICKET-12345".into(), + tokenized: false, + }], + ] + } + + // --- independent oracle: the ORIGINAL (pre-streaming, commit 583e39b) `write_stored` + // algorithm, duplicated here under a `reference_*` name so it does NOT call into + // `StoredStreamWriter`/`write_stored` (the current module's real writer, now a thin wrapper + // over the streaming writer) — comparing against those would make the byte-identity test + // compare the streaming writer against itself (tautological). `write_vint`/`write_modified_utf8` + // are reused from `crate::zsl::bytes` since they are untouched by the streaming refactor; + // `write_i64_be` (only used here) is imported locally for the same reason. + + fn reference_write_stored(docs_stored: &[Vec]) -> (Vec, Vec) { + use crate::zsl::bytes::write_i64_be; + let mut fdt = Vec::new(); + let mut fdx = Vec::new(); + for fields in docs_stored { + write_i64_be(&mut fdx, fdt.len() as i64); // offset of this doc's block in .fdt + write_vint(&mut fdt, fields.len() as u64); + for sf in fields { + write_vint(&mut fdt, sf.field_num as u64); + fdt.push(if sf.tokenized { 0x01 } else { 0x00 }); // never binary here + write_modified_utf8(&mut fdt, &sf.value); + } + } + (fdt, fdx) + } + + #[test] + fn stream_writer_matches_batch_writer_byte_for_byte() { + let docs = sample_docs(); + + // Independent oracle: the pre-streaming reference implementation (duplicated from git + // commit 583e39b, before `write_stored` became a thin wrapper over + // `StoredStreamWriter`), NOT `write_stored` itself (which now shares its code with the + // streaming writer under test and would make this comparison tautological). + let (expected_fdt, expected_fdx) = reference_write_stored(&docs); + + let mut fdt_buf = Vec::new(); + let mut fdx_buf = Vec::new(); + { + let mut writer = StoredStreamWriter::new(&mut fdt_buf, &mut fdx_buf); + for fields in &docs { + writer.add_doc(fields).unwrap(); + } + writer.finish().unwrap(); + } + + assert_eq!(fdt_buf, expected_fdt, "fdt mismatch"); + assert_eq!(fdx_buf, expected_fdx, "fdx mismatch"); + } + #[test] fn stored_roundtrips_through_reader_preserving_trailing_newline() { let stored = vec![ diff --git a/sdsearch-core/src/zsl/writer/terms.rs b/sdsearch-core/src/zsl/writer/terms.rs index 0c68fe4..8a9390b 100644 --- a/sdsearch-core/src/zsl/writer/terms.rs +++ b/sdsearch-core/src/zsl/writer/terms.rs @@ -12,6 +12,7 @@ use super::invert::TermPostings; use super::postings::write_term_postings; use crate::zsl::bytes::{write_i32_be, write_i64_be, write_modified_utf8, write_vint}; +use std::io::{self, Seek, SeekFrom, Write}; pub const INDEX_INTERVAL: u64 = 128; const MARKER: i32 = -3; // 0xFFFFFFFD @@ -25,55 +26,177 @@ pub struct DictFiles { pub prx: Vec, } -pub fn write_term_dict(terms: &[TermPostings]) -> DictFiles { - let mut tis = Vec::new(); - let mut tii = Vec::new(); - let mut frq = Vec::new(); - let mut prx = Vec::new(); +/// prev-term state for prefix-sharing + freq/prox pointer deltas: (text, field_num, freq_ptr, prox_ptr). +/// Owned (unlike the batch loop's `&str` borrow) because a streaming writer only sees one +/// term at a time and must remember the previous one across calls. +type PrevTerm = (String, usize, u64, u64); + +/// Streams the four dict sub-files. Terms MUST be added in ZSL canonical order +/// (fieldName·\0·text). docFreq is taken per term from `term.docs.len()`. `.tis`/`.tii` +/// need Seek to back-patch the 8-byte term-count in their headers on `finish`. +pub struct TermDictStreamWriter +where + Tis: Write + Seek, + Tii: Write + Seek, + Frq: Write, + Prx: Write, +{ + tis: Tis, + tii: Tii, + frq: Frq, + prx: Prx, + tis_len: u64, + frq_len: u64, + prx_len: u64, + prev: Option, + idx_prev: Option, + last_index_position: u64, + term_count: u64, +} + +impl TermDictStreamWriter +where + Tis: Write + Seek, + Tii: Write + Seek, + Frq: Write, + Prx: Write, +{ + pub fn new(mut tis: Tis, mut tii: Tii, frq: Frq, prx: Prx) -> io::Result { + let mut header = Vec::new(); + write_header(&mut header); + tis.write_all(&header)?; + tii.write_all(&header)?; + + // initial synthetic .tii entry (hand-written, NOT via dump_entry): + // the field number is a raw Int 0xFFFFFFFF, not a VInt. + let mut synthetic = Vec::new(); + write_vint(&mut synthetic, 0); // prefixChars + write_modified_utf8(&mut synthetic, ""); // empty suffix + write_i32_be(&mut synthetic, -1); // 0xFFFFFFFF + synthetic.push(0x0F); + write_vint(&mut synthetic, 0); // docFreq + write_vint(&mut synthetic, 0); // freqDelta + write_vint(&mut synthetic, 0); // proxDelta + write_vint(&mut synthetic, 24); // IndexDelta + tii.write_all(&synthetic)?; + + Ok(Self { + tis, + tii, + frq, + prx, + tis_len: header.len() as u64, // 24: header is already on disk + frq_len: 0, + prx_len: 0, + prev: None, + idx_prev: None, + last_index_position: 24, + term_count: 0, + }) + } + + /// Appends one term's dictionary entry (`.tis`, and `.tii` every `INDEX_INTERVAL`th + /// term) plus its postings (`.frq`/`.prx`). Terms must arrive in ZSL canonical order. + pub fn add_term(&mut self, term: &TermPostings) -> io::Result<()> { + // `write_term_postings` only depends on `term.docs` (it resets its own doc/position + // deltas per call), so writing into fresh local buffers reproduces the exact same + // bytes as appending into a running `.frq`/`.prx` buffer; the absolute pointers are + // then our own running lengths rather than the (always-zero) buffer-local ones. + let mut frq_buf = Vec::new(); + let mut prx_buf = Vec::new(); + write_term_postings(&mut frq_buf, &mut prx_buf, &term.docs); + let freq_ptr = self.frq_len; + let prox_ptr = self.prx_len; - write_header(&mut tis); - write_header(&mut tii); - - // initial synthetic .tii entry (hand-written, NOT via dump_entry): - // the field number is a raw Int 0xFFFFFFFF, not a VInt. - write_vint(&mut tii, 0); // prefixChars - write_modified_utf8(&mut tii, ""); // empty suffix - write_i32_be(&mut tii, -1); // 0xFFFFFFFF - tii.push(0x0F); - write_vint(&mut tii, 0); // docFreq - write_vint(&mut tii, 0); // freqDelta - write_vint(&mut tii, 0); // proxDelta - write_vint(&mut tii, 24); // IndexDelta - - // state of the previous term in the .tis (borrows term text from `terms`, no clone) - let mut prev: Option<(&str, usize, u64, u64)> = None; // (text, field, freqPtr, proxPtr) - // state of the last sample in the .tii - let mut idx_prev: Option<(&str, usize, u64, u64)> = None; - let mut last_index_position: u64 = 24; - - for (i, term) in terms.iter().enumerate() { - let (freq_ptr, prox_ptr) = write_term_postings(&mut frq, &mut prx, &term.docs); let doc_freq = term.doc_freq(); - dump_entry(&mut tis, prev, term, doc_freq, freq_ptr, prox_ptr); - prev = Some((term.text.as_str(), term.field_num, freq_ptr, prox_ptr)); + let mut tis_entry = Vec::new(); + dump_entry( + &mut tis_entry, + self.prev + .as_ref() + .map(|(t, f, fp, pp)| (t.as_str(), *f, *fp, *pp)), + term, + doc_freq, + freq_ptr, + prox_ptr, + ); + self.tis.write_all(&tis_entry)?; + self.tis_len += tis_entry.len() as u64; + self.prev = Some((term.text.clone(), term.field_num, freq_ptr, prox_ptr)); + + self.frq.write_all(&frq_buf)?; + self.frq_len += frq_buf.len() as u64; + self.prx.write_all(&prx_buf)?; + self.prx_len += prx_buf.len() as u64; + self.term_count += 1; // sample every indexInterval terms - if (i as u64 + 1).is_multiple_of(INDEX_INTERVAL) { - dump_entry(&mut tii, idx_prev, term, doc_freq, freq_ptr, prox_ptr); - let index_position = tis.len() as u64; - write_vint(&mut tii, index_position - last_index_position); - last_index_position = index_position; - idx_prev = Some((term.text.as_str(), term.field_num, freq_ptr, prox_ptr)); + if self.term_count.is_multiple_of(INDEX_INTERVAL) { + let mut tii_entry = Vec::new(); + dump_entry( + &mut tii_entry, + self.idx_prev + .as_ref() + .map(|(t, f, fp, pp)| (t.as_str(), *f, *fp, *pp)), + term, + doc_freq, + freq_ptr, + prox_ptr, + ); + let index_position = self.tis_len; + write_vint(&mut tii_entry, index_position - self.last_index_position); + self.last_index_position = index_position; + self.tii.write_all(&tii_entry)?; + self.idx_prev = Some((term.text.clone(), term.field_num, freq_ptr, prox_ptr)); } + + Ok(()) + } + + /// Back-patches the 8-byte term counts at offset 4 in `.tis`/`.tii` and flushes all + /// four sinks. + pub fn finish(mut self) -> io::Result<()> { + let term_count = self.term_count; + self.tis.seek(SeekFrom::Start(4))?; + self.tis.write_all(&term_count.to_be_bytes())?; + + let tii_count = (term_count - term_count % INDEX_INTERVAL) / INDEX_INTERVAL + 1; + self.tii.seek(SeekFrom::Start(4))?; + self.tii.write_all(&tii_count.to_be_bytes())?; + + self.tis.flush()?; + self.tii.flush()?; + self.frq.flush()?; + self.prx.flush()?; + Ok(()) } +} - let term_count = terms.len() as u64; - patch_long(&mut tis, 4, term_count); - let tii_count = (term_count - term_count % INDEX_INTERVAL) / INDEX_INTERVAL + 1; - patch_long(&mut tii, 4, tii_count); +pub fn write_term_dict(terms: &[TermPostings]) -> DictFiles { + let mut tis = io::Cursor::new(Vec::new()); + let mut tii = io::Cursor::new(Vec::new()); + let mut frq = Vec::new(); + let mut prx = Vec::new(); + { + let mut writer = TermDictStreamWriter::new(&mut tis, &mut tii, &mut frq, &mut prx) + .expect("writing to an in-memory Vec cannot fail"); + for term in terms { + writer + .add_term(term) + .expect("writing to an in-memory Vec cannot fail"); + } + writer + .finish() + .expect("writing to an in-memory Vec cannot fail"); + } - DictFiles { tis, tii, frq, prx } + DictFiles { + tis: tis.into_inner(), + tii: tii.into_inner(), + frq, + prx, + } } fn write_header(out: &mut Vec) { @@ -84,10 +207,6 @@ fn write_header(out: &mut Vec) { write_i32_be(out, MAX_SKIP_LEVELS); } -fn patch_long(buf: &mut [u8], offset: usize, v: u64) { - buf[offset..offset + 8].copy_from_slice(&v.to_be_bytes()); -} - /// Writes a term dict entry. Shares a prefix with `prev` only if it is of the SAME /// field; writes freq/prox as a delta relative to `prev`, or absolute if `prev` is None. fn dump_entry( @@ -236,4 +355,197 @@ mod tests { assert_eq!(dict.info("t", "w0000").unwrap().doc_freq, 1); assert_eq!(dict.info("t", "w0299").unwrap().doc_freq, 1); } + + /// Builds a hand-crafted, ZSL-canonically-sorted term list covering: multiple fields, + /// shared prefixes within a field, >128 terms (to exercise `.tii` sampling AND the + /// term-count back-patch), and multi-doc postings with positions. + fn multi_field_sample_terms() -> Vec { + let mut terms = Vec::new(); + + // field 0 ("body"): 150 terms sharing the "shared" prefix, sorted ascending, so + // consecutive terms within the field share a common prefix. Some have multi-doc + // postings with positions to exercise freq/prox deltas. + for i in 0..150usize { + let text = format!("shared{i:04}"); + let docs = if i % 10 == 0 { + vec![(i, vec![0u32, 3]), (i + 1000, vec![1u32])] + } else { + vec![(i, vec![0u32])] + }; + terms.push(TermPostings { + field_num: 0, + text, + docs, + }); + } + + // field 1 ("title"): a handful of terms, some sharing prefixes, to make sure + // prefix-sharing does NOT leak across fields (field 0's last term is "shared0149"). + for (text, docs) in [ + ("alpha".to_string(), vec![(2usize, vec![5u32])]), + ("alphabet".to_string(), vec![(3usize, vec![0u32, 1, 2])]), + ("beta".to_string(), vec![(4usize, vec![7u32])]), + ] { + terms.push(TermPostings { + field_num: 1, + text, + docs, + }); + } + + terms + } + + // --- independent oracle: the ORIGINAL (pre-streaming, commit 0654030) `write_term_dict` + // algorithm, duplicated here under `reference_*` names so it does NOT call into the + // current module's `dump_entry`/`common_prefix`/`write_header`/`write_term_dict` — those + // are now shared with `TermDictStreamWriter`, so comparing against them would make the + // byte-identity test compare the streaming writer against itself (tautological). Only + // `write_term_postings` (from `super`) is reused, since it is untouched by the streaming + // refactor and doesn't participate in the pointer/prefix-delta logic under test. + + fn reference_write_term_dict(terms: &[TermPostings]) -> DictFiles { + let mut tis = Vec::new(); + let mut tii = Vec::new(); + let mut frq = Vec::new(); + let mut prx = Vec::new(); + + reference_write_header(&mut tis); + reference_write_header(&mut tii); + + // initial synthetic .tii entry (hand-written, NOT via dump_entry): + // the field number is a raw Int 0xFFFFFFFF, not a VInt. + write_vint(&mut tii, 0); // prefixChars + write_modified_utf8(&mut tii, ""); // empty suffix + write_i32_be(&mut tii, -1); // 0xFFFFFFFF + tii.push(0x0F); + write_vint(&mut tii, 0); // docFreq + write_vint(&mut tii, 0); // freqDelta + write_vint(&mut tii, 0); // proxDelta + write_vint(&mut tii, 24); // IndexDelta + + // state of the previous term in the .tis + let mut prev: Option<(&str, usize, u64, u64)> = None; // (text, field, freqPtr, proxPtr) + // state of the last sample in the .tii + let mut idx_prev: Option<(&str, usize, u64, u64)> = None; + let mut last_index_position: u64 = 24; + + for (i, term) in terms.iter().enumerate() { + let (freq_ptr, prox_ptr) = write_term_postings(&mut frq, &mut prx, &term.docs); + let doc_freq = term.doc_freq(); + + reference_dump_entry(&mut tis, prev, term, doc_freq, freq_ptr, prox_ptr); + prev = Some((term.text.as_str(), term.field_num, freq_ptr, prox_ptr)); + + // sample every indexInterval terms + if (i as u64 + 1).is_multiple_of(INDEX_INTERVAL) { + reference_dump_entry(&mut tii, idx_prev, term, doc_freq, freq_ptr, prox_ptr); + let index_position = tis.len() as u64; + write_vint(&mut tii, index_position - last_index_position); + last_index_position = index_position; + idx_prev = Some((term.text.as_str(), term.field_num, freq_ptr, prox_ptr)); + } + } + + let term_count = terms.len() as u64; + reference_patch_long(&mut tis, 4, term_count); + let tii_count = (term_count - term_count % INDEX_INTERVAL) / INDEX_INTERVAL + 1; + reference_patch_long(&mut tii, 4, tii_count); + + DictFiles { tis, tii, frq, prx } + } + + fn reference_write_header(out: &mut Vec) { + write_i32_be(out, MARKER); + write_i64_be(out, 0); // placeholder termCount (back-patched at offset 4) + write_i32_be(out, INDEX_INTERVAL as i32); + write_i32_be(out, SKIP_INTERVAL); + write_i32_be(out, MAX_SKIP_LEVELS); + } + + fn reference_patch_long(buf: &mut [u8], offset: usize, v: u64) { + buf[offset..offset + 8].copy_from_slice(&v.to_be_bytes()); + } + + /// Writes a term dict entry. Shares a prefix with `prev` only if it is of the SAME + /// field; writes freq/prox as a delta relative to `prev`, or absolute if `prev` is None. + fn reference_dump_entry( + out: &mut Vec, + prev: Option<(&str, usize, u64, u64)>, + term: &TermPostings, + doc_freq: u32, + freq_ptr: u64, + prox_ptr: u64, + ) { + let (prefix_chars, prefix_bytes) = match prev { + Some((ptext, pfield, ..)) if pfield == term.field_num => { + reference_common_prefix(ptext, &term.text) + } + _ => (0, 0), + }; + write_vint(out, prefix_chars as u64); + write_modified_utf8(out, &term.text[prefix_bytes..]); + write_vint(out, term.field_num as u64); + write_vint(out, doc_freq as u64); + match prev { + Some((_, _, pf, pp)) => { + write_vint(out, freq_ptr - pf); + write_vint(out, prox_ptr - pp); + } + None => { + write_vint(out, freq_ptr); + write_vint(out, prox_ptr); + } + } + // skipOffset is omitted: docFreq is always < skipInterval. + } + + /// Common prefix in (chars, bytes). Matching chars ⟺ matching bytes (UTF-8), + /// so counting equal leading chars reproduces ZSL's byte-then-char calculation. + fn reference_common_prefix(a: &str, b: &str) -> (usize, usize) { + let mut chars = 0usize; + let mut bytes = 0usize; + for (ca, cb) in a.chars().zip(b.chars()) { + if ca == cb { + chars += 1; + bytes += ca.len_utf8(); + } else { + break; + } + } + (chars, bytes) + } + + #[test] + fn stream_writer_matches_batch_writer_byte_for_byte() { + let terms = multi_field_sample_terms(); + assert!( + terms.len() > 128, + "must exceed indexInterval to test .tii sampling + patch" + ); + + // Independent oracle: the pre-streaming reference implementation (duplicated from + // git commit 0654030, before `write_term_dict` became a thin wrapper over + // `TermDictStreamWriter`), NOT `write_term_dict` itself (which now shares its code + // with the streaming writer under test and would make this comparison tautological). + let expected = reference_write_term_dict(&terms); + + let mut tis = std::io::Cursor::new(Vec::new()); + let mut tii = std::io::Cursor::new(Vec::new()); + let mut frq = Vec::new(); + let mut prx = Vec::new(); + { + let mut writer = + TermDictStreamWriter::new(&mut tis, &mut tii, &mut frq, &mut prx).unwrap(); + for term in &terms { + writer.add_term(term).unwrap(); + } + writer.finish().unwrap(); + } + + assert_eq!(tis.into_inner(), expected.tis, "tis mismatch"); + assert_eq!(tii.into_inner(), expected.tii, "tii mismatch"); + assert_eq!(frq, expected.frq, "frq mismatch"); + assert_eq!(prx, expected.prx, "prx mismatch"); + } } diff --git a/sdsearch.stub.php b/sdsearch.stub.php index 03d0ed0..e7553d7 100644 --- a/sdsearch.stub.php +++ b/sdsearch.stub.php @@ -164,6 +164,13 @@ public function commit(): int {} * Commits, then merges all live segments (plus pending deletes) into a single * compacted segment. **Consumes** the writer, like {@see Writer::commit()}. * + * The merge is streaming and bounded-memory: peak heap is a per-term working set + * plus small per-document bookkeeping, independent of the index's total text volume + * (so optimizing a large index does not risk an OOM). While running it writes + * temporary files (`.fdt.tmp`/`.frq.tmp`/`.prx.tmp`) into the index dir. + * Prefer "open once → add_document* → optimize() once" per feed; the segment count + * only shrinks on optimize (no automatic merge policy). + * * @return int The document count the index will have after the optimize. * @throws \Exception if the writer is not open, or on error. */