use std::{
cmp::Ordering,
collections::BinaryHeap,
fs,
fs::File,
hash::{BuildHasher, Hash},
io::{self, BufReader, BufWriter, Error, ErrorKind, Read, Seek, SeekFrom, Write},
mem,
path::{Path, PathBuf},
sync::Arc,
time::{Duration, Instant},
vec::IntoIter,
};
use bumpalo::Bump;
use hashbrown::hash_map::{HashMap as HbHashMap, RawEntryMut};
use memmap2::Mmap;
use rustc_hash::{FxBuildHasher, FxHashMap};
use crate::superfile::{
BuildError,
format::{
self, FST_SEPARATOR,
checksum::{crc32c, crc32c_append},
},
fts::{
bm25,
dict::{DictBuilder, StreamingDictBuilder},
fst_value::{FstValue, INLINE_TF_MAX},
positions::{encode_run, read_varint, skip_run},
posting::{BLOCK_LEN, Block, EncodedBlock, encode_block},
tokenize::{AsciiLowerTokenizer, Tokenizer},
},
};
type TermIdMap = HbHashMap<&'static str, u32, FxBuildHasher>;
type DocTfMap = HbHashMap<&'static str, u32, FxBuildHasher>;
type DocPosHeadMap = HbHashMap<&'static str, u32, FxBuildHasher>;
const CHAIN_END: u32 = u32::MAX;
#[derive(Default)]
struct FinishProfile {
enabled: bool,
encode_calls: u64,
encode_df1: u64,
encode_pfor: u64,
encode_total: Duration,
encode_block_build: Duration,
encode_meta_write: Duration,
encode_skip_write: Duration,
encode_block_write: Duration,
fst_insert: Duration,
partition_flush: Duration,
lex_rank_build: Duration,
partition_sort: Duration,
mmap_open: Duration,
scratch_cleanup: Duration,
fst_close: Duration,
postings_close: Duration,
doc_lengths_emit: Duration,
blob_copy: Duration,
}
impl FinishProfile {
fn from_config() -> Self {
Self {
enabled: crate::config::global().diagnostics.fts_profile,
..Self::default()
}
}
}
pub(crate) const TERM_META_SIZE: usize = 20;
pub(crate) const TERM_META_POSITIONAL_SIZE: usize = 32;
pub(crate) const SKIP_ENTRY_SIZE: usize = 16;
pub(crate) const DOC_LENGTHS_ENTRY_SIZE: usize = 16;
pub const DEFAULT_SPILL_THRESHOLD_BYTES: usize = 256 * 1024 * 1024;
pub const DEFAULT_SPILL_PARTITIONS: usize = 128;
pub const DEFAULT_MAX_PARTITION_BYTES: u64 = 256 * 1024 * 1024;
const PARTITION_BUF_SIZE: usize = 64 * 1024;
const ACCUM_NEW_TERM_FIXED_BYTES: usize = 24;
const ACCUM_POSTING_BYTES: usize = 8;
const RADIX_SORT_MIN_TRIPLES: usize = 256;
const EXTERNAL_MERGE_CHUNK_CAP_TRIPLES: usize = 1024 * 1024;
const SORT_OUTPUT_BATCH_TRIPLES: usize = 4096;
struct ColumnState {
name: String,
doc_lengths: Vec<u32>,
total_tokens: u64,
positions: bool,
}
enum ColumnPostings {
InRam {
terms: FxHashMap<Box<str>, Vec<(u32, u32)>>,
pos_runs: FxHashMap<Box<str>, Vec<u8>>,
bytes: usize,
},
Spilled {
partitions: SpillStore,
term_to_id: TermIdMap,
id_to_term: Vec<&'static str>,
dense_doc_tf: Vec<u32>,
dense_doc_poshead: Vec<u32>,
updated_terms: Vec<u32>,
term_arena: Bump,
},
}
impl ColumnPostings {
fn new() -> Self {
Self::InRam {
terms: FxHashMap::default(),
pos_runs: FxHashMap::default(),
bytes: 0,
}
}
fn is_spilled(&self) -> bool {
matches!(self, Self::Spilled { .. })
}
}
const SPILL_BATCH_TRIPLES: usize = 341;
struct SpillPartition<const N: usize> {
path: PathBuf,
writer: Option<BufWriter<File>>,
batch: Vec<[u32; N]>,
}
const PLAIN_RECORD_LANES: usize = 3;
const POSITIONAL_RECORD_LANES: usize = 5;
#[inline(always)]
fn pos_off_lanes(off: u64) -> (u32, u32) {
(off as u32, (off >> 32) as u32)
}
#[inline(always)]
fn record_pos_off(rec: &[u32; POSITIONAL_RECORD_LANES]) -> u64 {
(rec[3] as u64) | ((rec[4] as u64) << 32)
}
struct PartitionPositions {
path: PathBuf,
writer: Option<BufWriter<File>>,
len: u64,
}
enum SpillStore {
Plain(Vec<SpillPartition<PLAIN_RECORD_LANES>>),
Positional {
partitions: Vec<SpillPartition<POSITIONAL_RECORD_LANES>>,
blobs: Vec<PartitionPositions>,
},
}
impl SpillStore {
fn n_partitions(&self) -> usize {
match self {
SpillStore::Plain(p) => p.len(),
SpillStore::Positional { partitions, .. } => partitions.len(),
}
}
}
#[inline(always)]
fn triple_term_id<const N: usize>(t: &[u32; N]) -> u32 {
t[0]
}
#[inline(always)]
fn triple_doc_id<const N: usize>(t: &[u32; N]) -> u32 {
t[1]
}
#[inline(always)]
fn triple_tf<const N: usize>(t: &[u32; N]) -> u32 {
t[2]
}
#[cfg(not(target_endian = "little"))]
#[inline(always)]
fn write_record<W: Write, const N: usize>(w: &mut W, rec: &[u32; N]) -> Result<(), BuildError> {
let mut buf = [0u8; MAX_RECORD_BYTES];
for (lane, v) in rec.iter().enumerate() {
buf[lane * 4..lane * 4 + 4].copy_from_slice(&v.to_le_bytes());
}
w.write_all(&buf[..mem::size_of::<[u32; N]>()])?;
Ok(())
}
#[inline(always)]
fn push_record_batched<const N: usize>(
partition: &mut SpillPartition<N>,
rec: [u32; N],
) -> Result<(), BuildError> {
partition.batch.push(rec);
if partition.batch.len() >= SPILL_BATCH_TRIPLES {
flush_partition_batch(partition)?;
}
Ok(())
}
#[inline]
fn flush_partition_batch<const N: usize>(
partition: &mut SpillPartition<N>,
) -> Result<(), BuildError> {
if partition.batch.is_empty() {
return Ok(());
}
let writer = partition
.writer
.as_mut()
.expect("partition writer is open before finish");
#[cfg(target_endian = "little")]
{
writer.write_all(bytemuck::cast_slice::<[u32; N], u8>(&partition.batch))?;
}
#[cfg(not(target_endian = "little"))]
{
for t in &partition.batch {
write_record(writer, t)?;
}
}
partition.batch.clear();
Ok(())
}
fn read_partition_records<const N: usize>(path: &Path) -> Result<Vec<[u32; N]>, BuildError> {
let mut bytes = Vec::new();
let mut f = File::open(path)?;
f.read_to_end(&mut bytes)?;
if bytes.is_empty() {
return Ok(Vec::new());
}
let rec_bytes = mem::size_of::<[u32; N]>();
if bytes.len() % rec_bytes != 0 {
return Err(BuildError::Io(Error::new(
ErrorKind::InvalidData,
format!(
"spill partition {path:?} length {} not a multiple of {}",
bytes.len(),
rec_bytes
),
)));
}
#[cfg(target_endian = "little")]
{
let records: &[[u32; N]] = bytemuck::try_cast_slice(&bytes).map_err(|_| {
BuildError::Io(Error::new(
ErrorKind::InvalidData,
"bytemuck: spill bytes failed alignment for the record slice",
))
})?;
Ok(records.to_vec())
}
#[cfg(not(target_endian = "little"))]
{
let n = bytes.len() / rec_bytes;
let mut out = Vec::with_capacity(n);
for i in 0..n {
let off = i * rec_bytes;
let mut rec = [0u32; N];
for (lane, slot) in rec.iter_mut().enumerate() {
let at = off + lane * 4;
*slot = u32::from_le_bytes(
bytes[at..at + 4]
.try_into()
.expect("invariant: 4-byte record lane"),
);
}
out.push(rec);
}
Ok(out)
}
}
fn build_lex_rank(id_to_term: &[&str]) -> (Vec<u32>, Vec<u32>) {
let n = id_to_term.len();
let mut by_lex: Vec<u32> = (0..n as u32).collect();
by_lex.sort_unstable_by(|&a, &b| {
id_to_term[a as usize]
.as_bytes()
.cmp(id_to_term[b as usize].as_bytes())
});
let mut rank = vec![0u32; n];
for (r, id) in by_lex.iter().enumerate() {
rank[*id as usize] = r as u32;
}
(rank, by_lex)
}
#[inline(always)]
fn compute_hash<Q: Hash + ?Sized, S: BuildHasher>(hash_builder: &S, key: &Q) -> u64 {
hash_builder.hash_one(key)
}
#[inline(always)]
fn intern_term_id(
term_to_id: &mut TermIdMap,
id_to_term: &mut Vec<&'static str>,
arena: &Bump,
term: &str,
) -> (u32, bool) {
let hash = compute_hash(term_to_id.hasher(), term);
match term_to_id
.raw_entry_mut()
.from_hash(hash, |existing| *existing == term)
{
RawEntryMut::Occupied(entry) => (*entry.get(), false),
RawEntryMut::Vacant(entry) => {
let id = id_to_term.len() as u32;
let arena_str: &str = arena.alloc_str(term);
let static_str: &'static str = unsafe { std::mem::transmute(arena_str) };
id_to_term.push(static_str);
entry.insert_hashed_nocheck(hash, static_str, id);
(id, true)
}
}
}
struct MergeEntry<const N: usize> {
sort_key: u64,
rec: [u32; N],
reader_idx: usize,
}
impl<const N: usize> PartialEq for MergeEntry<N> {
fn eq(&self, other: &Self) -> bool {
self.cmp(other) == Ordering::Equal
}
}
impl<const N: usize> Eq for MergeEntry<N> {}
impl<const N: usize> PartialOrd for MergeEntry<N> {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl<const N: usize> Ord for MergeEntry<N> {
fn cmp(&self, other: &Self) -> Ordering {
other
.sort_key
.cmp(&self.sort_key)
.then(other.reader_idx.cmp(&self.reader_idx))
}
}
#[inline(always)]
fn pack_sort_key(lex_rank: u32, doc_id: u32) -> u64 {
((lex_rank as u64) << 32) | (doc_id as u64)
}
enum PartitionIter<const N: usize> {
InMemory(IntoIter<[u32; N]>),
Merge {
readers: Vec<BufReader<File>>,
heap: BinaryHeap<MergeEntry<N>>,
_chunk_paths: Vec<PathBuf>,
},
}
impl<const N: usize> PartitionIter<N> {
fn next_with(&mut self, lex_rank: &[u32]) -> Option<Result<[u32; N], BuildError>> {
match self {
PartitionIter::InMemory(it) => it.next().map(Ok),
PartitionIter::Merge { readers, heap, .. } => {
let MergeEntry {
rec, reader_idx, ..
} = heap.pop()?;
match read_one_record::<_, N>(&mut readers[reader_idx]) {
Ok(Some(next_t)) => {
let next_id = triple_term_id(&next_t);
let next_doc = triple_doc_id(&next_t);
let key = pack_sort_key(lex_rank[next_id as usize], next_doc);
heap.push(MergeEntry {
sort_key: key,
rec: next_t,
reader_idx,
});
}
Ok(None) => { }
Err(e) => return Some(Err(e)),
}
Some(Ok(rec))
}
}
}
}
fn read_one_record<R: Read, const N: usize>(r: &mut R) -> Result<Option<[u32; N]>, BuildError> {
let mut buf = [0u8; MAX_RECORD_BYTES];
let rec_bytes = mem::size_of::<[u32; N]>();
match r.read_exact(&mut buf[..rec_bytes]) {
Ok(()) => {}
Err(e) if e.kind() == ErrorKind::UnexpectedEof => return Ok(None),
Err(e) => return Err(BuildError::Io(e)),
}
let mut rec = [0u32; N];
for (lane, slot) in rec.iter_mut().enumerate() {
let at = lane * 4;
*slot = u32::from_le_bytes(buf[at..at + 4].try_into().expect("slice len 4"));
}
Ok(Some(rec))
}
const MAX_RECORD_BYTES: usize = mem::size_of::<[u32; 5]>();
fn write_records_sorted<const N: usize>(
triples: &[[u32; N]],
path: &Path,
) -> Result<(), BuildError> {
let mut w = BufWriter::with_capacity(PARTITION_BUF_SIZE, File::create(path)?);
#[cfg(target_endian = "little")]
{
let bytes: &[u8] = bytemuck::cast_slice(triples);
w.write_all(bytes)?;
}
#[cfg(not(target_endian = "little"))]
{
for t in triples {
write_record(&mut w, t)?;
}
}
w.flush()?;
Ok(())
}
fn spill_sorted_chunk<const N: usize>(
chunk: &mut Vec<[u32; N]>,
scratch_dir: &Path,
partition_label: &str,
chunk_idx: usize,
lex_rank: &[u32],
out_paths: &mut Vec<PathBuf>,
) -> Result<(), BuildError> {
radix_sort_records_by_lex_rank(chunk, lex_rank);
let path = scratch_dir.join(format!("{partition_label}_sorted{chunk_idx}.bin"));
write_records_sorted(chunk, &path)?;
chunk.clear();
#[cfg(test)]
finish_debug::record_chunk_path(&path);
out_paths.push(path);
Ok(())
}
#[cfg(test)]
mod finish_debug {
use std::{
cell::RefCell,
path::{Path, PathBuf},
};
thread_local! {
static OBSERVED_CHUNKS: RefCell<Vec<PathBuf>> = const { RefCell::new(Vec::new()) };
}
pub fn reset() {
OBSERVED_CHUNKS.with(|c| c.borrow_mut().clear());
}
pub fn record_chunk_path(path: &Path) {
OBSERVED_CHUNKS.with(|c| c.borrow_mut().push(path.to_path_buf()));
}
pub fn observed() -> Vec<PathBuf> {
OBSERVED_CHUNKS.with(|c| c.borrow().clone())
}
}
fn radix_sort_records_by_lex_rank<const N: usize>(triples: &mut Vec<[u32; N]>, lex_rank: &[u32]) {
let n = triples.len();
if n < RADIX_SORT_MIN_TRIPLES {
triples.sort_unstable_by(|a, b| {
lex_rank[triple_term_id(a) as usize]
.cmp(&lex_rank[triple_term_id(b) as usize])
.then(triple_doc_id(a).cmp(&triple_doc_id(b)))
});
return;
}
let vocab_size = lex_rank.len();
let mut offsets: Vec<u32> = vec![0u32; vocab_size + 1];
for t in triples.iter() {
let rank = unsafe { *lex_rank.get_unchecked(t[0] as usize) } as usize;
offsets[rank] = offsets[rank].wrapping_add(1);
}
let mut sum: u32 = 0;
for c in offsets.iter_mut() {
let tmp = *c;
*c = sum;
sum = sum.wrapping_add(tmp);
}
debug_assert_eq!(sum as usize, n, "histogram total != triple count");
let mut out: Vec<[u32; N]> = vec![[0u32; N]; n];
for t in triples.iter() {
let rank = unsafe { *lex_rank.get_unchecked(t[0] as usize) } as usize;
let dst = unsafe { *offsets.get_unchecked(rank) } as usize;
unsafe {
*out.get_unchecked_mut(dst) = *t;
*offsets.get_unchecked_mut(rank) = (dst as u32).wrapping_add(1);
}
}
*triples = out;
}
fn open_partition_sorted<const N: usize>(
partition_path: &Path,
max_partition_bytes: u64,
scratch_dir: &Path,
partition_label: &str,
lex_rank: &[u32],
) -> Result<PartitionIter<N>, BuildError> {
let len = fs::metadata(partition_path)?.len();
if len <= max_partition_bytes {
let mut triples = read_partition_records::<N>(partition_path)?;
radix_sort_records_by_lex_rank(&mut triples, lex_rank);
return Ok(PartitionIter::InMemory(triples.into_iter()));
}
let chunk_triples = (max_partition_bytes as usize) / mem::size_of::<[u32; N]>();
let mut sorted_chunk_paths: Vec<PathBuf> = Vec::new();
let mut r = BufReader::with_capacity(PARTITION_BUF_SIZE, File::open(partition_path)?);
let mut chunk: Vec<[u32; N]> =
Vec::with_capacity(chunk_triples.min(EXTERNAL_MERGE_CHUNK_CAP_TRIPLES));
let mut chunk_idx: usize = 0;
while let Some(t) = read_one_record::<_, N>(&mut r)? {
chunk.push(t);
if chunk.len() >= chunk_triples {
spill_sorted_chunk(
&mut chunk,
scratch_dir,
partition_label,
chunk_idx,
lex_rank,
&mut sorted_chunk_paths,
)?;
chunk_idx += 1;
}
}
if !chunk.is_empty() {
spill_sorted_chunk(
&mut chunk,
scratch_dir,
partition_label,
chunk_idx,
lex_rank,
&mut sorted_chunk_paths,
)?;
}
let mut readers: Vec<BufReader<File>> = Vec::with_capacity(sorted_chunk_paths.len());
for p in &sorted_chunk_paths {
readers.push(BufReader::with_capacity(PARTITION_BUF_SIZE, File::open(p)?));
}
let mut heap: BinaryHeap<MergeEntry<N>> = BinaryHeap::with_capacity(readers.len());
for (idx, reader) in readers.iter_mut().enumerate() {
if let Some(t) = read_one_record::<_, N>(reader)? {
let term_id = triple_term_id(&t);
let doc_id = triple_doc_id(&t);
heap.push(MergeEntry {
sort_key: pack_sort_key(lex_rank[term_id as usize], doc_id),
rec: t,
reader_idx: idx,
});
}
}
Ok(PartitionIter::Merge {
readers,
heap,
_chunk_paths: sorted_chunk_paths,
})
}
pub struct FtsBuilder {
default_tokenizer: Arc<dyn Tokenizer>,
column_tokenizers: Vec<Arc<dyn Tokenizer>>,
columns: Vec<ColumnState>,
postings: Vec<ColumnPostings>,
scratch_dir: tempfile::TempDir,
spill_threshold_bytes: usize,
spill_partitions: usize,
max_partition_bytes: u64,
n_docs: u32,
doc_tf: DocTfMap,
doc_pos_head: DocPosHeadMap,
doc_pos_chain: Vec<(u32, u32)>,
pos_scratch: Vec<u32>,
run_scratch: Vec<u8>,
bump: Bump,
}
impl FtsBuilder {
pub fn new(tokenizer: Arc<dyn Tokenizer>) -> Self {
let scratch_dir = tempfile::tempdir().expect("create FtsBuilder scratch tempdir");
Self::from_parts(tokenizer, scratch_dir)
}
pub fn with_scratch(
tokenizer: Arc<dyn Tokenizer>,
scratch: PathBuf,
) -> Result<Self, BuildError> {
let scratch_dir = tempfile::Builder::new()
.prefix("infino-fts-")
.tempdir_in(&scratch)?;
Ok(Self::from_parts(tokenizer, scratch_dir))
}
fn from_parts(tokenizer: Arc<dyn Tokenizer>, scratch_dir: tempfile::TempDir) -> Self {
Self {
default_tokenizer: tokenizer,
column_tokenizers: Vec::new(),
columns: Vec::new(),
postings: Vec::new(),
scratch_dir,
spill_threshold_bytes: DEFAULT_SPILL_THRESHOLD_BYTES,
spill_partitions: DEFAULT_SPILL_PARTITIONS,
max_partition_bytes: DEFAULT_MAX_PARTITION_BYTES,
n_docs: 0,
doc_tf: DocTfMap::with_hasher(FxBuildHasher),
doc_pos_head: DocPosHeadMap::with_hasher(FxBuildHasher),
doc_pos_chain: Vec::new(),
pos_scratch: Vec::new(),
run_scratch: Vec::new(),
bump: Bump::new(),
}
}
pub fn set_spill_threshold_bytes(&mut self, threshold: usize) {
assert!(
threshold > 0,
"FtsBuilder: spill_threshold_bytes must be > 0"
);
self.spill_threshold_bytes = threshold;
}
pub fn set_spill_partitions(&mut self, n: usize) -> Result<(), BuildError> {
if !self.columns.is_empty() {
return Err(BuildError::Io(Error::new(
ErrorKind::InvalidInput,
"FtsBuilder::set_spill_partitions must be called before any register_column",
)));
}
if n == 0 {
return Err(BuildError::Io(Error::new(
ErrorKind::InvalidInput,
"FtsBuilder: spill_partitions must be ≥ 1",
)));
}
if !n.is_power_of_two() {
return Err(BuildError::Io(Error::new(
ErrorKind::InvalidInput,
format!("FtsBuilder: spill_partitions must be a power of two; got {n}"),
)));
}
self.spill_partitions = n;
Ok(())
}
pub fn set_max_partition_bytes(&mut self, bytes: u64) {
assert!(bytes > 0, "FtsBuilder: max_partition_bytes must be > 0");
self.max_partition_bytes = bytes;
}
pub fn register_column(&mut self, name: String, positions: bool) -> Result<u32, BuildError> {
let tokenizer = Arc::clone(&self.default_tokenizer);
self.register_column_with_tokenizer(name, positions, tokenizer)
}
pub fn register_column_with_tokenizer(
&mut self,
name: String,
positions: bool,
tokenizer: Arc<dyn Tokenizer>,
) -> Result<u32, BuildError> {
if name.as_bytes().contains(&FST_SEPARATOR) {
return Err(BuildError::ReservedSeparatorInColumnName(name));
}
if name.starts_with(format::RESERVED_PREFIX) {
return Err(BuildError::ReservedPrefixInColumnName(name));
}
if self.columns.iter().any(|c| c.name == name) {
return Err(BuildError::DuplicateColumnName(name));
}
let column_id = self.columns.len() as u32;
self.columns.push(ColumnState {
name,
doc_lengths: Vec::new(),
total_tokens: 0,
positions,
});
self.postings.push(ColumnPostings::new());
self.column_tokenizers.push(tokenizer);
Ok(column_id)
}
fn open_partitions_for_column<const N: usize>(
scratch_dir: &Path,
column_id: u32,
n_partitions: usize,
) -> Result<Vec<SpillPartition<N>>, BuildError> {
let mut partitions = Vec::with_capacity(n_partitions);
for partition in 0..n_partitions {
let path = scratch_dir.join(format!("fts_col{column_id}_part{partition}.bin"));
let file = File::create(&path)?;
partitions.push(SpillPartition {
path,
writer: Some(BufWriter::with_capacity(PARTITION_BUF_SIZE, file)),
batch: Vec::with_capacity(SPILL_BATCH_TRIPLES),
});
}
Ok(partitions)
}
fn open_position_blobs_for_column(
scratch_dir: &Path,
column_id: u32,
n_partitions: usize,
) -> Result<Vec<PartitionPositions>, BuildError> {
let mut blobs = Vec::with_capacity(n_partitions);
for partition in 0..n_partitions {
let path = scratch_dir.join(format!("fts_col{column_id}_part{partition}.pos.bin"));
let file = File::create(&path)?;
blobs.push(PartitionPositions {
path,
writer: Some(BufWriter::with_capacity(PARTITION_BUF_SIZE, file)),
len: 0,
});
}
Ok(blobs)
}
#[allow(clippy::too_many_arguments)]
fn flush_in_ram_to_positional_partitions(
terms: FxHashMap<Box<str>, Vec<(u32, u32)>>,
mut pos_runs: FxHashMap<Box<str>, Vec<u8>>,
partitions: &mut [SpillPartition<POSITIONAL_RECORD_LANES>],
blobs: &mut [PartitionPositions],
term_to_id: &mut TermIdMap,
id_to_term: &mut Vec<&'static str>,
arena: &Bump,
) -> Result<(), BuildError> {
let n_part = partitions.len();
debug_assert_eq!(n_part, blobs.len(), "one blob per partition");
debug_assert!(
n_part.is_power_of_two(),
"spill_partitions must be a power of 2; got {n_part}"
);
let mask = n_part - 1;
for (term, postings) in terms {
let (term_id, _is_new) = intern_term_id(term_to_id, id_to_term, arena, &term);
let p = (term_id as usize) & mask;
let runs = pos_runs
.remove(&term)
.expect("positional term accumulated a position run");
let blob = &mut blobs[p];
let writer = blob
.writer
.as_mut()
.expect("positions blob writer is open before finish");
let mut at: usize = 0;
for (doc_id, tf) in postings {
let run_start = at;
skip_run(&runs, &mut at, tf).expect("builder-encoded runs are well-formed");
writer.write_all(&runs[run_start..at])?;
let pos_off = blob.len;
blob.len += (at - run_start) as u64;
let (lo, hi) = pos_off_lanes(pos_off);
push_record_batched(&mut partitions[p], [term_id, doc_id, tf, lo, hi])?;
}
debug_assert_eq!(at, runs.len(), "runs must cover exactly the pairs");
}
Ok(())
}
fn flush_in_ram_to_partitions(
terms: FxHashMap<Box<str>, Vec<(u32, u32)>>,
partitions: &mut [SpillPartition<PLAIN_RECORD_LANES>],
term_to_id: &mut TermIdMap,
id_to_term: &mut Vec<&'static str>,
arena: &Bump,
) -> Result<(), BuildError> {
let n_part = partitions.len();
debug_assert!(
n_part.is_power_of_two(),
"spill_partitions must be a power of 2; got {n_part}"
);
let mask = n_part - 1;
for (term, postings) in terms {
let (term_id, _is_new) = intern_term_id(term_to_id, id_to_term, arena, &term);
let p = (term_id as usize) & mask;
for (doc_id, tf) in postings {
push_record_batched(&mut partitions[p], [term_id, doc_id, tf])?;
}
}
Ok(())
}
pub fn add_doc(
&mut self,
column_id: u32,
local_doc_id: u32,
text: &str,
) -> Result<(), BuildError> {
let col_idx = column_id as usize;
if col_idx >= self.columns.len() {
return Err(BuildError::FtsColumnTypeInvalid {
column: format!("(unregistered column_id {column_id})"),
actual: "n/a".to_string(),
});
}
debug_assert!(
local_doc_id as usize == self.columns[col_idx].doc_lengths.len(),
"FtsBuilder::add_doc: local_doc_id ({local_doc_id}) must equal \
this column's next index ({}); doc_ids must be consecutive \
from 0 within a column",
self.columns[col_idx].doc_lengths.len(),
);
if self.postings[col_idx].is_spilled() {
self.add_doc_spilled(col_idx, local_doc_id, text)
} else {
self.add_doc_inram(col_idx, local_doc_id, text)
}
}
#[inline(always)]
fn add_doc_spilled(
&mut self,
col_idx: usize,
local_doc_id: u32,
text: &str,
) -> Result<(), BuildError> {
let tokenizer = &self.column_tokenizers[col_idx];
let ascii_tok = tokenizer
.as_ref()
.as_any()
.downcast_ref::<AsciiLowerTokenizer>();
let mut tokens_in_doc: u64 = 0;
let positional = self.columns[col_idx].positions;
let col_post = &mut self.postings[col_idx];
let (
store,
term_to_id,
id_to_term,
dense_doc_tf,
dense_doc_poshead,
updated_terms,
term_arena,
) = match col_post {
ColumnPostings::Spilled {
partitions,
term_to_id,
id_to_term,
dense_doc_tf,
dense_doc_poshead,
updated_terms,
term_arena,
} => (
partitions,
term_to_id,
id_to_term,
dense_doc_tf,
dense_doc_poshead,
updated_terms,
term_arena,
),
ColumnPostings::InRam { .. } => {
unreachable!("add_doc_spilled called on InRam column")
}
};
let n_part = store.n_partitions();
debug_assert!(
n_part.is_power_of_two(),
"spill_partitions must be a power of 2"
);
let mask = n_part - 1;
updated_terms.clear();
let mut pos_overflow = false;
if !positional {
let mut on_token = |tok: &str| {
tokens_in_doc += 1;
let (term_id, is_new) = intern_term_id(term_to_id, id_to_term, term_arena, tok);
let idx = term_id as usize;
if is_new {
debug_assert_eq!(idx, dense_doc_tf.len());
dense_doc_tf.push(0);
}
let slot = unsafe { dense_doc_tf.get_unchecked_mut(idx) };
if *slot == 0 {
updated_terms.push(term_id);
}
*slot += 1;
};
if let Some(ascii) = ascii_tok {
ascii.tokenize_each_inline(text, &mut on_token);
} else {
tokenizer.tokenize_each(text, &mut on_token);
}
} else {
self.doc_pos_chain.clear();
let doc_pos_chain = &mut self.doc_pos_chain;
let mut record = |tok: &str, position: u64| {
let (term_id, is_new) = intern_term_id(term_to_id, id_to_term, term_arena, tok);
let idx = term_id as usize;
if is_new {
debug_assert_eq!(idx, dense_doc_tf.len());
dense_doc_tf.push(0);
dense_doc_poshead.push(CHAIN_END);
}
let slot = unsafe { dense_doc_tf.get_unchecked_mut(idx) };
if *slot == 0 {
updated_terms.push(term_id);
}
*slot += 1;
if position > u32::MAX as u64 {
pos_overflow = true;
return;
}
let chain_idx = doc_pos_chain.len() as u32;
let head = unsafe { dense_doc_poshead.get_unchecked_mut(idx) };
doc_pos_chain.push((position as u32, *head));
*head = chain_idx;
};
if let Some(ascii) = ascii_tok {
ascii.tokenize_each_inline_positioned(text, |tok, position| {
record(tok, position);
tokens_in_doc += 1;
});
} else {
tokenizer.tokenize_each(text, &mut |tok| {
record(tok, tokens_in_doc);
tokens_in_doc += 1;
});
}
}
if pos_overflow {
return Err(BuildError::PositionOverflow {
column: self.columns[col_idx].name.clone(),
});
}
let col = &mut self.columns[col_idx];
let dl_clamped: u32 = tokens_in_doc.min(u32::MAX as u64) as u32;
col.doc_lengths.push(dl_clamped);
col.total_tokens = col.total_tokens.saturating_add(tokens_in_doc);
let docs_now = local_doc_id.saturating_add(1);
if docs_now > self.n_docs {
self.n_docs = docs_now;
}
match store {
SpillStore::Plain(partitions) => {
for &term_id in updated_terms.iter() {
let idx = term_id as usize;
let slot = unsafe { dense_doc_tf.get_unchecked_mut(idx) };
let tf = *slot;
*slot = 0;
let p = (term_id as usize) & mask;
let partition = unsafe { partitions.get_unchecked_mut(p) };
push_record_batched(partition, [term_id, local_doc_id, tf])?;
}
}
SpillStore::Positional { partitions, blobs } => {
let doc_pos_chain = &self.doc_pos_chain;
let pos_scratch = &mut self.pos_scratch;
let run_buf = &mut self.run_scratch;
for &term_id in updated_terms.iter() {
let idx = term_id as usize;
let slot = unsafe { dense_doc_tf.get_unchecked_mut(idx) };
let tf = *slot;
*slot = 0;
let head_slot = unsafe { dense_doc_poshead.get_unchecked_mut(idx) };
let head = *head_slot;
*head_slot = CHAIN_END;
pos_scratch.clear();
let mut at = head;
while at != CHAIN_END {
let (pos, prev) = doc_pos_chain[at as usize];
pos_scratch.push(pos);
at = prev;
}
pos_scratch.reverse();
debug_assert_eq!(pos_scratch.len() as u32, tf, "chain length must equal tf");
run_buf.clear();
encode_run(run_buf, pos_scratch);
let p = (term_id as usize) & mask;
let blob = &mut blobs[p];
let pos_off = blob.len;
blob.writer
.as_mut()
.expect("positions blob writer is open before finish")
.write_all(run_buf)?;
blob.len += run_buf.len() as u64;
let (lo, hi) = pos_off_lanes(pos_off);
let partition = unsafe { partitions.get_unchecked_mut(p) };
push_record_batched(partition, [term_id, local_doc_id, tf, lo, hi])?;
}
}
}
Ok(())
}
#[inline(always)]
fn add_doc_inram(
&mut self,
col_idx: usize,
local_doc_id: u32,
text: &str,
) -> Result<(), BuildError> {
let tokenizer = &self.column_tokenizers[col_idx];
let ascii_tok = tokenizer
.as_ref()
.as_any()
.downcast_ref::<AsciiLowerTokenizer>();
let mut tokens_in_doc: u64 = 0;
let positional = self.columns[col_idx].positions;
self.doc_tf.clear();
self.bump.reset();
let bump = &self.bump;
let tf_per_term = &mut self.doc_tf;
if !positional {
let mut on_token = |tok: &str| {
tokens_in_doc += 1;
let hash = compute_hash(tf_per_term.hasher(), tok);
match tf_per_term
.raw_entry_mut()
.from_hash(hash, |existing| *existing == tok)
{
RawEntryMut::Occupied(mut e) => {
*e.get_mut() += 1;
}
RawEntryMut::Vacant(e) => {
let bumped: &str = bump.alloc_str(tok);
let extended: &'static str = unsafe { std::mem::transmute(bumped) };
e.insert_hashed_nocheck(hash, extended, 1);
}
}
};
if let Some(ascii) = ascii_tok {
ascii.tokenize_each_inline(text, &mut on_token);
} else {
tokenizer.tokenize_each(text, &mut on_token);
}
} else {
self.doc_pos_head.clear();
self.doc_pos_chain.clear();
let doc_pos_head = &mut self.doc_pos_head;
let doc_pos_chain = &mut self.doc_pos_chain;
let mut pos_overflow = false;
let mut record = |tok: &str, position: u64| {
let hash = compute_hash(tf_per_term.hasher(), tok);
let key: &'static str = match tf_per_term
.raw_entry_mut()
.from_hash(hash, |existing| *existing == tok)
{
RawEntryMut::Occupied(mut e) => {
*e.get_mut() += 1;
e.key()
}
RawEntryMut::Vacant(e) => {
let bumped: &str = bump.alloc_str(tok);
let extended: &'static str = unsafe { std::mem::transmute(bumped) };
e.insert_hashed_nocheck(hash, extended, 1);
extended
}
};
if position > u32::MAX as u64 {
pos_overflow = true;
return;
}
let idx = doc_pos_chain.len() as u32;
let prev = doc_pos_head.insert(key, idx).unwrap_or(CHAIN_END);
doc_pos_chain.push((position as u32, prev));
};
if let Some(ascii) = ascii_tok {
ascii.tokenize_each_inline_positioned(text, |tok, position| {
record(tok, position);
tokens_in_doc += 1;
});
} else {
tokenizer.tokenize_each(text, &mut |tok| {
record(tok, tokens_in_doc);
tokens_in_doc += 1;
});
}
if pos_overflow {
return Err(BuildError::PositionOverflow {
column: self.columns[col_idx].name.clone(),
});
}
}
let col = &mut self.columns[col_idx];
let dl_clamped: u32 = tokens_in_doc.min(u32::MAX as u64) as u32;
col.doc_lengths.push(dl_clamped);
col.total_tokens = col.total_tokens.saturating_add(tokens_in_doc);
let docs_now = local_doc_id.saturating_add(1);
if docs_now > self.n_docs {
self.n_docs = docs_now;
}
let column_id = col_idx as u32;
let col_post = &mut self.postings[col_idx];
let (terms, pos_runs, bytes) = match col_post {
ColumnPostings::InRam {
terms,
pos_runs,
bytes,
} => (terms, pos_runs, bytes),
ColumnPostings::Spilled { .. } => {
unreachable!("add_doc_inram called on Spilled column")
}
};
let doc_pos_head = &self.doc_pos_head;
let doc_pos_chain = &self.doc_pos_chain;
let pos_scratch = &mut self.pos_scratch;
let mut new_bytes: usize = 0;
for (term, tf) in tf_per_term.drain() {
let term_len = term.len();
match terms.get_mut(term) {
Some(acc) => {
acc.push((local_doc_id, tf));
new_bytes = new_bytes.saturating_add(ACCUM_POSTING_BYTES);
}
None => {
terms.insert(Box::<str>::from(term), vec![(local_doc_id, tf)]);
new_bytes = new_bytes.saturating_add(
ACCUM_NEW_TERM_FIXED_BYTES + term_len + ACCUM_POSTING_BYTES,
);
}
}
if positional {
let head = doc_pos_head.get(term).copied().unwrap_or(CHAIN_END);
debug_assert_ne!(head, CHAIN_END, "tf term missing from position chain");
pos_scratch.clear();
let mut at = head;
while at != CHAIN_END {
let (p, prev) = doc_pos_chain[at as usize];
pos_scratch.push(p);
at = prev;
}
pos_scratch.reverse();
debug_assert_eq!(pos_scratch.len() as u32, tf, "chain length must equal tf");
match pos_runs.get_mut(term) {
Some(run) => {
let before = run.len();
encode_run(run, pos_scratch);
new_bytes = new_bytes.saturating_add(run.len() - before);
}
None => {
let mut run = Vec::new();
encode_run(&mut run, pos_scratch);
new_bytes = new_bytes
.saturating_add(ACCUM_NEW_TERM_FIXED_BYTES + term_len + run.len());
pos_runs.insert(Box::<str>::from(term), run);
}
}
}
}
let new_total = bytes.saturating_add(new_bytes);
if new_total > self.spill_threshold_bytes {
let drained = mem::take(terms);
let drained_pos_runs = mem::take(pos_runs);
let term_arena = Bump::new();
let mut term_to_id: TermIdMap = TermIdMap::default();
let mut id_to_term: Vec<&'static str> = Vec::with_capacity(drained.len());
let store = match positional {
false => {
let mut partitions = Self::open_partitions_for_column(
self.scratch_dir.path(),
column_id,
self.spill_partitions,
)?;
Self::flush_in_ram_to_partitions(
drained,
&mut partitions,
&mut term_to_id,
&mut id_to_term,
&term_arena,
)?;
SpillStore::Plain(partitions)
}
true => {
let mut partitions = Self::open_partitions_for_column(
self.scratch_dir.path(),
column_id,
self.spill_partitions,
)?;
let mut blobs = Self::open_position_blobs_for_column(
self.scratch_dir.path(),
column_id,
self.spill_partitions,
)?;
Self::flush_in_ram_to_positional_partitions(
drained,
drained_pos_runs,
&mut partitions,
&mut blobs,
&mut term_to_id,
&mut id_to_term,
&term_arena,
)?;
SpillStore::Positional { partitions, blobs }
}
};
let dense_doc_tf = vec![0u32; id_to_term.len()];
let dense_doc_poshead = match positional {
true => vec![CHAIN_END; id_to_term.len()],
false => Vec::new(),
};
let updated_terms: Vec<u32> = Vec::new();
*col_post = ColumnPostings::Spilled {
partitions: store,
term_to_id,
id_to_term,
dense_doc_tf,
dense_doc_poshead,
updated_terms,
term_arena,
};
} else {
*bytes = new_total;
}
Ok(())
}
pub fn finish(self) -> Result<Vec<u8>, BuildError> {
let mut blob = Vec::new();
self.finish_to(&mut blob)?;
Ok(blob)
}
pub fn finish_to<W: Write>(self, w: W) -> Result<(), BuildError> {
if self.postings.iter().any(|c| c.is_spilled()) {
self.finish_to_spilled(w)
} else {
self.finish_to_inram(w)
}
}
fn finish_to_inram<W: Write>(self, mut w: W) -> Result<(), BuildError> {
let FtsBuilder {
default_tokenizer: _,
column_tokenizers: _,
columns,
postings,
scratch_dir,
spill_threshold_bytes: _,
spill_partitions: _,
max_partition_bytes: _,
n_docs,
doc_tf,
doc_pos_head,
doc_pos_chain: _,
pos_scratch: _,
run_scratch: _,
bump,
} = self;
drop(doc_tf);
drop(doc_pos_head);
drop(bump);
let n_columns = columns.len() as u32;
let mut n_terms_total_usize: usize = 0;
let mut work: Vec<(usize, ColumnState, ColumnPostings)> = columns
.into_iter()
.zip(postings)
.enumerate()
.map(|(orig_idx, (state, posting_state))| (orig_idx, state, posting_state))
.collect();
work.sort_unstable_by(|a, b| a.1.name.cmp(&b.1.name));
let mut avgdl_per_col: Vec<f32> = vec![0.0; n_columns as usize];
for (orig_idx, state, _) in &work {
let n = state.doc_lengths.len() as u64;
avgdl_per_col[*orig_idx] = if n == 0 {
0.0
} else {
(state.total_tokens as f32) / (n as f32)
};
}
let scratch_path = scratch_dir.path().to_path_buf();
let postings_path = scratch_path.join("infino_fts_postings.bin");
let mut postings_writer = BufWriter::new(File::create(&postings_path)?);
let mut postings_len: u64 = 0;
let mut postings_crc_acc: u32 = 0;
let mut positions_sink = PositionsSink::create(&scratch_path)?;
let mut key_buf: Vec<u8> = Vec::with_capacity(64);
let mut term_scratch = TermScratch::default();
let mut finish_profile = FinishProfile::from_config();
let mut fst_inram = DictBuilder::new();
let mut doc_lengths_by_orig_col: Vec<Option<Vec<u32>>> =
(0..n_columns as usize).map(|_| None).collect();
for (orig_col_idx, col_state, posting_state) in work.drain(..) {
let ColumnState {
name: col_name,
doc_lengths: col_doc_lengths_owned,
total_tokens: _,
positions: col_positions,
} = col_state;
let col_name_bytes = col_name.as_bytes();
let avgdl = avgdl_per_col[orig_col_idx];
let col_doc_lengths: &[u32] = &col_doc_lengths_owned;
let (terms, mut pos_runs) = match posting_state {
ColumnPostings::InRam {
terms,
pos_runs,
bytes: _,
} => (terms, pos_runs),
ColumnPostings::Spilled { .. } => unreachable!(
"finish_to_inram dispatched on !any_spilled; \
Spilled column cannot appear here"
),
};
type InRamEntries = Vec<(Box<str>, Vec<(u32, u32)>)>;
let mut entries: InRamEntries = terms.into_iter().collect();
entries.sort_unstable_by(|a, b| a.0.as_bytes().cmp(b.0.as_bytes()));
for (term, postings) in entries {
let term_runs: Vec<u8> = match col_positions {
true => pos_runs
.remove(&term)
.expect("positional term accumulated a position run"),
false => Vec::new(),
};
let term_positions = match col_positions {
true => Some((&mut positions_sink, term_runs.as_slice())),
false => None,
};
encode_and_emit_term(
&term,
&postings,
col_name_bytes,
col_doc_lengths,
avgdl,
n_docs,
&mut key_buf,
&mut postings_writer,
&mut postings_crc_acc,
&mut postings_len,
Some(&mut fst_inram),
None,
term_positions,
&mut finish_profile,
&mut term_scratch,
)?;
n_terms_total_usize += 1;
}
doc_lengths_by_orig_col[orig_col_idx] = Some(col_doc_lengths_owned);
}
assemble_and_write_blob(
BlobAssemblyInputs {
postings_writer,
postings_path,
postings_crc_acc,
postings_len,
positions_sink,
fst_sink: FstSinkFinish::InRam(fst_inram),
n_columns,
n_docs,
n_terms_total_usize,
avgdl_per_col,
doc_lengths_by_orig_col,
scratch_dir,
finish_profile,
},
&mut w,
)
}
fn finish_to_spilled<W: Write>(self, mut w: W) -> Result<(), BuildError> {
let FtsBuilder {
default_tokenizer: _,
column_tokenizers: _,
columns,
postings,
scratch_dir,
spill_threshold_bytes: _,
spill_partitions: _,
max_partition_bytes,
n_docs,
doc_tf,
doc_pos_head,
doc_pos_chain: _,
pos_scratch: _,
run_scratch: _,
bump,
} = self;
drop(doc_tf);
drop(doc_pos_head);
drop(bump);
let n_columns = columns.len() as u32;
let mut n_terms_total_usize: usize = 0;
let mut work: Vec<(usize, ColumnState, ColumnPostings)> = columns
.into_iter()
.zip(postings)
.enumerate()
.map(|(orig_idx, (state, posting_state))| (orig_idx, state, posting_state))
.collect();
work.sort_unstable_by(|a, b| a.1.name.cmp(&b.1.name));
let mut avgdl_per_col: Vec<f32> = vec![0.0; n_columns as usize];
for (orig_idx, state, _) in &work {
let n = state.doc_lengths.len() as u64;
avgdl_per_col[*orig_idx] = if n == 0 {
0.0
} else {
(state.total_tokens as f32) / (n as f32)
};
}
let scratch_path = scratch_dir.path().to_path_buf();
let postings_path = scratch_path.join("infino_fts_postings.bin");
let mut postings_writer = BufWriter::new(File::create(&postings_path)?);
let mut postings_len: u64 = 0;
let mut postings_crc_acc: u32 = 0;
let mut positions_sink = PositionsSink::create(&scratch_path)?;
let mut key_buf: Vec<u8> = Vec::with_capacity(64);
let mut term_scratch = TermScratch::default();
let mut finish_profile = FinishProfile::from_config();
let fst_streaming_path = scratch_path.join("infino_fts_dict.bin");
let mut fst_streaming = {
let fst_file = File::create(&fst_streaming_path)?;
let bw = BufWriter::new(fst_file);
StreamingDictBuilder::new(bw).map_err(map_fst_err)?
};
let partition_flush_start = finish_profile.enabled.then(Instant::now);
for (_, _, cp) in &mut work {
if let ColumnPostings::Spilled { partitions, .. } = cp {
match partitions {
SpillStore::Plain(parts) => {
for partition in parts {
flush_partition_batch(partition)?;
if let Some(mut writer) = partition.writer.take() {
writer.flush()?;
}
}
}
SpillStore::Positional { partitions, blobs } => {
for partition in partitions.iter_mut() {
flush_partition_batch(partition)?;
if let Some(mut writer) = partition.writer.take() {
writer.flush()?;
}
}
for blob in blobs.iter_mut() {
if let Some(mut writer) = blob.writer.take() {
writer.flush()?;
}
}
}
}
}
}
if let Some(t) = partition_flush_start {
finish_profile.partition_flush += t.elapsed();
}
let mut doc_lengths_by_orig_col: Vec<Option<Vec<u32>>> =
(0..n_columns as usize).map(|_| None).collect();
for (orig_col_idx, col_state, posting_state) in work.drain(..) {
let ColumnState {
name: col_name,
doc_lengths: col_doc_lengths_owned,
total_tokens: _,
positions: col_positions,
} = col_state;
let col_name_bytes = col_name.as_bytes();
let avgdl = avgdl_per_col[orig_col_idx];
let col_doc_lengths: &[u32] = &col_doc_lengths_owned;
match posting_state {
ColumnPostings::InRam {
terms,
pos_runs: mut col_pos_runs,
bytes: _,
} => {
type InRamEntries = Vec<(Box<str>, Vec<(u32, u32)>)>;
let mut entries: InRamEntries = terms.into_iter().collect();
entries.sort_unstable_by(|a, b| a.0.as_bytes().cmp(b.0.as_bytes()));
for (term, postings) in entries {
let term_runs: Vec<u8> = match col_positions {
true => col_pos_runs
.remove(&term)
.expect("positional term accumulated a position run"),
false => Vec::new(),
};
let term_positions = match col_positions {
true => Some((&mut positions_sink, term_runs.as_slice())),
false => None,
};
encode_and_emit_term(
&term,
&postings,
col_name_bytes,
col_doc_lengths,
avgdl,
n_docs,
&mut key_buf,
&mut postings_writer,
&mut postings_crc_acc,
&mut postings_len,
None,
Some(&mut fst_streaming),
term_positions,
&mut finish_profile,
&mut term_scratch,
)?;
n_terms_total_usize += 1;
}
}
ColumnPostings::Spilled {
partitions,
term_to_id,
id_to_term,
dense_doc_tf: _,
dense_doc_poshead: _,
updated_terms: _,
term_arena,
} => {
drop(term_to_id);
let lex_rank_start = finish_profile.enabled.then(Instant::now);
let (lex_rank, term_id_in_lex_order) = build_lex_rank(&id_to_term);
if let Some(t) = lex_rank_start {
finish_profile.lex_rank_build += t.elapsed();
}
let sort_start = finish_profile.enabled.then(Instant::now);
let mut sorted_files: Vec<PathBuf> =
Vec::with_capacity(partitions.n_partitions());
let partition_paths: Vec<PathBuf> = match &partitions {
SpillStore::Plain(parts) => parts.iter().map(|p| p.path.clone()).collect(),
SpillStore::Positional {
partitions: parts, ..
} => parts.iter().map(|p| p.path.clone()).collect(),
};
for (partition_idx, partition_path) in partition_paths.iter().enumerate() {
let sorted_path = scratch_path.join(format!(
"fts_col{orig_col_idx}_part{partition_idx}.sorted.bin"
));
match &partitions {
SpillStore::Plain(_) => sort_partition_to_file::<PLAIN_RECORD_LANES>(
partition_path,
&sorted_path,
max_partition_bytes,
&scratch_path,
&format!("c{orig_col_idx}_p{partition_idx}"),
&lex_rank,
)?,
SpillStore::Positional { .. } => {
sort_partition_to_file::<POSITIONAL_RECORD_LANES>(
partition_path,
&sorted_path,
max_partition_bytes,
&scratch_path,
&format!("c{orig_col_idx}_p{partition_idx}"),
&lex_rank,
)?
}
}
sorted_files.push(sorted_path);
}
if let Some(t) = sort_start {
finish_profile.partition_sort += t.elapsed();
}
let merge_profile_start = Instant::now();
let encode_calls_before = finish_profile.encode_calls;
let encode_df1_before = finish_profile.encode_df1;
let encode_pfor_before = finish_profile.encode_pfor;
let encode_total_before = finish_profile.encode_total;
let encode_block_build_before = finish_profile.encode_block_build;
let encode_meta_write_before = finish_profile.encode_meta_write;
let encode_skip_write_before = finish_profile.encode_skip_write;
let encode_block_write_before = finish_profile.encode_block_write;
let fst_insert_before = finish_profile.fst_insert;
let n_emitted = match &partitions {
SpillStore::Plain(_) => merge_sorted_spill::<PLAIN_RECORD_LANES, _>(
&sorted_files,
None,
&term_id_in_lex_order,
&id_to_term,
col_name_bytes,
col_doc_lengths,
avgdl,
n_docs,
&mut key_buf,
&mut postings_writer,
&mut postings_crc_acc,
&mut postings_len,
&mut fst_streaming,
&mut positions_sink,
&mut finish_profile,
&mut term_scratch,
)?,
SpillStore::Positional { blobs, .. } => {
let mut blob_mmaps: Vec<Option<Mmap>> = Vec::with_capacity(blobs.len());
for blob in blobs {
let f = File::open(&blob.path)?;
let mmap = match blob.len {
0 => None,
_ => Some(unsafe { Mmap::map(&f)? }),
};
blob_mmaps.push(mmap);
}
merge_sorted_spill::<POSITIONAL_RECORD_LANES, _>(
&sorted_files,
Some(&blob_mmaps),
&term_id_in_lex_order,
&id_to_term,
col_name_bytes,
col_doc_lengths,
avgdl,
n_docs,
&mut key_buf,
&mut postings_writer,
&mut postings_crc_acc,
&mut postings_len,
&mut fst_streaming,
&mut positions_sink,
&mut finish_profile,
&mut term_scratch,
)?
}
};
n_terms_total_usize += n_emitted;
if finish_profile.enabled {
let merge_total = merge_profile_start.elapsed();
let encode_total = finish_profile.encode_total - encode_total_before;
let non_encode = merge_total.saturating_sub(encode_total);
eprintln!(
"[fts-profile] col='{}' merge_total={:.3}s non_encode_merge={:.3}s encode_total={:.3}s calls={} df1={} pfor={} block_build={:.3}s meta_write={:.3}s skip_write={:.3}s block_write={:.3}s fst_insert={:.3}s",
col_name,
merge_total.as_secs_f64(),
non_encode.as_secs_f64(),
encode_total.as_secs_f64(),
finish_profile.encode_calls - encode_calls_before,
finish_profile.encode_df1 - encode_df1_before,
finish_profile.encode_pfor - encode_pfor_before,
(finish_profile.encode_block_build - encode_block_build_before)
.as_secs_f64(),
(finish_profile.encode_meta_write - encode_meta_write_before)
.as_secs_f64(),
(finish_profile.encode_skip_write - encode_skip_write_before)
.as_secs_f64(),
(finish_profile.encode_block_write - encode_block_write_before)
.as_secs_f64(),
(finish_profile.fst_insert - fst_insert_before).as_secs_f64(),
);
}
let cleanup_start = finish_profile.enabled.then(Instant::now);
for p in &sorted_files {
let _ = fs::remove_file(p);
}
drop(partitions);
drop(id_to_term);
drop(term_arena);
drop(lex_rank);
if let Some(t) = cleanup_start {
finish_profile.scratch_cleanup += t.elapsed();
}
}
}
doc_lengths_by_orig_col[orig_col_idx] = Some(col_doc_lengths_owned);
}
assemble_and_write_blob(
BlobAssemblyInputs {
postings_writer,
postings_path,
postings_crc_acc,
postings_len,
positions_sink,
fst_sink: FstSinkFinish::Streaming {
builder: fst_streaming,
path: fst_streaming_path,
},
n_columns,
n_docs,
n_terms_total_usize,
avgdl_per_col,
doc_lengths_by_orig_col,
scratch_dir,
finish_profile,
},
&mut w,
)
}
}
struct PositionsSink {
writer: BufWriter<File>,
path: PathBuf,
crc_acc: u32,
len: u64,
}
impl PositionsSink {
fn create(scratch_path: &Path) -> Result<Self, BuildError> {
let path = scratch_path.join("infino_fts_positions.bin");
Ok(Self {
writer: BufWriter::new(File::create(&path)?),
path,
crc_acc: 0,
len: 0,
})
}
fn write(&mut self, bytes: &[u8]) -> Result<(), BuildError> {
write_counted(&mut self.writer, &mut self.crc_acc, &mut self.len, bytes)
}
}
struct BlobAssemblyInputs {
postings_writer: BufWriter<File>,
postings_path: PathBuf,
postings_crc_acc: u32,
postings_len: u64,
positions_sink: PositionsSink,
fst_sink: FstSinkFinish,
n_columns: u32,
n_docs: u32,
n_terms_total_usize: usize,
avgdl_per_col: Vec<f32>,
doc_lengths_by_orig_col: Vec<Option<Vec<u32>>>,
scratch_dir: tempfile::TempDir,
finish_profile: FinishProfile,
}
enum FstSinkFinish {
InRam(DictBuilder),
Streaming {
builder: StreamingDictBuilder<BufWriter<File>>,
path: PathBuf,
},
}
fn assemble_and_write_blob<W: Write>(
inputs: BlobAssemblyInputs,
w: &mut W,
) -> Result<(), BuildError> {
let BlobAssemblyInputs {
mut postings_writer,
postings_path,
postings_crc_acc,
mut postings_len,
positions_sink,
fst_sink,
n_columns,
n_docs,
n_terms_total_usize,
avgdl_per_col,
mut doc_lengths_by_orig_col,
scratch_dir,
mut finish_profile,
} = inputs;
debug_assert!(
n_terms_total_usize <= u32::MAX as usize,
"term count overflows u32"
);
let n_terms_total = n_terms_total_usize as u32;
let postings_close_start = finish_profile.enabled.then(Instant::now);
let postings_crc = postings_crc_acc;
let postings_crc_le = postings_crc.to_le_bytes();
postings_writer.write_all(&postings_crc_le)?;
postings_writer.flush()?;
drop(postings_writer);
postings_len += postings_crc_le.len() as u64;
if let Some(t) = postings_close_start {
finish_profile.postings_close += t.elapsed();
}
let positions_region = {
let mut sink = positions_sink;
let crc_le = sink.crc_acc.to_le_bytes();
sink.writer.write_all(&crc_le)?;
sink.writer.flush()?;
drop(sink.writer);
(sink.path, sink.len + crc_le.len() as u64)
};
enum FstSource {
InRam(Vec<u8>),
Streamed { path: PathBuf, len: u64, crc: u32 },
}
let fst_close_start = finish_profile.enabled.then(Instant::now);
let fst_source = match fst_sink {
FstSinkFinish::InRam(db) => {
let mut bytes = db.finish();
let crc = crc32c(&bytes);
bytes.extend_from_slice(&crc.to_le_bytes());
FstSource::InRam(bytes)
}
FstSinkFinish::Streaming {
builder,
path: fst_streaming_path,
} => {
let mut bw = builder.finish().map_err(map_fst_err)?;
bw.flush()?;
let write_file = bw
.into_inner()
.map_err(|e| BuildError::Io(e.into_error()))?;
drop(write_file);
let mut read_file = File::open(&fst_streaming_path)?;
let fst_body_len = read_file.metadata()?.len();
read_file.seek(SeekFrom::Start(0))?;
let mut reader = BufReader::with_capacity(PARTITION_BUF_SIZE, read_file);
let mut crc: u32 = 0;
let mut buf = vec![0u8; PARTITION_BUF_SIZE];
loop {
let n = match reader.read(&mut buf) {
Ok(0) => break,
Ok(n) => n,
Err(e) if e.kind() == ErrorKind::Interrupted => continue,
Err(e) => return Err(BuildError::Io(e)),
};
crc = crc32c_append(crc, &buf[..n]);
}
drop(reader);
FstSource::Streamed {
path: fst_streaming_path,
len: fst_body_len + 4,
crc,
}
}
};
if let Some(t) = fst_close_start {
finish_profile.fst_close += t.elapsed();
}
let dl_emit_start = finish_profile.enabled.then(Instant::now);
let fst_total_len: u64 = match &fst_source {
FstSource::InRam(bytes) => bytes.len() as u64,
FstSource::Streamed { len, .. } => *len,
};
let header_size: u64 = format::fts::HEADER_SIZE_V2 as u64;
let fst_offset: u64 = header_size;
let postings_offset: u64 = fst_offset + fst_total_len;
let positions_offset: u64 = postings_offset + postings_len;
let doc_lengths_table_offset: u64 = positions_offset + positions_region.1;
let mut doc_lengths_array_offset: u64 =
doc_lengths_table_offset + (n_columns as u64) * (DOC_LENGTHS_ENTRY_SIZE as u64) + 4 ;
let mut dir_buf: Vec<u8> = Vec::with_capacity(n_columns as usize * DOC_LENGTHS_ENTRY_SIZE);
let mut arrays_buf: Vec<u8> = Vec::new();
for i in 0..n_columns as usize {
let avgdl_x1000 = (avgdl_per_col[i] * format::fts::AVGDL_FIXED_POINT_SCALE)
.max(0.0)
.min(u32::MAX as f32) as u32;
dir_buf.extend_from_slice(&(i as u32).to_le_bytes());
dir_buf.extend_from_slice(&doc_lengths_array_offset.to_le_bytes());
dir_buf.extend_from_slice(&avgdl_x1000.to_le_bytes());
let col_dls = doc_lengths_by_orig_col[i]
.take()
.expect("doc_lengths recorded for every registered column");
let array_start = arrays_buf.len();
#[cfg(target_endian = "little")]
arrays_buf.extend_from_slice(bytemuck::cast_slice::<u32, u8>(&col_dls));
#[cfg(not(target_endian = "little"))]
for &dl in &col_dls {
arrays_buf.extend_from_slice(&dl.to_le_bytes());
}
let array_bytes = &arrays_buf[array_start..];
let array_crc = crc32c(array_bytes);
arrays_buf.extend_from_slice(&array_crc.to_le_bytes());
doc_lengths_array_offset += (col_dls.len() as u64) * 4 + 4;
}
let dir_crc = crc32c(&dir_buf);
dir_buf.extend_from_slice(&dir_crc.to_le_bytes());
if let Some(t) = dl_emit_start {
finish_profile.doc_lengths_emit += t.elapsed();
}
let blob_copy_start = finish_profile.enabled.then(Instant::now);
let mut header = Vec::with_capacity(header_size as usize);
header.extend_from_slice(format::fts::MAGIC); header.extend_from_slice(&format::fts::VERSION_V2.to_le_bytes()); header.extend_from_slice(&n_columns.to_le_bytes()); header.extend_from_slice(&n_docs.to_le_bytes()); header.extend_from_slice(&n_terms_total.to_le_bytes()); header.extend_from_slice(&fst_offset.to_le_bytes()); header.extend_from_slice(&postings_offset.to_le_bytes()); header.extend_from_slice(&doc_lengths_table_offset.to_le_bytes()); header.extend_from_slice(&positions_offset.to_le_bytes()); debug_assert_eq!(header.len(), header_size as usize, "header size mismatch");
w.write_all(&header)?;
match fst_source {
FstSource::InRam(bytes) => w.write_all(&bytes)?,
FstSource::Streamed { path, crc, .. } => {
let mut reader = BufReader::with_capacity(PARTITION_BUF_SIZE, File::open(&path)?);
io::copy(&mut reader, w)?;
w.write_all(&crc.to_le_bytes())?;
}
}
let mut postings_reader =
BufReader::with_capacity(PARTITION_BUF_SIZE, File::open(&postings_path)?);
io::copy(&mut postings_reader, w)?;
drop(postings_reader);
{
let mut positions_reader =
BufReader::with_capacity(PARTITION_BUF_SIZE, File::open(&positions_region.0)?);
io::copy(&mut positions_reader, w)?;
}
drop(scratch_dir);
w.write_all(&dir_buf)?;
w.write_all(&arrays_buf)?;
if let Some(t) = blob_copy_start {
finish_profile.blob_copy += t.elapsed();
}
if finish_profile.enabled {
eprintln!(
"[fts-finish] partition_flush={:.3}s lex_rank={:.3}s partition_sort={:.3}s mmap_open={:.3}s scratch_cleanup={:.3}s postings_close={:.3}s fst_close={:.3}s doc_lengths_emit={:.3}s blob_copy={:.3}s",
finish_profile.partition_flush.as_secs_f64(),
finish_profile.lex_rank_build.as_secs_f64(),
finish_profile.partition_sort.as_secs_f64(),
finish_profile.mmap_open.as_secs_f64(),
finish_profile.scratch_cleanup.as_secs_f64(),
finish_profile.postings_close.as_secs_f64(),
finish_profile.fst_close.as_secs_f64(),
finish_profile.doc_lengths_emit.as_secs_f64(),
finish_profile.blob_copy.as_secs_f64(),
);
}
Ok(())
}
#[inline]
fn map_fst_err(e: fst::Error) -> BuildError {
BuildError::Io(Error::new(ErrorKind::InvalidData, e))
}
#[derive(Default)]
struct TermScratch {
doc_ids: Vec<u32>,
tfs: Vec<u32>,
block_ub_per_block: Vec<f32>,
encoded_blocks: Vec<EncodedBlock>,
term_buf: Vec<u8>,
pos_block_offsets: Vec<u32>,
}
#[allow(clippy::too_many_arguments)]
fn merge_sorted_spill<const N: usize, W: Write>(
sorted_files: &[PathBuf],
blob_mmaps: Option<&[Option<Mmap>]>,
term_id_in_lex_order: &[u32],
id_to_term: &[&'static str],
col_name_bytes: &[u8],
col_doc_lengths: &[u32],
avgdl: f32,
n_docs: u32,
key_buf: &mut Vec<u8>,
postings_writer: &mut W,
postings_crc_acc: &mut u32,
postings_len: &mut u64,
fst_streaming: &mut StreamingDictBuilder<BufWriter<File>>,
positions_sink: &mut PositionsSink,
finish_profile: &mut FinishProfile,
term_scratch: &mut TermScratch,
) -> Result<usize, BuildError> {
let mmap_start = finish_profile.enabled.then(Instant::now);
let mut mmaps: Vec<Mmap> = Vec::with_capacity(sorted_files.len());
for p in sorted_files {
let f = File::open(p)?;
let mmap = unsafe { Mmap::map(&f)? };
mmaps.push(mmap);
}
if let Some(t) = mmap_start {
finish_profile.mmap_open += t.elapsed();
}
let sorted_slices: Vec<&[[u32; N]]> = mmaps
.iter()
.map(|m| {
if m.is_empty() {
&[][..]
} else {
bytemuck::cast_slice::<u8, [u32; N]>(&m[..])
}
})
.collect();
let blob_slices: Vec<&[u8]> = match blob_mmaps {
Some(ms) => ms
.iter()
.map(|m| m.as_ref().map_or(&[][..], |m| &m[..]))
.collect(),
None => Vec::new(),
};
let mask = (sorted_slices.len() - 1) as u32;
let mut cursors: Vec<usize> = vec![0usize; sorted_slices.len()];
let mut group: Vec<(u32, u32)> = Vec::new();
let mut group_pos: Vec<u64> = Vec::new();
let mut term_run: Vec<u8> = Vec::new();
let mut n_emitted = 0usize;
for &term_id in term_id_in_lex_order {
let p = (term_id & mask) as usize;
let slice = sorted_slices[p];
let mut pos = cursors[p];
group.clear();
group_pos.clear();
while pos < slice.len() {
let t = &slice[pos];
if triple_term_id(t) != term_id {
break;
}
group.push((triple_doc_id(t), triple_tf(t)));
if blob_mmaps.is_some() {
debug_assert_eq!(N, POSITIONAL_RECORD_LANES);
let rec: &[u32; POSITIONAL_RECORD_LANES] =
t[..].try_into().expect("positional record width");
group_pos.push(record_pos_off(rec));
}
pos += 1;
}
cursors[p] = pos;
if group.is_empty() {
continue;
}
let term_positions = match blob_mmaps.is_some() {
true => {
term_run.clear();
let blob = blob_slices[p];
for (i, &(_, tf)) in group.iter().enumerate() {
let start = group_pos[i] as usize;
let mut at = start;
skip_run(blob, &mut at, tf).expect("builder-encoded blob runs are well-formed");
term_run.extend_from_slice(&blob[start..at]);
}
Some((&mut *positions_sink, term_run.as_slice()))
}
false => None,
};
let term_bytes: &str = id_to_term[term_id as usize];
encode_and_emit_term(
term_bytes,
&group,
col_name_bytes,
col_doc_lengths,
avgdl,
n_docs,
key_buf,
postings_writer,
postings_crc_acc,
postings_len,
None,
Some(fst_streaming),
term_positions,
finish_profile,
term_scratch,
)?;
n_emitted += 1;
}
debug_assert!(
cursors
.iter()
.zip(sorted_slices.iter())
.all(|(c, s)| *c == s.len()),
"lex-order partition traversal did not drain all records; \
partition assignment or sort invariant violated"
);
Ok(n_emitted)
}
#[allow(clippy::too_many_arguments)]
fn encode_and_emit_term<W: Write>(
term: &str,
pairs: &[(u32, u32)],
col_name_bytes: &[u8],
col_doc_lengths: &[u32],
avgdl: f32,
n_docs: u32,
key_buf: &mut Vec<u8>,
postings_writer: &mut W,
postings_crc_acc: &mut u32,
postings_len: &mut u64,
fst_entries_inram: Option<&mut DictBuilder>,
mut fst_streaming: Option<&mut StreamingDictBuilder<BufWriter<File>>>,
mut term_positions: Option<(&mut PositionsSink, &[u8])>,
profile: &mut FinishProfile,
scratch: &mut TermScratch,
) -> Result<(), BuildError> {
let encode_start = profile.enabled.then(Instant::now);
profile.encode_calls += 1;
key_buf.clear();
key_buf.extend_from_slice(col_name_bytes);
key_buf.push(FST_SEPARATOR);
key_buf.extend_from_slice(term.as_bytes());
debug_assert!(
pairs.windows(2).all(|w| w[0].0 < w[1].0),
"posting list not sorted by doc_id"
);
let df = pairs.len() as u64;
let inline_value: Option<u64> = if df == 1 {
let (doc_id, tf) = pairs[0];
match &term_positions {
None => Some(FstValue::pack_inline(doc_id, tf)),
Some((_, runs)) if tf == 1 => {
let mut at = 0;
let pos = read_varint(runs, &mut at).expect("builder-encoded run is well-formed");
(pos <= INLINE_TF_MAX).then(|| FstValue::pack_inline(doc_id, pos))
}
Some(_) => None,
}
} else {
None
};
let fst_value: u64 = if let Some(v) = inline_value {
profile.encode_df1 += 1;
v
} else {
profile.encode_pfor += 1;
let idf_t = bm25::idf(n_docs as u64, df);
let encoded_blocks = &mut scratch.encoded_blocks;
let block_ub_per_block = &mut scratch.block_ub_per_block;
encoded_blocks.clear();
block_ub_per_block.clear();
let block_build_start = profile.enabled.then(Instant::now);
let mut block_doc_ids = mem::take(&mut scratch.doc_ids);
let mut block_tfs = mem::take(&mut scratch.tfs);
if block_doc_ids.capacity() < BLOCK_LEN {
block_doc_ids.reserve(BLOCK_LEN - block_doc_ids.capacity());
}
if block_tfs.capacity() < BLOCK_LEN {
block_tfs.reserve(BLOCK_LEN - block_tfs.capacity());
}
for chunk in pairs.chunks(BLOCK_LEN) {
block_doc_ids.clear();
block_tfs.clear();
block_doc_ids.extend(chunk.iter().map(|&(d, _)| d));
block_tfs.extend(chunk.iter().map(|&(_, t)| t));
let block_ub = block_doc_ids
.iter()
.zip(block_tfs.iter())
.map(|(&d, &t)| bm25::score(idf_t, t, col_doc_lengths[d as usize], avgdl))
.fold(0.0f32, f32::max);
block_ub_per_block.push(block_ub);
let block = Block {
doc_ids: mem::take(&mut block_doc_ids),
tfs: mem::take(&mut block_tfs),
};
encoded_blocks.push(encode_block(&block));
block_doc_ids = block.doc_ids;
block_tfs = block.tfs;
}
scratch.doc_ids = block_doc_ids;
scratch.tfs = block_tfs;
if let Some(start) = block_build_start {
profile.encode_block_build += start.elapsed();
}
let num_blocks = encoded_blocks.len() as u32;
let metadata_offset = *postings_len;
let skip_table_size = encoded_blocks.len() * SKIP_ENTRY_SIZE;
let blocks_total_size: usize = encoded_blocks.iter().map(|b| b.bytes.len()).sum();
let term_meta_size = match term_positions {
Some(_) => TERM_META_POSITIONAL_SIZE,
None => TERM_META_SIZE,
};
let postings_length = (term_meta_size + skip_table_size + blocks_total_size) as u64;
let pos_block_offsets = &mut scratch.pos_block_offsets;
pos_block_offsets.clear();
if let Some((_, runs)) = &term_positions {
debug_assert!(
runs.len() <= u32::MAX as usize,
"single-term positions > 4 GiB"
);
let mut at: usize = 0;
for (i, &(_, tf)) in pairs.iter().enumerate() {
if i % BLOCK_LEN == 0 {
pos_block_offsets.push(at as u32);
}
skip_run(runs, &mut at, tf).expect("builder-encoded runs are well-formed");
}
debug_assert_eq!(at, runs.len(), "runs must cover exactly the pairs");
}
debug_assert!(df <= u32::MAX as u64, "df overflows u32");
debug_assert!(
postings_length <= u32::MAX as u64,
"single-term posting > 4 GiB"
);
let term_buf = &mut scratch.term_buf;
term_buf.clear();
if term_buf.capacity() < postings_length as usize {
term_buf.reserve(postings_length as usize - term_buf.capacity());
}
let meta_write_start = profile.enabled.then(Instant::now);
term_buf.extend_from_slice(&(df as u32).to_le_bytes());
term_buf.extend_from_slice(&metadata_offset.to_le_bytes());
term_buf.extend_from_slice(&(postings_length as u32).to_le_bytes());
term_buf.extend_from_slice(&num_blocks.to_le_bytes());
if let Some((sink, runs)) = &term_positions {
term_buf.extend_from_slice(&sink.len.to_le_bytes());
term_buf.extend_from_slice(&(runs.len() as u32).to_le_bytes());
}
debug_assert_eq!(term_buf.len(), term_meta_size);
if let Some(start) = meta_write_start {
profile.encode_meta_write += start.elapsed();
}
let mut block_offset: u32 = (term_meta_size + skip_table_size) as u32;
let skip_write_start = profile.enabled.then(Instant::now);
for (i, blk) in encoded_blocks.iter().enumerate() {
let max_bm25 = block_ub_per_block[i];
let max_bm25_x1000 = (max_bm25 * format::fts::BLOCK_MAX_BM25_FIXED_POINT_SCALE)
.ceil()
.max(0.0)
.min(u32::MAX as f32) as u32;
term_buf.extend_from_slice(&blk.last_doc_id.to_le_bytes());
term_buf.extend_from_slice(&block_offset.to_le_bytes());
term_buf.extend_from_slice(&max_bm25_x1000.to_le_bytes());
let pos_block_off = pos_block_offsets.get(i).copied().unwrap_or(0);
term_buf.extend_from_slice(&pos_block_off.to_le_bytes());
block_offset += blk.bytes.len() as u32;
}
if let Some(start) = skip_write_start {
profile.encode_skip_write += start.elapsed();
}
let block_write_start = profile.enabled.then(Instant::now);
for blk in encoded_blocks.iter() {
term_buf.extend_from_slice(&blk.bytes);
}
debug_assert_eq!(term_buf.len(), postings_length as usize);
write_counted(postings_writer, postings_crc_acc, postings_len, term_buf)?;
if let Some(start) = block_write_start {
profile.encode_block_write += start.elapsed();
}
if let Some((sink, runs)) = term_positions.as_mut() {
sink.write(runs)?;
}
FstValue::pack_pfor(metadata_offset, postings_length as u32)
};
let fst_insert_start = profile.enabled.then(Instant::now);
if let Some(db) = fst_entries_inram {
db.insert(key_buf, fst_value);
} else if let Some(sb) = fst_streaming.as_mut() {
sb.insert_sorted(key_buf, fst_value).map_err(map_fst_err)?;
}
if let Some(start) = fst_insert_start {
profile.fst_insert += start.elapsed();
}
if let Some(start) = encode_start {
profile.encode_total += start.elapsed();
}
Ok(())
}
fn write_counted<W: Write>(
w: &mut W,
crc_acc: &mut u32,
len: &mut u64,
bytes: &[u8],
) -> Result<(), BuildError> {
w.write_all(bytes)?;
*crc_acc = crc32c_append(*crc_acc, bytes);
*len += bytes.len() as u64;
Ok(())
}
fn sort_partition_to_file<const N: usize>(
in_path: &Path,
out_path: &Path,
max_partition_bytes: u64,
scratch_dir: &Path,
partition_label: &str,
lex_rank: &[u32],
) -> Result<(), BuildError> {
let mut iter = open_partition_sorted::<N>(
in_path,
max_partition_bytes,
scratch_dir,
partition_label,
lex_rank,
)?;
let mut w = BufWriter::with_capacity(PARTITION_BUF_SIZE, File::create(out_path)?);
let mut batch: Vec<[u32; N]> = Vec::with_capacity(SORT_OUTPUT_BATCH_TRIPLES);
while let Some(triple) = iter.next_with(lex_rank) {
let t = triple?;
batch.push(t);
if batch.len() == SORT_OUTPUT_BATCH_TRIPLES {
#[cfg(target_endian = "little")]
{
w.write_all(bytemuck::cast_slice::<[u32; N], u8>(&batch))?;
}
#[cfg(not(target_endian = "little"))]
for t in &batch {
write_record(&mut w, t)?;
}
batch.clear();
}
}
if !batch.is_empty() {
#[cfg(target_endian = "little")]
{
w.write_all(bytemuck::cast_slice::<[u32; N], u8>(&batch))?;
}
#[cfg(not(target_endian = "little"))]
for t in &batch {
write_record(&mut w, t)?;
}
}
w.flush()?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_helpers::default_tokenizer as tokenizer;
#[test]
fn register_column_returns_sequential_ids() {
let mut b = FtsBuilder::new(tokenizer());
assert_eq!(
b.register_column("title".into(), false)
.expect("register column"),
0
);
assert_eq!(
b.register_column("body".into(), false)
.expect("register column"),
1
);
assert_eq!(
b.register_column("tag".into(), false)
.expect("register column"),
2
);
}
#[test]
fn register_column_rejects_separator_byte() {
let mut b = FtsBuilder::new(tokenizer());
let bad = String::from("ti\x1Ftle");
let err = b.register_column(bad, false).expect_err("expected error");
assert!(matches!(err, BuildError::ReservedSeparatorInColumnName(_)));
}
#[test]
fn register_column_rejects_reserved_prefix() {
let mut b = FtsBuilder::new(tokenizer());
let err = b
.register_column("inf.title".into(), false)
.expect_err("expected error");
assert!(matches!(err, BuildError::ReservedPrefixInColumnName(_)));
}
#[test]
fn register_column_rejects_duplicates() {
let mut b = FtsBuilder::new(tokenizer());
b.register_column("title".into(), false)
.expect("register column");
let err = b
.register_column("title".into(), false)
.expect_err("expected error");
assert!(matches!(err, BuildError::DuplicateColumnName(_)));
}
#[test]
fn add_doc_unknown_column_id_errors() {
let mut b = FtsBuilder::new(tokenizer());
b.register_column("title".into(), false)
.expect("register column");
let err = b.add_doc(99, 0, "text").expect_err("expected error");
assert!(matches!(err, BuildError::FtsColumnTypeInvalid { .. }));
}
#[test]
fn add_doc_reuses_term_frequency_table_capacity() {
let mut b = FtsBuilder::new(tokenizer());
b.register_column("title".into(), false)
.expect("register column");
b.add_doc(0, 0, "alpha beta gamma delta epsilon zeta eta theta")
.expect("add first doc");
assert!(b.doc_tf.is_empty(), "per-document table must be drained");
let capacity = b.doc_tf.capacity();
assert!(capacity > 0, "first document must allocate table buckets");
b.add_doc(0, 1, "alpha beta alpha").expect("add second doc");
assert!(b.doc_tf.is_empty(), "per-document table must be drained");
assert_eq!(
b.doc_tf.capacity(),
capacity,
"subsequent documents must reuse the existing buckets"
);
}
#[tokio::test]
async fn add_doc_accumulates_tf_within_doc() {
use bytes::Bytes;
use crate::superfile::fts::reader::{BoolMode, FtsReader};
let mut b = FtsBuilder::new(tokenizer());
b.register_column("title".into(), false)
.expect("register column");
b.add_doc(0, 0, "rust rust rust async").expect("add doc");
let blob = Bytes::from(b.finish().expect("finish"));
let r =
FtsReader::open(blob, r#"[{"name":"title","tokenizer":"ascii_lower"}]"#).expect("open");
let rust_hits = r
.search("title", &["rust"], 10, BoolMode::Or)
.await
.expect("rust search");
let async_hits = r
.search("title", &["async"], 10, BoolMode::Or)
.await
.expect("async search");
assert_eq!(rust_hits.len(), 1);
assert_eq!(rust_hits[0].0, 0);
assert_eq!(async_hits.len(), 1);
assert_eq!(async_hits[0].0, 0);
}
#[tokio::test]
async fn cross_column_same_term_stays_isolated_through_round_trip() {
use bytes::Bytes;
use crate::superfile::fts::reader::{BoolMode, FtsReader};
let mut b = FtsBuilder::new(tokenizer());
let title_id = b
.register_column("title".into(), false)
.expect("register title");
let body_id = b
.register_column("body".into(), false)
.expect("register body");
b.add_doc(title_id, 0, "rust tokio")
.expect("add title doc 0");
b.add_doc(body_id, 0, "rust async").expect("add body doc 0");
b.add_doc(body_id, 1, "rust").expect("add body doc 1");
b.add_doc(title_id, 1, "rust").expect("add title doc 1");
let blob = Bytes::from(b.finish().expect("finish"));
let json = r#"[{"name":"title","tokenizer":"ascii_lower"},{"name":"body","tokenizer":"ascii_lower"}]"#;
let r = FtsReader::open(blob, json).expect("open");
let hits_t = r
.search("title", &["rust"], 10, BoolMode::Or)
.await
.expect("title search");
let ids_t: Vec<u32> = hits_t.iter().map(|(d, _)| *d).collect();
assert_eq!(ids_t.len(), 2, "title 'rust' hit count");
assert!(ids_t.contains(&0));
assert!(ids_t.contains(&1));
let hits_b = r
.search("body", &["rust"], 10, BoolMode::Or)
.await
.expect("body search");
let ids_b: Vec<u32> = hits_b.iter().map(|(d, _)| *d).collect();
assert_eq!(ids_b.len(), 2, "body 'rust' hit count");
assert!(ids_b.contains(&0));
assert!(ids_b.contains(&1));
let hits_async_in_title = r
.search("title", &["async"], 10, BoolMode::Or)
.await
.expect("title async search");
assert!(
hits_async_in_title.is_empty(),
"title must not return 'async' (lives only in body)"
);
let hits_tokio_in_body = r
.search("body", &["tokio"], 10, BoolMode::Or)
.await
.expect("body tokio search");
assert!(
hits_tokio_in_body.is_empty(),
"body must not return 'tokio' (lives only in title)"
);
}
#[test]
fn add_doc_tracks_doc_lengths_clamped() {
let mut b = FtsBuilder::new(tokenizer());
b.register_column("body".into(), false)
.expect("register column");
b.add_doc(0, 0, "alpha beta gamma").expect("add doc");
b.add_doc(0, 1, "").expect("add doc"); b.add_doc(0, 2, "delta").expect("add doc");
let col = &b.columns[0];
assert_eq!(col.doc_lengths, vec![3, 0, 1]);
assert_eq!(col.total_tokens, 4);
}
#[test]
fn add_doc_updates_n_docs_per_call() {
let mut b = FtsBuilder::new(tokenizer());
b.register_column("body".into(), false)
.expect("register column");
b.add_doc(0, 0, "a").expect("add doc");
b.add_doc(0, 1, "b").expect("add doc");
b.add_doc(0, 2, "c").expect("add doc");
assert_eq!(b.n_docs, 3);
}
#[test]
fn finish_emits_valid_header() {
let mut b = FtsBuilder::new(tokenizer());
b.register_column("title".into(), false)
.expect("register column");
b.add_doc(0, 0, "hello world").expect("add doc");
let blob = b.finish().expect("finish");
assert_eq!(&blob[0..8], format::fts::MAGIC);
let version = u32::from_le_bytes([blob[8], blob[9], blob[10], blob[11]]);
assert_eq!(version, format::fts::VERSION_V2);
let n_cols = u32::from_le_bytes([blob[12], blob[13], blob[14], blob[15]]);
assert_eq!(n_cols, 1);
let n_docs = u32::from_le_bytes([blob[16], blob[17], blob[18], blob[19]]);
assert_eq!(n_docs, 1);
let n_terms = u32::from_le_bytes([blob[20], blob[21], blob[22], blob[23]]);
assert_eq!(n_terms, 2);
let mut buf = [0u8; 8];
buf.copy_from_slice(&blob[24..32]);
let fst_off = u64::from_le_bytes(buf);
assert_eq!(fst_off, format::fts::HEADER_SIZE_V2 as u64);
}
#[test]
fn finish_to_matches_finish_byte_for_byte() {
fn build() -> FtsBuilder {
let mut b = FtsBuilder::new(tokenizer());
b.register_column("title".into(), false)
.expect("register title");
for (i, text) in [
"rust async rust",
"tokio runtime",
"rust search engine",
"async search",
]
.iter()
.enumerate()
{
b.add_doc(0, i as u32, text).expect("add doc");
}
b
}
let via_finish = build().finish().expect("finish");
let mut via_finish_to = Vec::new();
build()
.finish_to(&mut via_finish_to)
.expect("finish_to Vec");
assert_eq!(via_finish_to, via_finish);
}
#[tokio::test]
async fn finish_to_temp_file_round_trips_through_reader() {
use std::io::BufWriter;
use bytes::Bytes;
use crate::superfile::fts::reader::{BoolMode, FtsReader};
let mut b = FtsBuilder::new(tokenizer());
b.register_column("title".into(), false)
.expect("register title");
for i in 0..256u32 {
b.add_doc(0, i, &format!("common term{i:03}"))
.expect("add doc");
}
let tmp = tempfile::tempdir().expect("tempdir");
let path = tmp.path().join("fts.blob");
{
let file = File::create(&path).expect("create blob");
let writer = BufWriter::new(file);
b.finish_to(writer).expect("finish_to file");
}
let blob = fs::read(&path).expect("read blob");
let r = FtsReader::open(
Bytes::from(blob),
r#"[{"name":"title","tokenizer":"ascii_lower"}]"#,
)
.expect("open FTS reader");
let hits = r
.search("title", &["common"], 10, BoolMode::Or)
.await
.expect("search");
assert_eq!(hits.len(), 10);
}
#[test]
fn finish_with_no_docs_still_produces_valid_blob() {
let mut b = FtsBuilder::new(tokenizer());
b.register_column("title".into(), false)
.expect("register column");
let blob = b.finish().expect("finish");
assert_eq!(&blob[0..8], format::fts::MAGIC);
assert_eq!(
u32::from_le_bytes([blob[16], blob[17], blob[18], blob[19]]),
0
);
assert_eq!(
u32::from_le_bytes([blob[20], blob[21], blob[22], blob[23]]),
0
);
}
#[test]
fn small_build_stays_in_ram_no_spill_files_created() {
let parent = tempfile::tempdir().expect("parent");
let mut b = FtsBuilder::with_scratch(tokenizer(), parent.path().to_path_buf())
.expect("with_scratch");
b.register_column("body".into(), false)
.expect("register col");
for i in 0..100u32 {
b.add_doc(0, i, &format!("alpha beta gamma{i}"))
.expect("add doc");
}
for cp in &b.postings {
assert!(
!cp.is_spilled(),
"small build must not have spilled to disk"
);
}
let mut spill_files_found = 0usize;
for entry in walkdir_files(parent.path()) {
let name = entry.file_name().to_string_lossy().to_string();
if name.starts_with("fts_col") && name.ends_with(".bin") {
spill_files_found += 1;
}
}
assert_eq!(
spill_files_found, 0,
"small build must not pre-create posting spill files"
);
let _blob = b.finish().expect("finish");
}
#[test]
fn build_above_threshold_spills_and_matches_in_ram_byte_for_byte() {
fn build_corpus(b: &mut FtsBuilder) {
b.register_column("body".into(), false)
.expect("register col");
for i in 0..1000u32 {
b.add_doc(
0,
i,
&format!("common shared term{i:04} payload{i:04} extra word{i:04}"),
)
.expect("add doc");
}
}
let mut baseline = FtsBuilder::new(tokenizer());
build_corpus(&mut baseline);
for cp in &baseline.postings {
assert!(!cp.is_spilled(), "baseline must stay in RAM");
}
let baseline_blob = baseline.finish().expect("finish baseline");
let parent = tempfile::tempdir().expect("parent");
let mut spilled = FtsBuilder::with_scratch(tokenizer(), parent.path().to_path_buf())
.expect("with_scratch");
spilled.set_spill_threshold_bytes(16 * 1024);
build_corpus(&mut spilled);
let any_spilled = spilled.postings.iter().any(|c| c.is_spilled());
assert!(any_spilled, "low threshold must force spill");
let mut spill_files_found = 0usize;
for entry in walkdir_files(parent.path()) {
let name = entry.file_name().to_string_lossy().to_string();
if name.starts_with("fts_col") && name.ends_with(".bin") {
spill_files_found += 1;
}
}
assert!(
spill_files_found > 0,
"spilled build must materialise at least one fts_col*.bin partition file on disk"
);
let spilled_blob = spilled.finish().expect("finish spilled");
assert_eq!(
spilled_blob, baseline_blob,
"streaming-FST + spill path must produce byte-identical blob"
);
}
fn walkdir_files(root: &Path) -> Vec<fs::DirEntry> {
let mut out = Vec::new();
let mut stack: Vec<PathBuf> = vec![root.to_path_buf()];
while let Some(dir) = stack.pop() {
let rd = match fs::read_dir(&dir) {
Ok(r) => r,
Err(_) => continue,
};
for entry in rd.flatten() {
let ft = match entry.file_type() {
Ok(t) => t,
Err(_) => continue,
};
if ft.is_dir() {
stack.push(entry.path());
} else if ft.is_file() {
out.push(entry);
}
}
}
out
}
#[test]
fn external_merge_path_matches_in_memory_path_byte_for_byte() {
fn build_corpus(builder: &mut FtsBuilder) {
builder
.register_column("body".into(), false)
.expect("register col");
for i in 0..600u32 {
builder
.add_doc(0, i, &format!("common term{i:04} payload{i:04}"))
.expect("add doc");
}
}
let mut baseline = FtsBuilder::new(tokenizer());
baseline.set_spill_threshold_bytes(1);
build_corpus(&mut baseline);
let baseline_blob = baseline.finish().expect("finish baseline");
finish_debug::reset();
let mut tight = FtsBuilder::new(tokenizer());
tight.set_spill_threshold_bytes(1);
tight.set_max_partition_bytes(1024);
build_corpus(&mut tight);
let tight_blob = tight.finish().expect("finish tight");
assert_eq!(
tight_blob, baseline_blob,
"external-merge path must produce identical blob bytes"
);
let chunks = finish_debug::observed();
assert!(
!chunks.is_empty(),
"external-merge path must have written at least one sorted-chunk file; \
observed chunks were empty (test no longer exercises the over-budget branch)"
);
}
#[test]
fn scratch_dir_under_with_scratch_is_removed_after_finish() {
let parent = tempfile::tempdir().expect("parent tempdir");
let dir_count_before = fs::read_dir(parent.path()).expect("read parent").count();
let mut b = FtsBuilder::with_scratch(tokenizer(), parent.path().to_path_buf())
.expect("with_scratch");
b.register_column("body".into(), false)
.expect("register col");
b.add_doc(0, 0, "alpha beta gamma").expect("add doc");
let _blob = b.finish().expect("finish");
let dir_count_after = fs::read_dir(parent.path()).expect("read parent").count();
assert_eq!(
dir_count_after, dir_count_before,
"FtsBuilder scratch tempdir leaked under override path"
);
}
#[tokio::test]
async fn configurable_spill_partitions_round_trips_through_reader() {
use bytes::Bytes;
use crate::superfile::fts::reader::{BoolMode, FtsReader};
let mut b = FtsBuilder::new(tokenizer());
b.set_spill_partitions(256).expect("set partitions");
b.register_column("body".into(), false)
.expect("register col");
for i in 0..50u32 {
b.add_doc(0, i, &format!("alpha beta gamma{i:02}"))
.expect("add doc");
}
let blob = b.finish().expect("finish");
let r = FtsReader::open(
Bytes::from(blob),
r#"[{"name":"body","tokenizer":"ascii_lower"}]"#,
)
.expect("open reader");
let hits = r
.search("body", &["alpha"], 100, BoolMode::Or)
.await
.expect("search alpha");
assert_eq!(hits.len(), 50, "alpha is in every doc");
}
#[test]
fn finish_offsets_are_consistent() {
let mut b = FtsBuilder::new(tokenizer());
b.register_column("body".into(), false)
.expect("register column");
for i in 0..10 {
b.add_doc(0, i, &format!("term{i} common"))
.expect("add doc");
}
let blob = b.finish().expect("finish");
let mut buf = [0u8; 8];
buf.copy_from_slice(&blob[24..32]);
let fst_off = u64::from_le_bytes(buf) as usize;
buf.copy_from_slice(&blob[32..40]);
let postings_off = u64::from_le_bytes(buf) as usize;
buf.copy_from_slice(&blob[40..48]);
let dir_off = u64::from_le_bytes(buf) as usize;
assert_eq!(fst_off, format::fts::HEADER_SIZE_V2);
assert!(postings_off > fst_off, "postings after FST");
assert!(dir_off > postings_off, "directory after postings");
assert!(dir_off <= blob.len(), "directory offset within blob");
buf.copy_from_slice(&blob[48..56]);
let positions_off = u64::from_le_bytes(buf) as usize;
assert!(positions_off > postings_off, "positions after postings");
assert!(dir_off > positions_off, "directory after positions");
}
#[test]
fn set_spill_partitions_rejects_after_register_column() {
let mut b = FtsBuilder::new(tokenizer());
b.register_column("body".into(), false)
.expect("register col");
let err = b.set_spill_partitions(16).expect_err("expected error");
match err {
BuildError::Io(e) => {
assert!(e.to_string().contains("before any register_column"));
}
other => panic!("expected Io error, got {:?}", other),
}
}
#[test]
fn set_spill_partitions_rejects_zero() {
let mut b = FtsBuilder::new(tokenizer());
let err = b.set_spill_partitions(0).expect_err("expected error");
match err {
BuildError::Io(e) => assert!(e.to_string().contains("must be ≥ 1")),
other => panic!("expected Io error, got {:?}", other),
}
}
#[test]
fn set_spill_partitions_rejects_non_power_of_two() {
const NON_PO2: usize = 7;
let mut b = FtsBuilder::new(tokenizer());
let err = b.set_spill_partitions(NON_PO2).expect_err("expected error");
match err {
BuildError::Io(e) => assert!(e.to_string().contains("power of two")),
other => panic!("expected Io error, got {:?}", other),
}
}
#[test]
fn read_partition_triples_empty_file_is_empty() {
let dir = tempfile::tempdir().expect("tempdir");
let path = dir.path().join("empty.part");
fs::write(&path, []).expect("write empty file");
let triples = read_partition_records::<PLAIN_RECORD_LANES>(&path).expect("read empty");
assert!(triples.is_empty());
}
#[test]
fn read_partition_triples_round_trips_le_bytes() {
let dir = tempfile::tempdir().expect("tempdir");
let path = dir.path().join("good.part");
let mut bytes = Vec::new();
for field in [3u32, 4, 5, 6, 7, 8] {
bytes.extend_from_slice(&field.to_le_bytes());
}
fs::write(&path, &bytes).expect("write triples");
let triples = read_partition_records::<PLAIN_RECORD_LANES>(&path).expect("read triples");
assert_eq!(triples, vec![[3u32, 4, 5], [6u32, 7, 8]]);
}
#[test]
fn read_partition_triples_rejects_non_multiple_length() {
let dir = tempfile::tempdir().expect("tempdir");
let path = dir.path().join("ragged.part");
fs::write(
&path,
vec![0u8; mem::size_of::<[u32; PLAIN_RECORD_LANES]>() + 1],
)
.expect("write ragged file");
let err = read_partition_records::<PLAIN_RECORD_LANES>(&path).expect_err("expected error");
match err {
BuildError::Io(e) => {
assert_eq!(e.kind(), ErrorKind::InvalidData);
assert!(e.to_string().contains("not a multiple"));
}
other => panic!("expected Io error, got {:?}", other),
}
}
#[test]
fn merge_entry_orders_by_sort_key_reversed() {
let small = MergeEntry {
sort_key: 10,
rec: [1u32, 10, 1],
reader_idx: 0,
};
let large = MergeEntry {
sort_key: 20,
rec: [2u32, 20, 1],
reader_idx: 1,
};
assert_eq!(small.cmp(&large), Ordering::Greater);
assert_eq!(small.partial_cmp(&large), Some(Ordering::Greater));
assert!(small != large);
let small_dup = MergeEntry {
sort_key: 10,
rec: [9u32, 10, 9],
reader_idx: 0,
};
assert!(small == small_dup);
let mut heap: BinaryHeap<MergeEntry<PLAIN_RECORD_LANES>> = BinaryHeap::new();
heap.push(large);
heap.push(small);
assert_eq!(heap.pop().expect("non-empty heap").sort_key, 10);
}
fn synthesize_v1_blob(v2: &[u8]) -> Vec<u8> {
const V2_HEADER: usize = 56;
const V1_HEADER: usize = 48;
const HEADER_SHRINK: u64 = (V2_HEADER - V1_HEADER) as u64;
const EMPTY_REGION_CRC: u64 = 4;
let read_u64 = |at: usize| u64::from_le_bytes(v2[at..at + 8].try_into().expect("8 bytes"));
let read_u32 = |at: usize| u32::from_le_bytes(v2[at..at + 4].try_into().expect("4 bytes"));
assert_eq!(read_u32(8), format::fts::VERSION_V2);
let fst_off = read_u64(24);
let postings_off = read_u64(32);
let doc_lengths_off = read_u64(40);
let positions_off = read_u64(48);
assert_eq!(
doc_lengths_off - positions_off,
EMPTY_REGION_CRC,
"synthesis requires a positionless blob (empty region)"
);
let n_columns = read_u32(12) as usize;
let mut out = Vec::with_capacity(v2.len() - V2_HEADER + V1_HEADER);
out.extend_from_slice(&v2[0..8]); out.extend_from_slice(&format::fts::VERSION_V1_LEGACY.to_le_bytes());
out.extend_from_slice(&v2[12..24]); out.extend_from_slice(&(fst_off - HEADER_SHRINK).to_le_bytes());
out.extend_from_slice(&(postings_off - HEADER_SHRINK).to_le_bytes());
let v1_doc_lengths_off = doc_lengths_off - HEADER_SHRINK - EMPTY_REGION_CRC;
out.extend_from_slice(&v1_doc_lengths_off.to_le_bytes());
out.extend_from_slice(&v2[V2_HEADER..positions_off as usize]);
let dir_start = doc_lengths_off as usize;
let dir_size = n_columns * 16;
let mut dir = Vec::with_capacity(dir_size);
for c in 0..n_columns {
let e = dir_start + c * 16;
dir.extend_from_slice(&v2[e..e + 4]);
let arr_off = read_u64(e + 4) - HEADER_SHRINK - EMPTY_REGION_CRC;
dir.extend_from_slice(&arr_off.to_le_bytes());
dir.extend_from_slice(&v2[e + 12..e + 16]);
}
let dir_crc = crc32c(&dir);
out.extend_from_slice(&dir);
out.extend_from_slice(&dir_crc.to_le_bytes());
out.extend_from_slice(&v2[dir_start + dir_size + 4..]);
out
}
fn build_title_blob_spilled(
docs: &[String],
positional: bool,
max_partition_bytes: Option<u64>,
) -> bytes::Bytes {
let mut b = FtsBuilder::new(tokenizer());
b.set_spill_threshold_bytes(1);
if let Some(m) = max_partition_bytes {
b.set_max_partition_bytes(m);
}
b.register_column("title".into(), positional)
.expect("register column");
for (i, text) in docs.iter().enumerate() {
b.add_doc(0, i as u32, text).expect("add doc");
}
bytes::Bytes::from(b.finish().expect("finish"))
}
async fn assert_title_blobs_agree(
a: bytes::Bytes,
a_json: &str,
b: bytes::Bytes,
b_json: &str,
k: usize,
) {
use crate::superfile::fts::reader::{BoolMode, FtsReader};
let ra = FtsReader::open(a, a_json).expect("open a");
let rb = FtsReader::open(b, b_json).expect("open b");
let queries: &[(&[&str], BoolMode)] = &[
(&["common"], BoolMode::Or),
(&["uniqueonce"], BoolMode::Or),
(&["dupdup"], BoolMode::Or),
(&["medium", "uniqueonce"], BoolMode::Or),
(&["common", "medium"], BoolMode::And),
];
for (terms, mode) in queries {
let ha = ra.search("title", terms, k, *mode).await.expect("search a");
let hb = rb.search("title", terms, k, *mode).await.expect("search b");
assert_eq!(ha, hb, "results diverged for {terms:?} ({mode:?})");
assert!(!ha.is_empty(), "corpus sanity: {terms:?} matches");
}
for term in ["common", "medium", "uniqueonce", "dupdup"] {
assert_eq!(
ra.term_df("title", term).await.expect("df a"),
rb.term_df("title", term).await.expect("df b"),
"df diverged for {term}"
);
}
}
#[tokio::test]
async fn spilled_positional_build_searches_identically() {
let docs = positional_corpus();
let k = docs.len();
let spilled_pos = build_title_blob_spilled(&docs, true, None);
assert_eq!(
u32::from_le_bytes(spilled_pos[8..12].try_into().expect("version bytes")),
format::fts::VERSION_V2
);
let inram_pos = build_title_blob(&docs, true);
let inram_plain = build_title_blob(&docs, false);
assert_title_blobs_agree(
spilled_pos.clone(),
title_json(true),
inram_pos,
title_json(true),
k,
)
.await;
assert_title_blobs_agree(
spilled_pos,
title_json(true),
inram_plain,
title_json(false),
k,
)
.await;
}
#[tokio::test]
async fn spilled_positional_external_merge_searches_identically() {
let docs = positional_corpus();
let k = docs.len();
let spilled = build_title_blob_spilled(&docs, true, Some(64));
let inram = build_title_blob(&docs, true);
assert_title_blobs_agree(spilled, title_json(true), inram, title_json(true), k).await;
}
#[tokio::test]
async fn new_code_reads_synthesized_v1_blob() {
use crate::superfile::fts::reader::{BoolMode, FtsReader};
let docs = positional_corpus();
let v2_blob = build_title_blob(&docs, false);
let v1_blob = bytes::Bytes::from(synthesize_v1_blob(&v2_blob));
assert_eq!(
u32::from_le_bytes(v1_blob[8..12].try_into().expect("version bytes")),
format::fts::VERSION_V1_LEGACY
);
let v1 = FtsReader::open(v1_blob, title_json(false)).expect("v1 opens");
let v2 = FtsReader::open(v2_blob, title_json(false)).expect("v2 opens");
let queries: &[&[&str]] = &[&["common"], &["uniqueonce"], &["common", "medium"]];
for terms in queries {
let a = v1
.search("title", terms, docs.len(), BoolMode::Or)
.await
.expect("v1 search");
let b = v2
.search("title", terms, docs.len(), BoolMode::Or)
.await
.expect("v2 search");
assert_eq!(a, b, "v1/v2 results diverged for {terms:?}");
assert!(!a.is_empty());
}
assert_eq!(
v1.term_df("title", "common").await.expect("v1 df"),
v2.term_df("title", "common").await.expect("v2 df"),
);
}
fn title_json(positional: bool) -> &'static str {
match positional {
true => r#"[{"name":"title","tokenizer":"ascii_lower","positions":true}]"#,
false => r#"[{"name":"title","tokenizer":"ascii_lower"}]"#,
}
}
fn build_title_blob(docs: &[String], positional: bool) -> bytes::Bytes {
let mut b = FtsBuilder::new(tokenizer());
b.register_column("title".into(), positional)
.expect("register column");
for (i, text) in docs.iter().enumerate() {
b.add_doc(0, i as u32, text).expect("add doc");
}
bytes::Bytes::from(b.finish().expect("finish"))
}
fn positional_corpus() -> Vec<String> {
let n_docs = 3 * BLOCK_LEN + 7;
let mut docs = Vec::with_capacity(n_docs);
for i in 0..n_docs {
let mut t = String::from("common filler");
if i % 5 == 0 {
t.push_str(" medium medium");
}
if i == 42 {
t.push_str(" uniqueonce");
}
if i == 43 {
t.push_str(" dupdup dupdup dupdup");
}
docs.push(t);
}
docs
}
#[test]
fn every_build_writes_v2_and_positionless_region_is_empty() {
let docs = positional_corpus();
let plain = build_title_blob(&docs, false);
let positional = build_title_blob(&docs, true);
let version_of = |blob: &bytes::Bytes| {
u32::from_le_bytes(blob[8..12].try_into().expect("4 header bytes"))
};
assert_eq!(version_of(&plain), format::fts::VERSION_V2);
assert_eq!(version_of(&positional), format::fts::VERSION_V2);
let read_u64_plain =
|at: usize| u64::from_le_bytes(plain[at..at + 8].try_into().expect("8 header bytes"));
let region_len = read_u64_plain(40) - read_u64_plain(48);
assert_eq!(region_len, 4, "positionless region = 4-byte CRC only");
let read_u64 = |blob: &bytes::Bytes, at: usize| {
u64::from_le_bytes(blob[at..at + 8].try_into().expect("8 header bytes"))
};
let postings_off = read_u64(&positional, 32);
let doc_lengths_off = read_u64(&positional, 40);
let positions_off = read_u64(&positional, 48);
assert!(postings_off < positions_off, "positions follow postings");
assert!(
positions_off < doc_lengths_off,
"doc-lengths directory follows positions"
);
assert!(doc_lengths_off - positions_off > 4, "region has a body");
}
#[tokio::test]
async fn positional_build_searches_identically_to_positionless() {
use crate::superfile::fts::reader::{BoolMode, FtsReader};
let docs = positional_corpus();
let v1 = FtsReader::open(build_title_blob(&docs, false), title_json(false)).expect("v1");
let v2 = FtsReader::open(build_title_blob(&docs, true), title_json(true)).expect("v2");
let queries: &[(&[&str], BoolMode)] = &[
(&["common"], BoolMode::Or),
(&["uniqueonce"], BoolMode::Or),
(&["dupdup"], BoolMode::Or),
(&["medium", "uniqueonce"], BoolMode::Or),
(&["common", "medium"], BoolMode::And),
];
let k = docs.len();
for (terms, mode) in queries {
let a = v1
.search("title", terms, k, *mode)
.await
.expect("v1 search");
let b = v2
.search("title", terms, k, *mode)
.await
.expect("v2 search");
assert_eq!(a, b, "results diverged for {terms:?} ({mode:?})");
assert!(!a.is_empty(), "corpus sanity: {terms:?} matches");
}
for term in ["common", "medium", "uniqueonce", "dupdup"] {
let a = v1.term_df("title", term).await.expect("v1 df");
let b = v2.term_df("title", term).await.expect("v2 df");
assert_eq!(a, b, "df diverged for {term}");
let ca = v1
.token_match_count("title", &[term], BoolMode::Or)
.await
.expect("v1 count");
let cb = v2
.token_match_count("title", &[term], BoolMode::Or)
.await
.expect("v2 count");
assert_eq!(ca, cb, "count diverged for {term}");
}
}
#[tokio::test]
async fn mixed_columns_only_positional_column_pays() {
use crate::superfile::fts::reader::{BoolMode, FtsReader};
let mut b = FtsBuilder::new(tokenizer());
b.register_column("body".into(), false).expect("register");
b.register_column("title".into(), true).expect("register");
for i in 0..(BLOCK_LEN as u32 + 9) {
b.add_doc(0, i, "shared bodyterm").expect("body doc");
b.add_doc(1, i, "shared titleterm titleterm")
.expect("title doc");
}
let blob = bytes::Bytes::from(b.finish().expect("finish"));
assert_eq!(
u32::from_le_bytes(blob[8..12].try_into().expect("version bytes")),
format::fts::VERSION_V2
);
let json = r#"[{"name":"body","tokenizer":"ascii_lower"},{"name":"title","tokenizer":"ascii_lower","positions":true}]"#;
let r = FtsReader::open(blob, json).expect("open");
let body_hits = r
.search("body", &["bodyterm"], 10, BoolMode::Or)
.await
.expect("body search");
let title_hits = r
.search("title", &["titleterm"], 10, BoolMode::Or)
.await
.expect("title search");
assert_eq!(body_hits.len(), 10);
assert_eq!(title_hits.len(), 10);
let shared_body = r
.token_match_count("body", &["shared"], BoolMode::Or)
.await
.expect("shared body");
assert_eq!(shared_body, BLOCK_LEN as u64 + 9);
}
}