use alloc::vec::Vec;
use crate::util::varint;
#[cfg(test)]
pub(crate) static INDEX_ROUTE_HITS: core::sync::atomic::AtomicUsize =
core::sync::atomic::AtomicUsize::new(0);
const MAIN_PREFIX: u8 = b'0';
pub(crate) const AVERAGES_ROWID: i64 = 1;
pub(crate) const STRUCTURE_ROWID: i64 = 10;
pub(crate) fn segment_leaf_rowid(segid: i64, pgno: i64) -> i64 {
(segid << 37) | pgno
}
pub(crate) fn dlidx_rowid(segid: i64, height: i64, pgno: i64) -> i64 {
(segid << 37) | (1 << 36) | (height << 31) | pgno
}
const MIN_DLIDX_SIZE: usize = 4;
fn dlidx_first_rowid(page: &[u8]) -> i64 {
let mut pos = 1usize; if read_varint(page, &mut pos).is_none() {
return 0;
}
read_varint(page, &mut pos).unwrap_or(0) as i64
}
fn put_varint(out: &mut Vec<u8>, v: u64) {
let mut buf = [0u8; varint::MAX_LEN];
let n = varint::encode(v, &mut buf);
out.extend_from_slice(&buf[..n]);
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct Posting {
pub rowid: i64,
pub cols: Vec<Vec<u32>>,
pub del: bool,
}
fn collist(positions: &[u32]) -> Vec<u8> {
let mut out = Vec::new();
let mut prev = 0u32;
for (i, &pos) in positions.iter().enumerate() {
put_varint(
&mut out,
((if i == 0 { pos } else { pos - prev }) as u64) + 2,
);
prev = pos;
}
out
}
fn poslist(p: &Posting) -> Vec<u8> {
let mut content = Vec::new();
for (c, positions) in p.cols.iter().enumerate() {
if positions.is_empty() {
continue;
}
if c != 0 {
content.push(0x01);
put_varint(&mut content, c as u64);
}
content.extend_from_slice(&collist(positions));
}
let mut out = Vec::new();
put_varint(&mut out, (content.len() as u64) * 2 + u64::from(p.del));
out.extend_from_slice(&content);
out
}
fn poslist_prefix(buf: &[u8], n_max: isize) -> usize {
let mut ret = varint::decode(buf).map(|(_, n)| n).unwrap_or(1);
if (ret as isize) < n_max {
while let Some((_, i)) = varint::decode(&buf[ret..]) {
if (ret + i) as isize > n_max {
break;
}
ret += i;
}
}
ret
}
fn term_key(term: &[u8]) -> Vec<u8> {
let mut key = Vec::with_capacity(term.len() + 1);
key.push(MAIN_PREFIX);
key.extend_from_slice(term);
key
}
fn pgidx(offsets: &[usize]) -> Vec<u8> {
let mut out = Vec::new();
let mut prev = 0usize;
for (i, &off) in offsets.iter().enumerate() {
put_varint(&mut out, (if i == 0 { off } else { off - prev }) as u64);
prev = off;
}
out
}
fn separator(prev_last: &[u8], first: &[u8]) -> Vec<u8> {
let mut i = 0;
while i < prev_last.len() && i < first.len() && prev_last[i] == first[i] {
i += 1;
}
first[..=i.min(first.len() - 1)].to_vec()
}
pub(crate) struct IdxRow {
pub segid: i64,
pub term: Vec<u8>,
pub pgno: i64,
}
type SegParts = (Vec<Vec<u8>>, Vec<IdxRow>, Vec<(i64, Vec<u8>)>);
pub(crate) type TokenizedDocs = (Vec<(Vec<u8>, Vec<Posting>)>, Vec<u64>, Vec<(i64, Vec<u64>)>);
struct SegWriter {
pgsz: usize,
segid: i64,
leaves: Vec<Vec<u8>>,
idx: Vec<IdxRow>,
body: Vec<u8>,
term_offsets: Vec<usize>,
first_rowid_off: usize,
prev_term_key: Option<Vec<u8>>,
prev_rowid: i64,
leaf_first_term: Option<Vec<u8>>,
leaf_last_term: Option<Vec<u8>>,
prev_leaf_last_term: Option<Vec<u8>>,
pgno: i64,
dlidx_data: Vec<(i64, Vec<u8>)>,
span_pages: Vec<(i64, i64)>,
first_rowid_in_page: bool,
first_rowid_in_doclist: bool,
merge_mode: bool,
}
impl SegWriter {
fn new(pgsz: usize, segid: i64) -> Self {
SegWriter {
pgsz,
segid,
leaves: Vec::new(),
idx: Vec::new(),
body: Vec::new(),
term_offsets: Vec::new(),
first_rowid_off: 0,
prev_term_key: None,
prev_rowid: 0,
leaf_first_term: None,
leaf_last_term: None,
prev_leaf_last_term: None,
pgno: 1,
dlidx_data: Vec::new(),
span_pages: Vec::new(),
first_rowid_in_page: true,
first_rowid_in_doclist: true,
merge_mode: false,
}
}
fn finish_leaf(&self) -> Vec<u8> {
let footer_off = 4 + self.body.len();
let mut leaf = Vec::new();
leaf.extend_from_slice(&(self.first_rowid_off as u16).to_be_bytes());
leaf.extend_from_slice(&(footer_off as u16).to_be_bytes());
leaf.extend_from_slice(&self.body);
leaf.extend_from_slice(&pgidx(&self.term_offsets));
leaf
}
fn flush(&mut self) {
self.leaves.push(self.finish_leaf());
if let Some(ft) = self.leaf_first_term.take() {
let term = match &self.prev_leaf_last_term {
Some(p) => separator(p, &ft),
None => Vec::new(),
};
self.idx.push(IdxRow {
segid: self.segid,
term,
pgno: self.pgno << 1,
});
}
if let Some(lt) = self.leaf_last_term.take() {
self.prev_leaf_last_term = Some(lt);
}
self.body.clear();
self.term_offsets.clear();
self.first_rowid_off = 0;
self.prev_term_key = None;
self.prev_rowid = 0;
self.pgno += 1;
self.first_rowid_in_page = true;
}
fn term_record(&self, key: &[u8]) -> Vec<u8> {
let mut rec = Vec::new();
match &self.prev_term_key {
None => {
put_varint(&mut rec, key.len() as u64);
rec.extend_from_slice(key);
}
Some(prev) => {
let n_common = key
.iter()
.zip(prev.iter())
.take_while(|(a, b)| a == b)
.count();
put_varint(&mut rec, n_common as u64);
put_varint(&mut rec, (key.len() - n_common) as u64);
rec.extend_from_slice(&key[n_common..]);
}
}
rec
}
fn pgidx_len(&self) -> usize {
pgidx(&self.term_offsets).len()
}
fn append_rowid(&mut self, rowid: i64, term_start_leaf: i64) {
if self.merge_mode && 4 + self.body.len() + self.pgidx_len() >= self.pgsz {
self.flush();
}
if self.first_rowid_in_page {
self.first_rowid_off = 4 + self.body.len();
if self.pgno != term_start_leaf {
self.span_pages.push((self.pgno, rowid));
}
}
if self.first_rowid_in_doclist || self.first_rowid_in_page {
put_varint(&mut self.body, rowid as u64);
} else {
put_varint(&mut self.body, (rowid - self.prev_rowid) as u64);
}
self.prev_rowid = rowid;
self.first_rowid_in_doclist = false;
self.first_rowid_in_page = false;
}
fn append_poslist_data(&mut self, data: &[u8]) {
let n_copy = data.len();
if self.merge_mode {
let mut a = data;
while 4 + self.body.len() + self.pgidx_len() + a.len() >= self.pgsz {
let n_req = self.pgsz as isize - (4 + self.body.len() + self.pgidx_len()) as isize;
let mut n_c = 0usize;
while (n_c as isize) < n_req {
match varint::decode(&a[n_c..]) {
Some((_, len)) => n_c += len,
None => break,
}
}
if n_c > a.len() {
n_c = a.len();
}
self.body.extend_from_slice(&a[..n_c]);
a = &a[n_c..];
self.flush();
}
if !a.is_empty() {
self.body.extend_from_slice(a);
}
return;
}
if 4 + self.body.len() + self.pgidx_len() + n_copy <= self.pgsz {
self.body.extend_from_slice(data);
return;
}
let mut i_pos = 0usize;
loop {
let n_space = self.pgsz as isize - (4 + self.body.len() + self.pgidx_len()) as isize;
let n = if (n_copy - i_pos) as isize <= n_space {
n_copy - i_pos
} else {
poslist_prefix(&data[i_pos..], n_space)
};
self.body.extend_from_slice(&data[i_pos..i_pos + n]);
i_pos += n;
if 4 + self.body.len() + self.pgidx_len() >= self.pgsz {
self.flush();
}
if i_pos >= n_copy {
break;
}
}
}
fn add_term(&mut self, term: &[u8], postings: &[Posting]) {
let key = term_key(term);
self.add_key(&key, postings);
}
fn add_key(&mut self, key: &[u8], postings: &[Posting]) {
if 4 + self.body.len() + self.pgidx_len() + key.len() + 2 >= self.pgsz
&& !self.body.is_empty()
{
self.flush();
}
let rec = self.term_record(key);
self.term_offsets.push(4 + self.body.len());
if self.leaf_first_term.is_none() {
self.leaf_first_term = Some(key.to_vec());
}
self.leaf_last_term = Some(key.to_vec());
self.body.extend_from_slice(&rec);
self.prev_term_key = Some(key.to_vec());
self.first_rowid_in_page = false;
self.first_rowid_in_doclist = true;
let term_start_leaf = self.pgno;
self.span_pages.clear();
self.prev_rowid = 0;
for p in postings {
self.append_rowid(p.rowid, term_start_leaf);
let pl = poslist(p);
if self.merge_mode {
let size_len = varint::decode(&pl).map(|(_, n)| n).unwrap_or(pl.len());
self.body.extend_from_slice(&pl[..size_len]);
if size_len < pl.len() {
self.append_poslist_data(&pl[size_len..]);
}
} else {
self.append_poslist_data(&pl);
}
}
self.finish_term_dlidx(term_start_leaf);
}
fn finish_term_dlidx(&mut self, term_start_leaf: i64) {
if self.span_pages.is_empty() {
return;
}
let last_leaf = self.span_pages.last().map(|&(pg, _)| pg).unwrap_or(0);
let n_empty = (last_leaf - term_start_leaf) as usize;
if n_empty < MIN_DLIDX_SIZE {
self.span_pages.clear();
return;
}
let span = core::mem::take(&mut self.span_pages);
let pgsz = self.pgsz;
struct DlidxLvl {
buf: Vec<u8>,
pgno: i64,
prev: i64,
prev_valid: bool,
}
let new_lvl = |pgno: i64| DlidxLvl {
buf: Vec::new(),
pgno,
prev: 0,
prev_valid: false,
};
let mut lvls: Vec<DlidxLvl> = alloc::vec![new_lvl(term_start_leaf)];
let mut out: Vec<(i64, i64, Vec<u8>)> = Vec::new();
let mut prev_leaf = term_start_leaf;
for &(leaf_pgno, rowid) in &span {
for _ in 0..(leaf_pgno - prev_leaf - 1).max(0) {
lvls[0].buf.push(0x00);
}
prev_leaf = leaf_pgno;
let mut i = 0usize;
let mut b_done = false;
while !b_done {
if i >= lvls.len() {
lvls.push(new_lvl(0));
}
if lvls[i].buf.len() >= pgsz {
lvls[i].buf[0] = 0x01;
let flushed = core::mem::take(&mut lvls[i].buf);
let flushed_pgno = lvls[i].pgno;
out.push((i as i64, flushed_pgno, flushed.clone()));
if i + 1 >= lvls.len() {
lvls.push(new_lvl(0));
}
if lvls[i + 1].buf.is_empty() {
let first = dlidx_first_rowid(&flushed);
lvls[i + 1].pgno = flushed_pgno;
let parent = &mut lvls[i + 1];
parent.buf.push(0x00);
put_varint(&mut parent.buf, flushed_pgno as u64);
put_varint(&mut parent.buf, first as u64);
parent.prev = first;
parent.prev_valid = true;
}
lvls[i].prev_valid = false;
lvls[i].pgno += 1;
} else {
b_done = true;
}
if lvls[i].prev_valid {
let d = (rowid - lvls[i].prev) as u64;
put_varint(&mut lvls[i].buf, d);
} else {
let ref_pgno = if i == 0 { leaf_pgno } else { lvls[i - 1].pgno };
let flag = u8::from(!b_done);
let lvl = &mut lvls[i];
lvl.buf.push(flag);
put_varint(&mut lvl.buf, ref_pgno as u64);
put_varint(&mut lvl.buf, rowid as u64);
}
lvls[i].prev = rowid;
lvls[i].prev_valid = true;
i += 1;
}
}
let top = lvls.iter().rposition(|l| !l.buf.is_empty()).unwrap_or(0);
for (i, lvl) in lvls.iter_mut().enumerate() {
if lvl.buf.is_empty() {
continue;
}
let mut page = core::mem::take(&mut lvl.buf);
page[0] = if i == top { 0x00 } else { 0x01 };
out.push((i as i64, lvl.pgno, page));
}
for (height, pgno, page) in out {
self.dlidx_data
.push((dlidx_rowid(self.segid, height, pgno), page));
}
let want = term_start_leaf << 1;
for row in self.idx.iter_mut() {
if row.pgno == want {
row.pgno |= 1;
break;
}
}
}
fn finish(mut self) -> SegParts {
self.flush();
(self.leaves, self.idx, self.dlidx_data)
}
}
fn structure(n_leaves: i64, cookie: u32) -> Vec<u8> {
let mut out = cookie.to_be_bytes().to_vec();
if n_leaves == 0 {
out.extend_from_slice(&[0, 0, 0]); return out;
}
for v in [1, 1, n_leaves as u64, 0, 1, 1, 1, n_leaves as u64] {
put_varint(&mut out, v);
}
out
}
pub(crate) fn encode_averages(n_rows: u64, col_totals: &[u64]) -> Vec<u8> {
let mut out = Vec::new();
if n_rows > 0 {
put_varint(&mut out, n_rows);
for &t in col_totals {
put_varint(&mut out, t);
}
}
out
}
pub(crate) fn encode_averages_full(n_rows: u64, col_totals: &[u64]) -> Vec<u8> {
let mut out = Vec::new();
put_varint(&mut out, n_rows);
for &t in col_totals {
put_varint(&mut out, t);
}
out
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct StructSeg {
pub segid: i64,
pub pgno_first: i64,
pub pgno_last: i64,
}
impl StructSeg {
pub(crate) fn size(&self) -> i64 {
1 + self.pgno_last - self.pgno_first
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct StructLevel {
pub n_merge: i64,
pub segs: Vec<StructSeg>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct SegStructure {
pub cookie: u32,
pub write_counter: u64,
pub levels: Vec<StructLevel>,
}
impl SegStructure {
pub(crate) fn parse(buf: &[u8]) -> Option<SegStructure> {
if buf.len() < 4 {
return None;
}
let cookie = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
let mut pos = 4usize;
let n_level = read_varint(buf, &mut pos)?;
let n_segment = read_varint(buf, &mut pos)?;
let write_counter = read_varint(buf, &mut pos)?;
let mut levels = Vec::with_capacity(n_level as usize);
let mut total = 0u64;
for _ in 0..n_level {
let n_merge = read_varint(buf, &mut pos)? as i64;
let n_seg = read_varint(buf, &mut pos)?;
let mut segs = Vec::with_capacity(n_seg as usize);
for _ in 0..n_seg {
let segid = read_varint(buf, &mut pos)? as i64;
let pgno_first = read_varint(buf, &mut pos)? as i64;
let pgno_last = read_varint(buf, &mut pos)? as i64;
segs.push(StructSeg {
segid,
pgno_first,
pgno_last,
});
}
total += n_seg;
levels.push(StructLevel { n_merge, segs });
}
if total != n_segment {
return None;
}
Some(SegStructure {
cookie,
write_counter,
levels,
})
}
pub(crate) fn encode(&self) -> Vec<u8> {
let mut out = self.cookie.to_be_bytes().to_vec();
let n_segment: u64 = self.levels.iter().map(|l| l.segs.len() as u64).sum();
put_varint(&mut out, self.levels.len() as u64);
put_varint(&mut out, n_segment);
put_varint(&mut out, self.write_counter);
for lvl in &self.levels {
put_varint(&mut out, lvl.n_merge as u64);
put_varint(&mut out, lvl.segs.len() as u64);
for s in &lvl.segs {
put_varint(&mut out, s.segid as u64);
put_varint(&mut out, s.pgno_first as u64);
put_varint(&mut out, s.pgno_last as u64);
}
}
out
}
pub(crate) fn allocate_segid(&self) -> i64 {
let mut used = alloc::collections::BTreeSet::new();
for lvl in &self.levels {
for s in &lvl.segs {
if s.segid > 0 {
used.insert(s.segid);
}
}
}
let mut id = 1;
while used.contains(&id) {
id += 1;
}
id
}
pub(crate) fn append_level0(&mut self, segid: i64, n_leaves: i64) {
if self.levels.is_empty() {
self.levels.push(StructLevel {
n_merge: 0,
segs: Vec::new(),
});
}
self.levels[0].segs.push(StructSeg {
segid,
pgno_first: 1,
pgno_last: n_leaves,
});
self.write_counter += n_leaves as u64;
self.promote(0);
}
pub(crate) fn promote_after_merge(&mut self, i_lvl: usize) {
if i_lvl < self.levels.len() {
self.promote(i_lvl);
}
}
fn promote(&mut self, i_lvl: usize) {
if self.levels[i_lvl].segs.is_empty() {
return;
}
let sz_seg = self.levels[i_lvl].segs.last().unwrap().size();
let mut i_tst: isize = i_lvl as isize - 1;
while i_tst >= 0 && self.levels[i_tst as usize].segs.is_empty() {
i_tst -= 1;
}
let mut i_promote = i_lvl;
let mut sz_promote = sz_seg;
if i_tst >= 0 {
let tst = &self.levels[i_tst as usize];
let sz_max = tst.segs.iter().map(|s| s.size()).max().unwrap_or(0);
if sz_max >= sz_seg {
i_promote = i_tst as usize;
sz_promote = sz_max;
}
}
self.promote_to(i_promote, sz_promote);
}
fn promote_to(&mut self, i_promote: usize, sz_promote: i64) {
if self.levels[i_promote].n_merge != 0 {
return;
}
let mut il = i_promote + 1;
while il < self.levels.len() {
if self.levels[il].n_merge != 0 {
return;
}
let mut is = self.levels[il].segs.len();
while is > 0 {
is -= 1;
let sz = self.levels[il].segs[is].size();
if sz > sz_promote {
return;
}
let seg = self.levels[il].segs.remove(is);
self.levels[i_promote].segs.insert(0, seg);
}
il += 1;
}
}
}
pub(crate) struct Segment {
pub data: Vec<(i64, Vec<u8>)>,
pub idx: Vec<IdxRow>,
pub docsize: Vec<(i64, Vec<u8>)>,
}
pub(crate) fn build_segment(
terms: &[(Vec<u8>, Vec<Posting>)],
n_docs: u64,
col_totals: &[u64],
doc_sizes: &[(i64, Vec<u64>)],
pgsz: usize,
cookie: u32,
) -> Segment {
build_segment_prefixed(terms, n_docs, col_totals, doc_sizes, pgsz, cookie, &[])
}
fn prefix_bytelen(t: &[u8], n_char: usize) -> Option<usize> {
let mut n = 0usize;
for i in 0..n_char {
if n >= t.len() {
return None; }
if t[n] >= 0xc0 {
n += 1;
if n >= t.len() {
return None;
}
while t[n] & 0xc0 == 0x80 {
n += 1;
if n >= t.len() {
if i + 1 == n_char {
break;
}
return None;
}
}
} else {
n += 1;
}
}
Some(n)
}
fn merge_prefix_postings(groups: &[&[Posting]]) -> Vec<Posting> {
use alloc::collections::BTreeMap;
let mut by_rowid: BTreeMap<i64, (bool, Vec<Vec<u32>>)> = BTreeMap::new();
for postings in groups {
for p in *postings {
let entry = by_rowid
.entry(p.rowid)
.or_insert_with(|| (false, alloc::vec![Vec::new(); p.cols.len()]));
entry.0 |= p.del;
if entry.1.len() < p.cols.len() {
entry.1.resize(p.cols.len(), Vec::new());
}
for (c, positions) in p.cols.iter().enumerate() {
entry.1[c].extend_from_slice(positions);
}
}
}
by_rowid
.into_iter()
.map(|(rowid, (del, mut cols))| {
for col in &mut cols {
col.sort_unstable();
col.dedup();
}
Posting { rowid, cols, del }
})
.collect()
}
pub(crate) fn build_segment_prefixed(
terms: &[(Vec<u8>, Vec<Posting>)],
n_docs: u64,
col_totals: &[u64],
doc_sizes: &[(i64, Vec<u64>)],
pgsz: usize,
cookie: u32,
prefixes: &[usize],
) -> Segment {
let segid = 1;
let (leaves, idx, dlidx) = build_segment_leaves(terms, pgsz, segid, prefixes);
let mut data: Vec<(i64, Vec<u8>)> = Vec::new();
let mut avg = Vec::new();
if n_docs > 0 {
put_varint(&mut avg, n_docs);
for &t in col_totals {
put_varint(&mut avg, t);
}
}
data.push((AVERAGES_ROWID, avg));
data.push((STRUCTURE_ROWID, structure(leaves.len() as i64, cookie)));
for (i, leaf) in leaves.iter().enumerate() {
data.push((segment_leaf_rowid(segid, i as i64 + 1), leaf.clone()));
}
data.extend(dlidx);
let docsize = build_docsize(doc_sizes);
Segment { data, idx, docsize }
}
pub(crate) fn build_docsize(doc_sizes: &[(i64, Vec<u64>)]) -> Vec<(i64, Vec<u8>)> {
doc_sizes
.iter()
.map(|(rowid, sizes)| {
let mut sz = Vec::new();
for &s in sizes {
put_varint(&mut sz, s);
}
(*rowid, sz)
})
.collect()
}
fn build_segment_leaves(
terms: &[(Vec<u8>, Vec<Posting>)],
pgsz: usize,
segid: i64,
prefixes: &[usize],
) -> SegParts {
use alloc::collections::BTreeMap;
let mut w = SegWriter::new(pgsz.max(16), segid);
for (term, postings) in terms {
w.add_term(term, postings);
}
for (i, &n_char) in prefixes.iter().enumerate() {
let idx_byte = MAIN_PREFIX + (i as u8) + 1;
let mut groups: BTreeMap<&[u8], Vec<&[Posting]>> = BTreeMap::new();
for (term, postings) in terms {
if let Some(nb) = prefix_bytelen(term, n_char) {
groups.entry(&term[..nb]).or_default().push(postings);
}
}
for (pfx, parts) in groups {
let merged = merge_prefix_postings(&parts);
let mut key = Vec::with_capacity(pfx.len() + 1);
key.push(idx_byte);
key.extend_from_slice(pfx);
w.add_key(&key, &merged);
}
}
if terms.is_empty() {
(Vec::new(), Vec::new(), Vec::new())
} else {
w.finish()
}
}
pub(crate) struct SegmentBlock {
pub data: Vec<(i64, Vec<u8>)>,
pub idx: Vec<IdxRow>,
pub docsize: Vec<(i64, Vec<u8>)>,
pub n_leaves: i64,
}
pub(crate) fn build_segment_block(
terms: &[(Vec<u8>, Vec<Posting>)],
doc_sizes: &[(i64, Vec<u64>)],
pgsz: usize,
segid: i64,
prefixes: &[usize],
) -> SegmentBlock {
let (leaves, idx, dlidx) = build_segment_leaves(terms, pgsz, segid, prefixes);
let mut data: Vec<(i64, Vec<u8>)> = Vec::new();
for (i, leaf) in leaves.iter().enumerate() {
data.push((segment_leaf_rowid(segid, i as i64 + 1), leaf.clone()));
}
data.extend(dlidx);
SegmentBlock {
data,
idx,
docsize: build_docsize(doc_sizes),
n_leaves: leaves.len() as i64,
}
}
pub(crate) fn build_merged_segment_block(
terms: &[(Vec<u8>, Vec<Posting>)],
pgsz: usize,
segid: i64,
) -> SegmentBlock {
let (leaves, idx, dlidx) = {
let mut w = SegWriter::new(pgsz.max(16), segid);
w.merge_mode = true;
for (term, postings) in terms {
w.add_term(term, postings);
}
if terms.is_empty() {
(Vec::new(), Vec::new(), Vec::new())
} else {
w.finish()
}
};
let mut data: Vec<(i64, Vec<u8>)> = Vec::new();
for (i, leaf) in leaves.iter().enumerate() {
data.push((segment_leaf_rowid(segid, i as i64 + 1), leaf.clone()));
}
data.extend(dlidx);
SegmentBlock {
data,
idx,
docsize: Vec::new(),
n_leaves: leaves.len() as i64,
}
}
pub(crate) fn build_merged_segment_block_full(
terms: &[(Vec<u8>, Vec<Posting>)],
pgsz: usize,
segid: i64,
) -> SegmentBlock {
let (leaves, idx, dlidx) = {
let mut w = SegWriter::new(pgsz.max(16), segid);
w.merge_mode = true;
for (key, postings) in terms {
w.add_key(key, postings);
}
if terms.is_empty() {
(Vec::new(), Vec::new(), Vec::new())
} else {
w.finish()
}
};
let mut data: Vec<(i64, Vec<u8>)> = Vec::new();
for (i, leaf) in leaves.iter().enumerate() {
data.push((segment_leaf_rowid(segid, i as i64 + 1), leaf.clone()));
}
data.extend(dlidx);
SegmentBlock {
data,
idx,
docsize: Vec::new(),
n_leaves: leaves.len() as i64,
}
}
fn read_varint(buf: &[u8], pos: &mut usize) -> Option<u64> {
let (v, n) = varint::decode(buf.get(*pos..)?)?;
*pos += n;
Some(v)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct DecodedPosting {
pub rowid: i64,
pub cols: Vec<Vec<u32>>,
}
fn decode_poslist(buf: &[u8], pos: &mut usize) -> Option<Vec<Vec<u32>>> {
let size2 = read_varint(buf, pos)?;
if size2 & 1 != 0 {
return None;
}
let content_len = (size2 / 2) as usize;
let end = pos.checked_add(content_len)?;
if end > buf.len() {
return None;
}
let mut cols: Vec<Vec<u32>> = Vec::new();
let mut col = 0usize;
let mut p = *pos;
cols.push(Vec::new());
while p < end {
if buf[p] == 0x01 {
p += 1;
let c = read_varint(buf, &mut p)? as usize;
col = c;
while cols.len() <= col {
cols.push(Vec::new());
}
} else {
let raw = read_varint(buf, &mut p)?;
if raw < 2 {
return None;
}
let delta = (raw - 2) as u32;
let next = if cols[col].is_empty() {
delta
} else {
cols[col].last().copied()?.checked_add(delta)?
};
cols[col].push(next);
}
}
*pos = end;
Some(cols)
}
struct TermRec {
key: Vec<u8>,
rec_start: usize,
doclist_start: usize,
}
struct LeafView {
first_rowid_off: usize,
footer_off: usize,
terms: Vec<TermRec>,
}
fn parse_leaf(leaf: &[u8]) -> Option<LeafView> {
if leaf.len() < 4 {
return None;
}
let first_rowid_off = u16::from_be_bytes([leaf[0], leaf[1]]) as usize;
let footer_off = u16::from_be_bytes([leaf[2], leaf[3]]) as usize;
if footer_off < 4 || footer_off > leaf.len() {
return None;
}
if first_rowid_off != 0 && (first_rowid_off < 4 || first_rowid_off > footer_off) {
return None;
}
let mut term_offs: Vec<usize> = Vec::new();
{
let mut p = footer_off;
let mut prev = 0usize;
let mut first = true;
while p < leaf.len() {
let d = read_varint(leaf, &mut p)? as usize;
let off = if first { d } else { prev.checked_add(d)? };
first = false;
if off >= footer_off || off < 4 {
return None;
}
term_offs.push(off);
prev = off;
}
}
let mut terms = Vec::with_capacity(term_offs.len());
let mut prev_key: Vec<u8> = Vec::new();
for (i, &off) in term_offs.iter().enumerate() {
let mut p = off;
let key = if i == 0 {
let keylen = read_varint(leaf, &mut p)? as usize;
let end = p.checked_add(keylen)?;
if end > footer_off {
return None;
}
let key = leaf.get(p..end)?.to_vec();
p = end;
key
} else {
let n_common = read_varint(leaf, &mut p)? as usize;
let n_new = read_varint(leaf, &mut p)? as usize;
let end = p.checked_add(n_new)?;
if end > footer_off || n_common > prev_key.len() {
return None;
}
let mut key = prev_key.get(..n_common)?.to_vec();
key.extend_from_slice(leaf.get(p..end)?);
p = end;
key
};
if p > footer_off {
return None;
}
terms.push(TermRec {
key: key.clone(),
rec_start: off,
doclist_start: p,
});
prev_key = key;
}
Some(LeafView {
first_rowid_off,
footer_off,
terms,
})
}
struct DoclistRun<'a> {
bytes: &'a [u8],
abs_start: bool,
}
fn decode_spanning_doclist(runs: &[DoclistRun]) -> Option<Vec<DecodedPosting>> {
let mut buf: Vec<u8> = Vec::new();
let mut abs_at: Vec<usize> = Vec::new();
for run in runs {
if run.abs_start && !run.bytes.is_empty() {
abs_at.push(buf.len());
}
buf.extend_from_slice(run.bytes);
}
let end = buf.len();
let mut pos = 0usize;
let mut out = Vec::new();
let mut rowid = 0i64;
let mut first = true;
while pos < end {
let absolute = first || abs_at.contains(&pos);
let d = read_varint(&buf, &mut pos)? as i64;
rowid = if absolute { d } else { rowid.wrapping_add(d) };
first = false;
let cols = decode_poslist(&buf, &mut pos)?;
if pos > end {
return None;
}
out.push(DecodedPosting { rowid, cols });
}
if pos != end {
return None;
}
Some(out)
}
fn gather_doclist_runs<'a>(
leaves: &'a [&'a [u8]],
start_leaf: usize,
start_ti: usize,
start_off: usize,
leaf_views: &[LeafView],
) -> Option<Vec<DoclistRun<'a>>> {
let mut runs: Vec<DoclistRun<'a>> = Vec::new();
let first_view = &leaf_views[start_leaf];
let first_next_term = first_view.terms.get(start_ti + 1).map(|r| r.rec_start);
let first_end = first_next_term.unwrap_or(first_view.footer_off);
if first_end < start_off || first_end > first_view.footer_off {
return None;
}
runs.push(DoclistRun {
bytes: leaves[start_leaf].get(start_off..first_end)?,
abs_start: true, });
if first_next_term.is_some() {
return Some(runs); }
let mut li = start_leaf + 1;
while li < leaves.len() {
let view = &leaf_views[li];
let next_term = view.terms.first().map(|r| r.rec_start);
let boundary = next_term.unwrap_or(view.footer_off);
let tail_end = if view.first_rowid_off == 0 {
boundary
} else {
view.first_rowid_off
};
if tail_end < 4 || tail_end > boundary || boundary > view.footer_off {
return None;
}
runs.push(DoclistRun {
bytes: leaves[li].get(4..tail_end)?,
abs_start: false,
});
if view.first_rowid_off != 0 {
runs.push(DoclistRun {
bytes: leaves[li].get(view.first_rowid_off..boundary)?,
abs_start: true,
});
}
if next_term.is_some() {
break;
}
li += 1;
}
Some(runs)
}
#[cfg(test)]
pub(crate) fn decode_term(leaves: &[&[u8]], term: &[u8]) -> Option<Vec<DecodedPosting>> {
let key = term_key(term);
let mut views: Vec<LeafView> = Vec::with_capacity(leaves.len());
for leaf in leaves {
views.push(parse_leaf(leaf)?);
}
for (li, view) in views.iter().enumerate() {
for (ti, rec) in view.terms.iter().enumerate() {
if rec.key != key {
continue;
}
let runs = gather_doclist_runs(leaves, li, ti, rec.doclist_start, &views)?;
return decode_spanning_doclist(&runs);
}
}
None
}
fn decode_prefix_strict(leaves: &[&[u8]], prefix: &[u8]) -> SegDecode {
let want = term_key(prefix);
let mut views: Vec<LeafView> = Vec::with_capacity(leaves.len());
for leaf in leaves {
match parse_leaf(leaf) {
Some(v) => views.push(v),
None => return SegDecode::Bail,
}
}
let mut out: Vec<DecodedPosting> = Vec::new();
for (li, view) in views.iter().enumerate() {
for (ti, rec) in view.terms.iter().enumerate() {
if !rec.key.starts_with(&want) {
continue;
}
match gather_doclist_runs(leaves, li, ti, rec.doclist_start, &views)
.and_then(|runs| decode_spanning_doclist(&runs))
{
Some(postings) => out.extend(postings),
None => return SegDecode::Bail, }
}
}
SegDecode::Postings(out)
}
struct SegmentLoc {
segid: i64,
pgno_first: i64,
pgno_last: i64,
}
fn all_segments(structure: &[u8]) -> Option<Vec<SegmentLoc>> {
let mut pos = 4usize; let n_level = read_varint(structure, &mut pos)?;
let n_segment = read_varint(structure, &mut pos)?;
let _n_write_counter = read_varint(structure, &mut pos)?;
if n_level == 0 || n_segment == 0 {
return None; }
let mut segs: Vec<SegmentLoc> = Vec::new();
for _ in 0..n_level {
let _n_merge = read_varint(structure, &mut pos)?;
let n_seg = read_varint(structure, &mut pos)?;
for _ in 0..n_seg {
let segid = read_varint(structure, &mut pos)? as i64;
let pgno_first = read_varint(structure, &mut pos)? as i64;
let pgno_last = read_varint(structure, &mut pos)? as i64;
if segid <= 0 || pgno_first < 1 || pgno_last < pgno_first {
return None;
}
segs.push(SegmentLoc {
segid,
pgno_first,
pgno_last,
});
}
}
if segs.len() as u64 != n_segment {
return None;
}
Some(segs)
}
fn segments_leaves(data: &[(i64, Vec<u8>)]) -> Option<Vec<Vec<&[u8]>>> {
let structure = &data.iter().find(|(id, _)| *id == STRUCTURE_ROWID)?.1;
let locs = all_segments(structure)?;
let mut out: Vec<Vec<&[u8]>> = Vec::with_capacity(locs.len());
for loc in &locs {
let mut leaves: Vec<&[u8]> = Vec::new();
for pgno in loc.pgno_first..=loc.pgno_last {
let rid = segment_leaf_rowid(loc.segid, pgno);
let blob = &data.iter().find(|(id, _)| *id == rid)?.1;
leaves.push(blob.as_slice());
}
out.push(leaves);
}
#[cfg(test)]
INDEX_ROUTE_HITS.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
Some(out)
}
enum SegDecode {
Bail,
Postings(Vec<DecodedPosting>),
}
fn decode_term_strict(leaves: &[&[u8]], term: &[u8]) -> SegDecode {
let key = term_key(term);
let mut views: Vec<LeafView> = Vec::with_capacity(leaves.len());
for leaf in leaves {
match parse_leaf(leaf) {
Some(v) => views.push(v),
None => return SegDecode::Bail, }
}
for (li, view) in views.iter().enumerate() {
for (ti, rec) in view.terms.iter().enumerate() {
if rec.key != key {
continue;
}
return match gather_doclist_runs(leaves, li, ti, rec.doclist_start, &views)
.and_then(|runs| decode_spanning_doclist(&runs))
{
Some(postings) => SegDecode::Postings(postings),
None => SegDecode::Bail,
};
}
}
SegDecode::Postings(Vec::new()) }
fn merge_segments(
data: &[(i64, Vec<u8>)],
decode: impl Fn(&[&[u8]]) -> SegDecode,
) -> Option<Vec<DecodedPosting>> {
let segments = segments_leaves(data)?;
let mut all: Vec<DecodedPosting> = Vec::new();
let mut seen_segidx: Vec<(i64, usize)> = Vec::new(); for (si, leaves) in segments.iter().enumerate() {
match decode(leaves) {
SegDecode::Bail => return None,
SegDecode::Postings(postings) => {
for p in postings {
seen_segidx.push((p.rowid, si));
all.push(p);
}
}
}
}
seen_segidx.sort_unstable();
for w in seen_segidx.windows(2) {
if w[0].0 == w[1].0 && w[0].1 != w[1].1 {
return None;
}
}
all.sort_by_key(|p| p.rowid);
all.dedup_by_key(|p| p.rowid);
Some(all)
}
fn decode_poslist_keepdel(buf: &[u8], pos: &mut usize) -> Option<(Vec<Vec<u32>>, bool)> {
let size2 = read_varint(buf, pos)?;
let del = (size2 & 1) != 0;
let content_len = (size2 / 2) as usize;
let end = pos.checked_add(content_len)?;
if end > buf.len() {
return None;
}
let mut cols: Vec<Vec<u32>> = Vec::new();
let mut col = 0usize;
let mut p = *pos;
cols.push(Vec::new());
while p < end {
if buf[p] == 0x01 {
p += 1;
let c = read_varint(buf, &mut p)? as usize;
col = c;
while cols.len() <= col {
cols.push(Vec::new());
}
} else {
let raw = read_varint(buf, &mut p)?;
if raw < 2 {
return None;
}
let delta = (raw - 2) as u32;
let next = if cols[col].is_empty() {
delta
} else {
cols[col].last().copied()?.checked_add(delta)?
};
cols[col].push(next);
}
}
*pos = end;
Some((cols, del))
}
fn decode_doclist_keepdel(runs: &[DoclistRun]) -> Option<Vec<Posting>> {
let mut buf: Vec<u8> = Vec::new();
let mut abs_at: Vec<usize> = Vec::new();
for run in runs {
if run.abs_start && !run.bytes.is_empty() {
abs_at.push(buf.len());
}
buf.extend_from_slice(run.bytes);
}
let end = buf.len();
let mut pos = 0usize;
let mut out = Vec::new();
let mut rowid = 0i64;
let mut first = true;
while pos < end {
let absolute = first || abs_at.contains(&pos);
let d = read_varint(&buf, &mut pos)? as i64;
rowid = if absolute { d } else { rowid.wrapping_add(d) };
first = false;
let (cols, del) = decode_poslist_keepdel(&buf, &mut pos)?;
if pos > end {
return None;
}
out.push(Posting { rowid, cols, del });
}
if pos != end {
return None;
}
Some(out)
}
pub(crate) fn read_segment_postings(leaves: &[&[u8]]) -> Option<Vec<(Vec<u8>, Vec<Posting>)>> {
let mut views: Vec<LeafView> = Vec::with_capacity(leaves.len());
for leaf in leaves {
views.push(parse_leaf(leaf)?);
}
let mut out: Vec<(Vec<u8>, Vec<Posting>)> = Vec::new();
for (li, view) in views.iter().enumerate() {
for (ti, rec) in view.terms.iter().enumerate() {
if rec.key.first() != Some(&MAIN_PREFIX) {
return None;
}
let term = rec.key.get(1..)?.to_vec();
let runs = gather_doclist_runs(leaves, li, ti, rec.doclist_start, &views)?;
let postings = decode_doclist_keepdel(&runs)?;
out.push((term, postings));
}
}
Some(out)
}
pub(crate) fn read_segment_postings_full(leaves: &[&[u8]]) -> Option<Vec<(Vec<u8>, Vec<Posting>)>> {
let mut views: Vec<LeafView> = Vec::with_capacity(leaves.len());
for leaf in leaves {
views.push(parse_leaf(leaf)?);
}
let mut out: Vec<(Vec<u8>, Vec<Posting>)> = Vec::new();
for (li, view) in views.iter().enumerate() {
for (ti, rec) in view.terms.iter().enumerate() {
if rec.key.is_empty() {
return None; }
let key = rec.key.clone();
let runs = gather_doclist_runs(leaves, li, ti, rec.doclist_start, &views)?;
let postings = decode_doclist_keepdel(&runs)?;
out.push((key, postings));
}
}
Some(out)
}
fn merge_level_postings(
segs: &[Vec<(Vec<u8>, Vec<Posting>)>],
b_oldest: bool,
) -> Vec<(Vec<u8>, Vec<Posting>)> {
use alloc::collections::BTreeMap;
let mut map: BTreeMap<Vec<u8>, BTreeMap<i64, Posting>> = BTreeMap::new();
for seg in segs {
for (term, postings) in seg {
let e = map.entry(term.clone()).or_default();
for p in postings {
e.insert(p.rowid, p.clone());
}
}
}
let mut out: Vec<(Vec<u8>, Vec<Posting>)> = Vec::new();
for (term, by_rowid) in map {
let mut ps: Vec<Posting> = Vec::new();
for (_rowid, p) in by_rowid {
let empty = p.cols.iter().all(|c| c.is_empty());
if empty && (b_oldest || !p.del) {
continue;
}
ps.push(p);
}
if !ps.is_empty() {
out.push((term, ps));
}
}
out
}
pub(crate) fn merge_segments_keepdel(
seg_leaves: &[Vec<&[u8]>],
b_oldest: bool,
) -> Option<Vec<(Vec<u8>, Vec<Posting>)>> {
let mut segs: Vec<Vec<(Vec<u8>, Vec<Posting>)>> = Vec::with_capacity(seg_leaves.len());
for leaves in seg_leaves {
segs.push(read_segment_postings(leaves)?);
}
Some(merge_level_postings(&segs, b_oldest))
}
pub(crate) fn merge_segments_keepdel_full(
seg_leaves: &[Vec<&[u8]>],
b_oldest: bool,
) -> Option<Vec<(Vec<u8>, Vec<Posting>)>> {
let mut segs: Vec<Vec<(Vec<u8>, Vec<Posting>)>> = Vec::with_capacity(seg_leaves.len());
for leaves in seg_leaves {
segs.push(read_segment_postings_full(leaves)?);
}
Some(merge_level_postings(&segs, b_oldest))
}
fn read_main_segment_postings(leaves: &[&[u8]]) -> Option<Vec<(Vec<u8>, Vec<Posting>)>> {
let mut views: Vec<LeafView> = Vec::with_capacity(leaves.len());
for leaf in leaves {
views.push(parse_leaf(leaf)?);
}
let mut out: Vec<(Vec<u8>, Vec<Posting>)> = Vec::new();
for (li, view) in views.iter().enumerate() {
for (ti, rec) in view.terms.iter().enumerate() {
match rec.key.first() {
Some(&MAIN_PREFIX) => {}
Some(_) => continue,
None => return None, }
let term = rec.key.get(1..)?.to_vec();
let runs = gather_doclist_runs(leaves, li, ti, rec.doclist_start, &views)?;
let postings = decode_doclist_keepdel(&runs)?;
out.push((term, postings));
}
}
Some(out)
}
pub(crate) enum MainIndexScan {
Clean(Vec<(Vec<u8>, Vec<Posting>)>),
Skip,
Malformed,
}
pub(crate) fn scan_main_index(data: &[(i64, Vec<u8>)]) -> MainIndexScan {
use alloc::collections::BTreeMap;
let Some(structure) = data
.iter()
.find(|(id, _)| *id == STRUCTURE_ROWID)
.map(|(_, b)| b.as_slice())
else {
return MainIndexScan::Skip; };
let mut pos = 4usize; let (Some(n_level), Some(n_segment), Some(_wc)) = (
read_varint(structure, &mut pos),
read_varint(structure, &mut pos),
read_varint(structure, &mut pos),
) else {
return MainIndexScan::Skip;
};
if n_level == 0 || n_segment == 0 {
return MainIndexScan::Clean(Vec::new()); }
let mut locs: Vec<SegmentLoc> = Vec::new();
for _ in 0..n_level {
let (Some(_n_merge), Some(n_seg)) = (
read_varint(structure, &mut pos),
read_varint(structure, &mut pos),
) else {
return MainIndexScan::Skip;
};
for _ in 0..n_seg {
let (Some(segid), Some(pgno_first), Some(pgno_last)) = (
read_varint(structure, &mut pos),
read_varint(structure, &mut pos),
read_varint(structure, &mut pos),
) else {
return MainIndexScan::Skip;
};
let (segid, pgno_first, pgno_last) =
(segid as i64, pgno_first as i64, pgno_last as i64);
if segid <= 0 || pgno_first < 1 || pgno_last < pgno_first {
return MainIndexScan::Skip;
}
locs.push(SegmentLoc {
segid,
pgno_first,
pgno_last,
});
}
}
if locs.len() as u64 != n_segment {
return MainIndexScan::Skip;
}
let mut per_term: BTreeMap<Vec<u8>, Vec<Posting>> = BTreeMap::new();
let mut seen: Vec<(Vec<u8>, i64, usize)> = Vec::new(); for (si, loc) in locs.iter().enumerate() {
let mut leaves: Vec<&[u8]> = Vec::new();
for pgno in loc.pgno_first..=loc.pgno_last {
let rid = segment_leaf_rowid(loc.segid, pgno);
match data.iter().find(|(id, _)| *id == rid) {
Some((_, b)) => leaves.push(b.as_slice()),
None => return MainIndexScan::Malformed, }
}
let Some(terms) = read_main_segment_postings(&leaves) else {
return MainIndexScan::Skip; };
for (term, postings) in terms {
for p in &postings {
if p.del {
return MainIndexScan::Skip; }
seen.push((term.clone(), p.rowid, si));
}
per_term.entry(term).or_default().extend(postings);
}
}
seen.sort();
for w in seen.windows(2) {
if w[0].0 == w[1].0 && w[0].1 == w[1].1 && w[0].2 != w[1].2 {
return MainIndexScan::Skip;
}
}
let mut out: Vec<(Vec<u8>, Vec<Posting>)> = Vec::with_capacity(per_term.len());
for (term, mut ps) in per_term {
ps.sort_by_key(|p| p.rowid);
out.push((term, ps));
}
MainIndexScan::Clean(out)
}
pub(crate) fn lookup_term_rowids(data: &[(i64, Vec<u8>)], term: &[u8]) -> Option<Vec<i64>> {
decode_term_in_data(data, term).map(|postings| postings.into_iter().map(|p| p.rowid).collect())
}
pub(crate) fn lookup_term_rowids_in_column(
data: &[(i64, Vec<u8>)],
term: &[u8],
column: usize,
) -> Option<Vec<i64>> {
decode_term_in_data(data, term).map(|postings| {
postings
.into_iter()
.filter(|p| p.cols.get(column).is_some_and(|c| !c.is_empty()))
.map(|p| p.rowid)
.collect()
})
}
fn prefix_rowids(
mut postings: Vec<DecodedPosting>,
keep: impl Fn(&DecodedPosting) -> bool,
) -> Vec<i64> {
let mut rowids: Vec<i64> = postings
.drain(..)
.filter(|p| keep(p))
.map(|p| p.rowid)
.collect();
rowids.sort_unstable();
rowids.dedup();
rowids
}
pub(crate) fn lookup_prefix_rowids(data: &[(i64, Vec<u8>)], prefix: &[u8]) -> Option<Vec<i64>> {
let postings = merge_segments(data, |leaves| decode_prefix_strict(leaves, prefix))?;
Some(prefix_rowids(postings, |_| true))
}
pub(crate) fn lookup_prefix_rowids_in_column(
data: &[(i64, Vec<u8>)],
prefix: &[u8],
column: usize,
) -> Option<Vec<i64>> {
let postings = merge_segments(data, |leaves| decode_prefix_strict(leaves, prefix))?;
Some(prefix_rowids(postings, |p| {
p.cols.get(column).is_some_and(|c| !c.is_empty())
}))
}
fn decode_term_in_data(data: &[(i64, Vec<u8>)], term: &[u8]) -> Option<Vec<DecodedPosting>> {
merge_segments(data, |leaves| decode_term_strict(leaves, term))
}
fn col(p: &DecodedPosting, c: usize) -> &[u32] {
p.cols.get(c).map(Vec::as_slice).unwrap_or(&[])
}
fn phrase_intersect(
postings_a: &[DecodedPosting],
postings_b: &[DecodedPosting],
adj: impl Fn(&DecodedPosting, &DecodedPosting) -> bool,
) -> Vec<i64> {
let mut out = Vec::new();
let (mut i, mut j) = (0usize, 0usize);
while i < postings_a.len() && j < postings_b.len() {
let (a, b) = (&postings_a[i], &postings_b[j]);
match a.rowid.cmp(&b.rowid) {
core::cmp::Ordering::Less => i += 1,
core::cmp::Ordering::Greater => j += 1,
core::cmp::Ordering::Equal => {
if adj(a, b) {
out.push(a.rowid);
}
i += 1;
j += 1;
}
}
}
out
}
fn phrase_run_matches(cols: &[Vec<&[u32]>], column: Option<usize>) -> bool {
debug_assert!(!cols.is_empty());
let k = cols.len();
let ncols = cols.iter().map(Vec::len).max().unwrap_or(0);
let in_col = |c: usize| -> bool {
let first = cols[0].get(c).copied().unwrap_or(&[]);
first.iter().any(|&p| {
(1..k).all(|i| {
let want = match p.checked_add(i as u32) {
Some(w) => w,
None => return false,
};
let list = cols[i].get(c).copied().unwrap_or(&[]);
list.binary_search(&want).is_ok()
})
})
};
match column {
Some(c) => in_col(c),
None => (0..ncols).any(in_col),
}
}
fn phrase_intersect_k(postings: &[Vec<DecodedPosting>], column: Option<usize>) -> Vec<i64> {
let k = postings.len();
if k == 0 || postings.iter().any(Vec::is_empty) {
return Vec::new();
}
let mut out = Vec::new();
let mut idx = alloc::vec![0usize; k];
loop {
let mut min_rowid = i64::MAX;
for (i, p) in postings.iter().enumerate() {
if idx[i] >= p.len() {
return out;
}
min_rowid = min_rowid.min(p[idx[i]].rowid);
}
let all_equal = postings
.iter()
.zip(&idx)
.all(|(p, &j)| p[j].rowid == min_rowid);
if all_equal {
let cols: Vec<Vec<&[u32]>> = postings
.iter()
.zip(&idx)
.map(|(p, &j)| p[j].cols.iter().map(Vec::as_slice).collect())
.collect();
if phrase_run_matches(&cols, column) {
out.push(min_rowid);
}
}
for (i, p) in postings.iter().enumerate() {
if p[idx[i]].rowid == min_rowid {
idx[i] += 1;
}
}
}
}
fn decode_terms_multiseg(
data: &[(i64, Vec<u8>)],
terms: &[&[u8]],
) -> Option<Vec<Vec<DecodedPosting>>> {
let segments = segments_leaves(data)?;
let mut per_term: Vec<Vec<DecodedPosting>> = alloc::vec![Vec::new(); terms.len()];
let mut seen_segidx: Vec<(i64, usize)> = Vec::new();
for (si, leaves) in segments.iter().enumerate() {
for (ti, term) in terms.iter().enumerate() {
match decode_term_strict(leaves, term) {
SegDecode::Bail => return None,
SegDecode::Postings(postings) => {
for p in postings {
seen_segidx.push((p.rowid, si));
per_term[ti].push(p);
}
}
}
}
}
seen_segidx.sort_unstable();
for w in seen_segidx.windows(2) {
if w[0].0 == w[1].0 && w[0].1 != w[1].1 {
return None;
}
}
for postings in &mut per_term {
postings.sort_by_key(|p| p.rowid);
}
Some(per_term)
}
pub(crate) fn lookup_phrase_rowids_k(data: &[(i64, Vec<u8>)], terms: &[&[u8]]) -> Option<Vec<i64>> {
if terms.len() < 2 {
return None;
}
let postings = decode_terms_multiseg(data, terms)?;
Some(phrase_intersect_k(&postings, None))
}
pub(crate) fn lookup_phrase_rowids_in_column_k(
data: &[(i64, Vec<u8>)],
terms: &[&[u8]],
column: usize,
) -> Option<Vec<i64>> {
if terms.len() < 2 {
return None;
}
let postings = decode_terms_multiseg(data, terms)?;
Some(phrase_intersect_k(&postings, Some(column)))
}
fn near_within_in_column(pa: &[u32], pb: &[u32], n: u32) -> bool {
let limit = n.saturating_add(1);
let (mut i, mut j) = (0usize, 0usize);
while i < pa.len() && j < pb.len() {
let (a, b) = (pa[i], pb[j]);
let gap = a.abs_diff(b);
if gap <= limit {
return true;
}
if a < b {
i += 1;
} else {
j += 1;
}
}
false
}
fn near_matches(a: &DecodedPosting, b: &DecodedPosting, n: u32) -> bool {
let ncols = a.cols.len().max(b.cols.len());
(0..ncols).any(|c| near_within_in_column(col(a, c), col(b, c), n))
}
pub(crate) fn lookup_near_rowids(
data: &[(i64, Vec<u8>)],
term_a: &[u8],
term_b: &[u8],
n: u32,
) -> Option<Vec<i64>> {
let terms: [&[u8]; 2] = [term_a, term_b];
let postings = decode_terms_multiseg(data, &terms)?;
Some(phrase_intersect(&postings[0], &postings[1], |a, b| {
near_matches(a, b, n)
}))
}
fn rowids_intersect(a: &[i64], b: &[i64]) -> Vec<i64> {
let mut out = Vec::new();
let (mut i, mut j) = (0usize, 0usize);
while i < a.len() && j < b.len() {
match a[i].cmp(&b[j]) {
core::cmp::Ordering::Less => i += 1,
core::cmp::Ordering::Greater => j += 1,
core::cmp::Ordering::Equal => {
out.push(a[i]);
i += 1;
j += 1;
}
}
}
out
}
fn rowids_union(a: &[i64], b: &[i64]) -> Vec<i64> {
let mut out = Vec::with_capacity(a.len() + b.len());
let (mut i, mut j) = (0usize, 0usize);
while i < a.len() && j < b.len() {
match a[i].cmp(&b[j]) {
core::cmp::Ordering::Less => {
out.push(a[i]);
i += 1;
}
core::cmp::Ordering::Greater => {
out.push(b[j]);
j += 1;
}
core::cmp::Ordering::Equal => {
out.push(a[i]);
i += 1;
j += 1;
}
}
}
out.extend_from_slice(&a[i..]);
out.extend_from_slice(&b[j..]);
out
}
fn rowids_difference(a: &[i64], b: &[i64]) -> Vec<i64> {
let mut out = Vec::new();
let (mut i, mut j) = (0usize, 0usize);
while i < a.len() {
if j >= b.len() {
out.extend_from_slice(&a[i..]);
break;
}
match a[i].cmp(&b[j]) {
core::cmp::Ordering::Less => {
out.push(a[i]);
i += 1;
}
core::cmp::Ordering::Greater => j += 1,
core::cmp::Ordering::Equal => {
i += 1;
j += 1;
}
}
}
out
}
#[cfg(feature = "fts5")]
fn eval_bool_tree(data: &[(i64, Vec<u8>)], tree: &crate::vtab::Fts5BoolTree) -> Option<Vec<i64>> {
use crate::vtab::{Fts5BoolOp, Fts5BoolTree};
match tree {
Fts5BoolTree::Leaf(term) => Some(
decode_term_in_data(data, term)?
.into_iter()
.map(|p| p.rowid)
.collect(),
),
Fts5BoolTree::Op(op, a, b) => {
let ra = eval_bool_tree(data, a)?;
let rb = eval_bool_tree(data, b)?;
Some(match op {
Fts5BoolOp::And => rowids_intersect(&ra, &rb),
Fts5BoolOp::Or => rowids_union(&ra, &rb),
Fts5BoolOp::Not => rowids_difference(&ra, &rb),
})
}
}
}
pub(crate) fn lookup_bool_tree_rowids(
data: &[(i64, Vec<u8>)],
tree: &crate::vtab::Fts5BoolTree,
) -> Option<Vec<i64>> {
eval_bool_tree(data, tree)
}
#[cfg(test)]
mod tests {
use super::*;
use alloc::{format, string::ToString, vec};
#[test]
fn segstructure_roundtrips_and_allocates_segid() {
let raw = &[
0x00, 0x00, 0x00, 0x00, 0x01, 0x02, 0x02, 0x00, 0x02, 0x01, 0x01, 0x01, 0x02, 0x01, 0x01, ];
let s = SegStructure::parse(raw).unwrap();
assert_eq!(s.levels.len(), 1);
assert_eq!(s.levels[0].segs.len(), 2);
assert_eq!(s.write_counter, 2);
assert_eq!(s.encode(), raw);
assert_eq!(s.allocate_segid(), 3);
}
#[test]
fn segstructure_append_and_promote_down() {
let mut s = SegStructure {
cookie: 0,
write_counter: 0,
levels: Vec::new(),
};
s.append_level0(1, 1);
assert_eq!(s.levels.len(), 1);
assert_eq!(
s.levels[0].segs,
vec![StructSeg {
segid: 1,
pgno_first: 1,
pgno_last: 1
}]
);
assert_eq!(s.write_counter, 1);
let mut s = SegStructure {
cookie: 0,
write_counter: 16,
levels: vec![
StructLevel {
n_merge: 0,
segs: Vec::new(),
},
StructLevel {
n_merge: 0,
segs: vec![StructSeg {
segid: 17,
pgno_first: 1,
pgno_last: 1,
}],
},
],
};
s.append_level0(1, 1); assert_eq!(
s.levels[0].segs,
vec![
StructSeg {
segid: 17,
pgno_first: 1,
pgno_last: 1
},
StructSeg {
segid: 1,
pgno_first: 1,
pgno_last: 1
},
],
"promoted-down segment precedes the newly appended one"
);
assert!(s.levels[1].segs.is_empty());
}
fn p(rowid: i64, cols: &[&[u32]]) -> Posting {
Posting {
rowid,
cols: cols.iter().map(|c| c.to_vec()).collect(),
del: false,
}
}
#[test]
fn empty_table_structure_and_averages() {
let seg = build_segment(&[], 0, &[0], &[], 1000, 0);
assert_eq!(seg.data[0], (AVERAGES_ROWID, Vec::new()));
assert_eq!(seg.data[1], (STRUCTURE_ROWID, vec![0, 0, 0, 0, 0, 0, 0]));
assert_eq!(seg.data.len(), 2); assert!(seg.idx.is_empty());
}
#[test]
fn single_term_single_doc_matches_known_bytes() {
let terms = vec![(b"a".to_vec(), vec![p(1, &[&[0]])])];
let seg = build_segment(&terms, 1, &[1], &[(1, vec![1])], 1000, 0);
let leaf = &seg
.data
.iter()
.find(|(id, _)| *id == segment_leaf_rowid(1, 1))
.unwrap()
.1;
assert_eq!(
leaf,
&vec![0, 0, 0, 0x0A, 0x02, 0x30, 0x61, 0x01, 0x02, 0x02, 0x04]
);
assert_eq!(seg.data[0].1, vec![0x01, 0x01]);
assert_eq!(seg.idx.len(), 1);
assert_eq!(seg.idx[0].pgno, 2);
assert!(seg.idx[0].term.is_empty());
assert_eq!(seg.docsize, vec![(1, vec![1])]);
}
#[test]
fn multi_column_poslist_bytes() {
let terms = vec![("hello".to_string().into_bytes(), vec![p(1, &[&[0], &[0]])])];
let seg = build_segment(&terms, 1, &[1, 1], &[(1, vec![1, 1])], 1000, 0);
let leaf = &seg
.data
.iter()
.find(|(id, _)| *id == segment_leaf_rowid(1, 1))
.unwrap()
.1;
let expected = vec![
0, 0, 0, 0x11, 0x06, 0x30, b'h', b'e', b'l', b'l', b'o', 0x01, 0x08, 0x02, 0x01, 0x01,
0x02, 0x04,
];
assert_eq!(leaf, &expected);
}
fn leaves_of(seg: &Segment) -> Vec<Vec<u8>> {
let mut out = Vec::new();
let mut pgno = 1i64;
loop {
let rid = segment_leaf_rowid(1, pgno);
match seg.data.iter().find(|(id, _)| *id == rid) {
Some((_, blob)) => out.push(blob.clone()),
None => break,
}
pgno += 1;
}
out
}
fn decode(seg: &Segment, term: &[u8]) -> Option<Vec<DecodedPosting>> {
let leaves = leaves_of(seg);
let refs: Vec<&[u8]> = leaves.iter().map(|l| l.as_slice()).collect();
decode_term(&refs, term)
}
fn dp(rowid: i64, cols: &[&[u32]]) -> DecodedPosting {
DecodedPosting {
rowid,
cols: cols.iter().map(|c| c.to_vec()).collect(),
}
}
fn dlidx_pages(seg: &Segment) -> Vec<(i64, Vec<u8>)> {
let mut out: Vec<(i64, Vec<u8>)> = seg
.data
.iter()
.filter(|(id, _)| (*id & (1 << 36)) != 0)
.map(|(id, b)| (*id, b.clone()))
.collect();
out.sort_by_key(|(id, _)| *id);
out
}
type Corpus = (
Vec<(Vec<u8>, Vec<Posting>)>,
u64,
Vec<u64>,
Vec<(i64, Vec<u64>)>,
);
fn single_token_corpus(tok: &[u8], n: i64) -> Corpus {
let postings: Vec<Posting> = (1..=n).map(|r| p(r, &[&[0]])).collect();
let terms = vec![(tok.to_vec(), postings)];
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=n).map(|r| (r, vec![1])).collect();
(terms, n as u64, vec![n as u64], doc_sizes)
}
#[test]
fn dlidx_absent_when_doclist_fits_one_leaf() {
let (terms, ndoc, tot, sizes) = single_token_corpus(b"shared", 50);
let seg = build_segment(&terms, ndoc, &tot, &sizes, 4050, 0);
assert!(dlidx_pages(&seg).is_empty());
assert_eq!(seg.idx.len(), 1);
assert_eq!(seg.idx[0].pgno & 1, 0);
}
#[test]
fn dlidx_absent_below_threshold() {
let (terms, ndoc, tot, sizes) = single_token_corpus(b"shared", 2500);
let seg = build_segment(&terms, ndoc, &tot, &sizes, 4050, 0);
let leaves = leaves_of(&seg);
assert!(leaves.len() >= 2, "expected a multi-leaf spill");
assert!(
dlidx_pages(&seg).is_empty(),
"fewer than {MIN_DLIDX_SIZE} continuation leaves must not add a dlidx"
);
assert_eq!(seg.idx.iter().filter(|r| r.pgno & 1 == 1).count(), 0);
assert_eq!(decode(&seg, b"shared").unwrap().len(), 2500);
}
#[test]
fn dlidx_bytes_match_sqlite_6000_docs() {
let (terms, ndoc, tot, sizes) = single_token_corpus(b"shared", 6000);
let seg = build_segment(&terms, ndoc, &tot, &sizes, 4050, 0);
assert_eq!(leaves_of(&seg).len(), 5);
let dl = dlidx_pages(&seg);
assert_eq!(dl.len(), 1);
assert_eq!(dl[0].0, dlidx_rowid(1, 0, 1));
assert_eq!(dl[0].1, hex("00028A438A458A458A45"));
assert_eq!(seg.idx.len(), 1);
assert_eq!(seg.idx[0].pgno, (1 << 1) | 1); assert!(seg.idx[0].term.is_empty());
let got = decode(&seg, b"shared").unwrap();
assert_eq!(got.len(), 6000);
assert_eq!(got.first().unwrap().rowid, 1);
assert_eq!(got.last().unwrap().rowid, 6000);
}
#[test]
fn dlidx_bytes_match_sqlite_8000_docs() {
let (terms, ndoc, tot, sizes) = single_token_corpus(b"shared", 8000);
let seg = build_segment(&terms, ndoc, &tot, &sizes, 4050, 0);
assert_eq!(leaves_of(&seg).len(), 6);
let dl = dlidx_pages(&seg);
assert_eq!(dl.len(), 1);
assert_eq!(dl[0].0, dlidx_rowid(1, 0, 1));
assert_eq!(dl[0].1, hex("00028A438A458A458A458A45"));
assert_eq!(seg.idx[0].pgno & 1, 1);
assert_eq!(decode(&seg, b"shared").unwrap().len(), 8000);
}
#[test]
fn high_frequency_spanning_term_takes_index_route() {
let (terms, ndoc, tot, sizes) = single_token_corpus(b"shared", 8000);
let seg = build_segment(&terms, ndoc, &tot, &sizes, 4050, 0);
assert!(leaves_of(&seg).len() > MIN_DLIDX_SIZE);
assert_eq!(dlidx_pages(&seg).len(), 1);
assert_eq!(
seg.idx[0].pgno & 1,
1,
"term-start leaf carries a dlidx flag"
);
let before = INDEX_ROUTE_HITS.load(core::sync::atomic::Ordering::Relaxed);
let rowids = lookup_term_rowids(&seg.data, b"shared")
.expect("a dlidx-bearing single-segment index must be servable, not scanned");
assert_eq!(rowids.len(), 8000);
assert_eq!(*rowids.first().unwrap(), 1);
assert_eq!(*rowids.last().unwrap(), 8000);
assert!(
INDEX_ROUTE_HITS.load(core::sync::atomic::Ordering::Relaxed) > before,
"the high-frequency spanning term must take the index route, not scan _content"
);
}
fn hex(s: &str) -> Vec<u8> {
(0..s.len())
.step_by(2)
.map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
.collect()
}
fn dlidx_level_entries(page: &[u8]) -> Vec<(i64, i64)> {
let mut out = Vec::new();
let mut pos = 1usize; let mut pgno = read_varint(page, &mut pos).unwrap() as i64;
let mut rowid = read_varint(page, &mut pos).unwrap() as i64;
out.push((pgno, rowid));
while pos < page.len() {
while pos < page.len() && page[pos] == 0 {
pgno += 1;
pos += 1;
}
if pos >= page.len() {
break;
}
let d = read_varint(page, &mut pos).unwrap() as i64;
pgno += 1;
rowid += d;
out.push((pgno, rowid));
}
out
}
fn dlidx_page(seg: &Segment, height: i64, pgno: i64) -> Option<&[u8]> {
let rid = dlidx_rowid(1, height, pgno);
seg.data
.iter()
.find(|(id, _)| *id == rid)
.map(|(_, b)| b.as_slice())
}
fn leaf_first_rowids(seg: &Segment) -> alloc::collections::BTreeMap<i64, i64> {
let mut out = alloc::collections::BTreeMap::new();
for (id, blob) in &seg.data {
if *id <= STRUCTURE_ROWID || (*id & (1 << 36)) != 0 {
continue;
}
let pgno = *id & 0x7fff_ffff;
let fro = u16::from_be_bytes([blob[0], blob[1]]) as usize;
if fro == 0 {
continue;
}
let mut pos = fro;
if let Some(r) = read_varint(blob, &mut pos) {
out.insert(pgno, r as i64);
}
}
out
}
#[test]
fn dlidx_multi_level_is_consistent() {
let postings: Vec<Posting> = (1..=4000).map(|i| p(i * 1000, &[&[0]])).collect();
let terms = vec![(b"w".to_vec(), postings)];
let sizes: Vec<(i64, Vec<u64>)> = (1..=4000).map(|i| (i * 1000, vec![1])).collect();
let seg = build_segment(&terms, 4000, &[4000], &sizes, 64, 0);
let n_h0 = seg
.data
.iter()
.filter(|(id, _)| (*id & (1 << 36)) != 0 && ((*id >> 31) & 0x1f) == 0)
.count();
let n_h1 = seg
.data
.iter()
.filter(|(id, _)| (*id & (1 << 36)) != 0 && ((*id >> 31) & 0x1f) == 1)
.count();
assert!(
n_h0 > 1,
"expected multiple height-0 dlidx pages, got {n_h0}"
);
assert_eq!(n_h1, 1, "expected exactly one height-1 root");
let root = dlidx_page(&seg, 1, 1).expect("root at height 1, pgno 1");
let root_entries = dlidx_level_entries(root);
let mut leaf_index: alloc::collections::BTreeMap<i64, i64> =
alloc::collections::BTreeMap::new();
for &(child_pgno, root_rowid) in &root_entries {
let child = dlidx_page(&seg, 0, child_pgno)
.unwrap_or_else(|| panic!("height-0 dlidx page {child_pgno} referenced by root"));
let entries = dlidx_level_entries(child);
assert_eq!(
entries[0].1, root_rowid,
"root's rowid for child {child_pgno} must equal that child's first rowid"
);
for (leaf_pgno, rowid) in entries {
leaf_index.insert(leaf_pgno, rowid);
}
}
let actual = leaf_first_rowids(&seg);
for (&leaf_pgno, &rowid) in &leaf_index {
assert_eq!(
actual.get(&leaf_pgno),
Some(&rowid),
"dlidx entry for leaf {leaf_pgno} must match the leaf's first rowid"
);
}
assert_eq!(leaf_index.len(), actual.len());
assert_eq!(leaf_index.keys().next(), Some(&2));
let got = decode(&seg, b"w").unwrap();
assert_eq!(got.len(), 4000);
assert_eq!(got.first().unwrap().rowid, 1000);
assert_eq!(got.last().unwrap().rowid, 4_000_000);
}
#[test]
fn decode_single_term_single_doc() {
let terms = vec![(b"a".to_vec(), vec![p(1, &[&[0]])])];
let seg = build_segment(&terms, 1, &[1], &[(1, vec![1])], 1000, 0);
assert_eq!(decode(&seg, b"a"), Some(vec![dp(1, &[&[0]])]));
assert_eq!(decode(&seg, b"z"), None);
assert_eq!(decode(&seg, b""), None);
}
#[test]
fn decode_multi_doc_rowid_deltas() {
let terms = vec![(
b"cat".to_vec(),
vec![p(1, &[&[0]]), p(3, &[&[2]]), p(7, &[&[1]])],
)];
let seg = build_segment(
&terms,
3,
&[3],
&[(1, vec![1]), (3, vec![3]), (7, vec![2])],
1000,
0,
);
assert_eq!(
decode(&seg, b"cat"),
Some(vec![dp(1, &[&[0]]), dp(3, &[&[2]]), dp(7, &[&[1]])])
);
}
#[test]
fn decode_term_multiple_positions_one_doc() {
let terms = vec![(b"the".to_vec(), vec![p(1, &[&[0, 2]])])];
let seg = build_segment(&terms, 1, &[3], &[(1, vec![3])], 1000, 0);
assert_eq!(decode(&seg, b"the"), Some(vec![dp(1, &[&[0, 2]])]));
}
#[test]
fn decode_prefix_compressed_terms() {
let terms = vec![
(b"apple".to_vec(), vec![p(1, &[&[0]])]),
(b"apply".to_vec(), vec![p(2, &[&[0]])]),
];
let seg = build_segment(&terms, 2, &[2], &[(1, vec![1]), (2, vec![1])], 1000, 0);
assert_eq!(decode(&seg, b"apple"), Some(vec![dp(1, &[&[0]])]));
assert_eq!(decode(&seg, b"apply"), Some(vec![dp(2, &[&[0]])]));
assert_eq!(decode(&seg, b"appl"), None);
}
#[test]
fn decode_multi_column_positions() {
let terms = vec![
(b"hello".to_vec(), vec![p(1, &[&[0], &[0]])]),
(b"there".to_vec(), vec![p(1, &[&[], &[1]])]),
];
let seg = build_segment(&terms, 1, &[1, 2], &[(1, vec![1, 2])], 1000, 0);
assert_eq!(decode(&seg, b"hello"), Some(vec![dp(1, &[&[0], &[0]])]));
assert_eq!(decode(&seg, b"there"), Some(vec![dp(1, &[&[], &[1]])]));
}
#[test]
fn decode_many_terms_one_leaf() {
let words: &[&[u8]] = &[b"alpha", b"beta", b"delta", b"gamma", b"omega"];
let terms: Vec<(Vec<u8>, Vec<Posting>)> = words
.iter()
.enumerate()
.map(|(i, w)| (w.to_vec(), vec![p(i as i64 + 1, &[&[0]])]))
.collect();
let doc_sizes: Vec<(i64, Vec<u64>)> =
(1..=words.len() as i64).map(|r| (r, vec![1])).collect();
let seg = build_segment(
&terms,
words.len() as u64,
&[words.len() as u64],
&doc_sizes,
1000,
0,
);
for (i, w) in words.iter().enumerate() {
assert_eq!(
decode(&seg, w),
Some(vec![dp(i as i64 + 1, &[&[0]])]),
"{w:?}"
);
}
assert_eq!(decode(&seg, b"missing"), None);
}
fn leaf_count(seg: &Segment) -> usize {
leaves_of(seg).len()
}
#[test]
fn decode_multi_leaf_term_pagination() {
let n = 40usize;
let words: Vec<Vec<u8>> = (0..n).map(|i| format!("term{i:03}").into_bytes()).collect();
let terms: Vec<(Vec<u8>, Vec<Posting>)> = words
.iter()
.enumerate()
.map(|(i, w)| (w.clone(), vec![p(i as i64 + 1, &[&[0]])]))
.collect();
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=n as i64).map(|r| (r, vec![1])).collect();
let seg = build_segment(&terms, n as u64, &[n as u64], &doc_sizes, 64, 0);
assert!(leaf_count(&seg) > 1, "pgsz 64 must split into many leaves");
for (i, w) in words.iter().enumerate() {
assert_eq!(
decode(&seg, w),
Some(vec![dp(i as i64 + 1, &[&[0]])]),
"term {w:?} on leaf pagination"
);
}
assert_eq!(decode(&seg, b"term999"), None);
}
#[test]
fn decode_doclist_spanning_leaves() {
let n = 40i64;
let postings: Vec<Posting> = (1..=n).map(|r| p(r, &[&[0]])).collect();
let terms = vec![(b"x".to_vec(), postings)];
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=n).map(|r| (r, vec![1])).collect();
let seg = build_segment(&terms, n as u64, &[n as u64], &doc_sizes, 64, 0);
assert!(leaf_count(&seg) > 1, "pgsz 64 must span the doclist");
let want: Vec<DecodedPosting> = (1..=n).map(|r| dp(r, &[&[0]])).collect();
assert_eq!(decode(&seg, b"x"), Some(want));
assert_eq!(decode(&seg, b"y"), None);
}
#[test]
fn decode_doclist_spanning_multi_position() {
let n = 30i64;
let postings: Vec<Posting> = (1..=n).map(|r| p(r, &[&[0, 3, 9, 15]])).collect();
let terms = vec![(b"w".to_vec(), postings)];
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=n).map(|r| (r, vec![16])).collect();
let seg = build_segment(&terms, n as u64, &[(16 * n) as u64], &doc_sizes, 48, 0);
assert!(leaf_count(&seg) > 1, "pgsz 48 must span the doclist");
let want: Vec<DecodedPosting> = (1..=n).map(|r| dp(r, &[&[0, 3, 9, 15]])).collect();
assert_eq!(decode(&seg, b"w"), Some(want));
}
#[test]
fn decode_mixed_pagination_and_spanning() {
let mut terms: Vec<(Vec<u8>, Vec<Posting>)> = Vec::new();
terms.push((b"heavy".to_vec(), (1..=25).map(|r| p(r, &[&[0]])).collect()));
for i in 0..20 {
let w = format!("light{i:02}").into_bytes();
terms.push((w, vec![p(100 + i as i64, &[&[1]])]));
}
terms.sort_by(|a, b| a.0.cmp(&b.0));
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=120).map(|r| (r, vec![1])).collect();
let seg = build_segment(&terms, 120, &[120], &doc_sizes, 56, 0);
assert!(leaf_count(&seg) > 2, "expected several leaves");
for (term, postings) in &terms {
let want: Vec<DecodedPosting> = postings
.iter()
.map(|p| DecodedPosting {
rowid: p.rowid,
cols: p.cols.clone(),
})
.collect();
assert_eq!(decode(&seg, term), Some(want), "term {term:?}");
}
assert_eq!(decode(&seg, b"absent"), None);
}
#[test]
fn lookup_rowids_single_segment_present_and_absent() {
let terms = vec![(
b"cat".to_vec(),
vec![p(1, &[&[0]]), p(3, &[&[2]]), p(7, &[&[1]])],
)];
let seg = build_segment(
&terms,
3,
&[3],
&[(1, vec![1]), (3, vec![3]), (7, vec![2])],
1000,
0,
);
assert_eq!(lookup_term_rowids(&seg.data, b"cat"), Some(vec![1, 3, 7]));
assert_eq!(lookup_term_rowids(&seg.data, b"dog"), Some(Vec::new()));
}
#[test]
fn lookup_rowids_empty_index_falls_back() {
let seg = build_segment(&[], 0, &[0], &[], 1000, 0);
assert_eq!(lookup_term_rowids(&seg.data, b"anything"), None);
}
#[test]
fn lookup_rowids_multi_leaf_segment() {
let n = 40i64;
let postings: Vec<Posting> = (1..=n).map(|r| p(r, &[&[0]])).collect();
let terms = vec![(b"x".to_vec(), postings)];
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=n).map(|r| (r, vec![1])).collect();
let seg = build_segment(&terms, n as u64, &[n as u64], &doc_sizes, 64, 0);
assert!(leaf_count(&seg) > 1, "pgsz 64 must span the doclist");
let want: Vec<i64> = (1..=n).collect();
assert_eq!(lookup_term_rowids(&seg.data, b"x"), Some(want));
assert_eq!(lookup_term_rowids(&seg.data, b"y"), Some(Vec::new()));
}
#[test]
fn lookup_bool_tree_n_operand_set_ops() {
use crate::vtab::{Fts5BoolOp, Fts5BoolTree};
use alloc::boxed::Box;
let terms = vec![
(
b"a".to_vec(),
(1..=5).map(|r| p(r, &[&[0]])).collect::<Vec<_>>(),
),
(
b"b".to_vec(),
vec![2, 4, 6, 8]
.into_iter()
.map(|r| p(r, &[&[1]]))
.collect(),
),
(
b"c".to_vec(),
vec![3, 4, 5, 6]
.into_iter()
.map(|r| p(r, &[&[2]]))
.collect(),
),
];
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=8).map(|r| (r, vec![3])).collect();
let seg = build_segment(&terms, 8, &[8], &doc_sizes, 1000, 0);
let leaf = |t: &[u8]| Fts5BoolTree::Leaf(t.to_vec());
let op = |o, l, r| Fts5BoolTree::Op(o, Box::new(l), Box::new(r));
let t = op(
Fts5BoolOp::And,
op(Fts5BoolOp::And, leaf(b"a"), leaf(b"b")),
leaf(b"c"),
);
assert_eq!(lookup_bool_tree_rowids(&seg.data, &t), Some(vec![4]));
let t = op(
Fts5BoolOp::Or,
op(Fts5BoolOp::Or, leaf(b"a"), leaf(b"b")),
leaf(b"c"),
);
assert_eq!(
lookup_bool_tree_rowids(&seg.data, &t),
Some(vec![1, 2, 3, 4, 5, 6, 8])
);
let t = op(
Fts5BoolOp::Or,
leaf(b"a"),
op(Fts5BoolOp::And, leaf(b"b"), leaf(b"c")),
);
assert_eq!(
lookup_bool_tree_rowids(&seg.data, &t),
Some(vec![1, 2, 3, 4, 5, 6])
);
let t = op(
Fts5BoolOp::Not,
op(Fts5BoolOp::Or, leaf(b"a"), leaf(b"b")),
leaf(b"c"),
);
assert_eq!(lookup_bool_tree_rowids(&seg.data, &t), Some(vec![1, 2, 8]));
let t = op(Fts5BoolOp::And, leaf(b"a"), leaf(b"missing"));
assert_eq!(lookup_bool_tree_rowids(&seg.data, &t), Some(Vec::new()));
assert_eq!(
lookup_bool_tree_rowids(&seg.data, &leaf(b"c")),
Some(vec![3, 4, 5, 6])
);
}
#[test]
fn lookup_bool_tree_empty_index_falls_back() {
use crate::vtab::Fts5BoolTree;
let seg = build_segment(&[], 0, &[0], &[], 1000, 0);
let t = Fts5BoolTree::Leaf(b"x".to_vec());
assert_eq!(lookup_bool_tree_rowids(&seg.data, &t), None);
}
#[test]
fn lookup_rowids_in_column_filters_by_column() {
let terms = vec![(
b"word".to_vec(),
vec![
p(1, &[&[0], &[]]),
p(2, &[&[], &[0]]),
p(3, &[&[0], &[1]]),
p(4, &[&[2], &[]]),
],
)];
let doc_sizes: Vec<(i64, Vec<u64>)> = vec![
(1, vec![1, 0]),
(2, vec![0, 1]),
(3, vec![1, 2]),
(4, vec![3, 0]),
];
let seg = build_segment(&terms, 4, &[3, 2], &doc_sizes, 1000, 0);
assert_eq!(
lookup_term_rowids_in_column(&seg.data, b"word", 0),
Some(vec![1, 3, 4])
);
assert_eq!(
lookup_term_rowids_in_column(&seg.data, b"word", 1),
Some(vec![2, 3])
);
assert_eq!(
lookup_term_rowids(&seg.data, b"word"),
Some(vec![1, 2, 3, 4])
);
assert_eq!(
lookup_term_rowids_in_column(&seg.data, b"missing", 0),
Some(Vec::new())
);
assert_eq!(
lookup_term_rowids_in_column(&seg.data, b"word", 9),
Some(Vec::new())
);
}
#[test]
fn lookup_rowids_in_column_multi_leaf() {
let n = 40i64;
let postings: Vec<Posting> = (1..=n)
.map(|r| {
if r % 2 == 0 {
p(r, &[&[0], &[]])
} else {
p(r, &[&[], &[0]])
}
})
.collect();
let terms = vec![(b"x".to_vec(), postings)];
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=n).map(|r| (r, vec![1, 1])).collect();
let seg = build_segment(&terms, n as u64, &[20, 20], &doc_sizes, 64, 0);
assert!(leaf_count(&seg) > 1, "pgsz 64 must span the doclist");
let even: Vec<i64> = (1..=n).filter(|r| r % 2 == 0).collect();
let odd: Vec<i64> = (1..=n).filter(|r| r % 2 == 1).collect();
assert_eq!(lookup_term_rowids_in_column(&seg.data, b"x", 0), Some(even));
assert_eq!(lookup_term_rowids_in_column(&seg.data, b"x", 1), Some(odd));
}
#[test]
fn lookup_prefix_rowids_unions_matching_terms() {
let terms = vec![
(b"apex".to_vec(), vec![p(3, &[&[1]])]),
(b"apple".to_vec(), vec![p(1, &[&[0]]), p(6, &[&[0]])]),
(b"apply".to_vec(), vec![p(4, &[&[0]]), p(6, &[&[1]])]),
(b"banana".to_vec(), vec![p(2, &[&[0]]), p(5, &[&[0]])]),
];
let doc_sizes: Vec<(i64, Vec<u64>)> = vec![
(1, vec![1]),
(2, vec![1]),
(3, vec![2]),
(4, vec![1]),
(5, vec![1]),
(6, vec![2]),
];
let seg = build_segment(&terms, 6, &[8], &doc_sizes, 1000, 0);
assert_eq!(
lookup_prefix_rowids(&seg.data, b"ap"),
Some(vec![1, 3, 4, 6])
);
assert_eq!(
lookup_prefix_rowids(&seg.data, b"appl"),
Some(vec![1, 4, 6])
);
assert_eq!(lookup_prefix_rowids(&seg.data, b"apple"), Some(vec![1, 6]));
assert_eq!(lookup_prefix_rowids(&seg.data, b"ban"), Some(vec![2, 5]));
assert_eq!(lookup_prefix_rowids(&seg.data, b"zzz"), Some(Vec::new()));
let empty = build_segment(&[], 0, &[0], &[], 1000, 0);
assert_eq!(lookup_prefix_rowids(&empty.data, b"ap"), None);
}
#[test]
fn lookup_prefix_rowids_multi_leaf() {
let n = 40usize;
let terms: Vec<(Vec<u8>, Vec<Posting>)> = (0..n)
.map(|i| {
(
format!("word{i:03}").into_bytes(),
vec![p(i as i64 + 1, &[&[0]])],
)
})
.collect();
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=n as i64).map(|r| (r, vec![1])).collect();
let seg = build_segment(&terms, n as u64, &[n as u64], &doc_sizes, 64, 0);
assert!(leaf_count(&seg) > 1, "pgsz 64 must split into many leaves");
assert_eq!(
lookup_prefix_rowids(&seg.data, b"word"),
Some((1..=n as i64).collect::<Vec<_>>())
);
assert_eq!(
lookup_prefix_rowids(&seg.data, b"word01"),
Some((11..=20).collect::<Vec<_>>())
);
assert_eq!(lookup_prefix_rowids(&seg.data, b"zzz"), Some(Vec::new()));
}
#[test]
fn lookup_prefix_rowids_in_column_filters() {
let terms = vec![
(b"fort".to_vec(), vec![p(3, &[&[0], &[]])]),
(
b"fox".to_vec(),
vec![p(1, &[&[0], &[]]), p(2, &[&[], &[0]]), p(4, &[&[], &[1]])],
),
];
let doc_sizes: Vec<(i64, Vec<u64>)> = vec![
(1, vec![1, 0]),
(2, vec![0, 1]),
(3, vec![1, 0]),
(4, vec![0, 2]),
];
let seg = build_segment(&terms, 4, &[2, 3], &doc_sizes, 1000, 0);
assert_eq!(
lookup_prefix_rowids(&seg.data, b"fo"),
Some(vec![1, 2, 3, 4])
);
assert_eq!(
lookup_prefix_rowids_in_column(&seg.data, b"fo", 0),
Some(vec![1, 3])
);
assert_eq!(
lookup_prefix_rowids_in_column(&seg.data, b"fo", 1),
Some(vec![2, 4])
);
assert_eq!(
lookup_prefix_rowids_in_column(&seg.data, b"fo", 9),
Some(Vec::new())
);
}
#[test]
fn lookup_phrase_rowids_adjacency() {
let terms = vec![
(
b"a".to_vec(),
vec![
p(1, &[&[0]]),
p(2, &[&[0]]),
p(3, &[&[1]]),
p(4, &[&[0, 3]]),
p(5, &[&[2]]),
],
),
(
b"b".to_vec(),
vec![
p(1, &[&[1]]),
p(2, &[&[2]]),
p(3, &[&[0]]),
p(4, &[&[1, 5]]),
p(6, &[&[1]]),
],
),
];
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=6).map(|r| (r, vec![8])).collect::<Vec<_>>();
let seg = build_segment(&terms, 6, &[40], &doc_sizes, 1000, 0);
assert_eq!(
lookup_phrase_rowids_k(&seg.data, &[b"a", b"b"]),
Some(vec![1, 4])
);
assert_eq!(
lookup_phrase_rowids_k(&seg.data, &[b"b", b"a"]),
Some(vec![3])
);
assert_eq!(
lookup_phrase_rowids_k(&seg.data, &[b"a", b"zzz"]),
Some(Vec::new())
);
}
#[test]
fn lookup_phrase_repeated_word() {
let terms = vec![(b"a".to_vec(), vec![p(1, &[&[0, 1]]), p(2, &[&[0, 2]])])];
let seg = build_segment(&terms, 2, &[4], &[(1, vec![2]), (2, vec![3])], 1000, 0);
assert_eq!(
lookup_phrase_rowids_k(&seg.data, &[b"a", b"a"]),
Some(vec![1])
);
}
#[test]
fn lookup_phrase_in_column_requires_same_column() {
let terms = vec![
(
b"a".to_vec(),
vec![
p(1, &[&[0], &[]]),
p(2, &[&[0], &[]]),
p(3, &[&[], &[2]]),
p(4, &[&[0], &[5]]),
],
),
(
b"b".to_vec(),
vec![
p(1, &[&[1], &[]]),
p(2, &[&[], &[1]]),
p(3, &[&[], &[3]]),
p(4, &[&[1], &[6]]),
],
),
];
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=4).map(|r| (r, vec![8, 8])).collect::<Vec<_>>();
let seg = build_segment(&terms, 4, &[40, 40], &doc_sizes, 1000, 0);
assert_eq!(
lookup_phrase_rowids_k(&seg.data, &[b"a", b"b"]),
Some(vec![1, 3, 4])
);
assert_eq!(
lookup_phrase_rowids_in_column_k(&seg.data, &[b"a", b"b"], 0),
Some(vec![1, 4])
);
assert_eq!(
lookup_phrase_rowids_in_column_k(&seg.data, &[b"a", b"b"], 1),
Some(vec![3, 4])
);
}
#[test]
fn lookup_phrase_multi_leaf() {
let n = 40i64;
let a_post: Vec<Posting> = (1..=n).map(|r| p(r, &[&[0]])).collect();
let b_post: Vec<Posting> = (1..=n)
.map(|r| {
if r % 2 == 0 {
p(r, &[&[1]])
} else {
p(r, &[&[3]])
}
})
.collect();
let terms = vec![(b"a".to_vec(), a_post), (b"b".to_vec(), b_post)];
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=n).map(|r| (r, vec![8])).collect();
let seg = build_segment(&terms, n as u64, &[8 * n as u64], &doc_sizes, 64, 0);
assert!(leaf_count(&seg) > 1, "pgsz 64 must span the doclists");
let even: Vec<i64> = (1..=n).filter(|r| r % 2 == 0).collect();
assert_eq!(lookup_phrase_rowids_k(&seg.data, &[b"a", b"b"]), Some(even));
}
#[test]
fn lookup_phrase_empty_index_falls_back() {
let seg = build_segment(&[], 0, &[0], &[], 1000, 0);
assert_eq!(lookup_phrase_rowids_k(&seg.data, &[b"a", b"b"]), None);
assert_eq!(
lookup_phrase_rowids_in_column_k(&seg.data, &[b"a", b"b"], 0),
None
);
}
#[test]
fn lookup_phrase_k_consecutive_run() {
let terms = vec![
(
b"a".to_vec(),
vec![
p(1, &[&[0]]),
p(2, &[&[0]]),
p(3, &[&[0]]),
p(4, &[&[0, 3]]),
p(5, &[&[0]]),
],
),
(
b"b".to_vec(),
vec![
p(1, &[&[1]]),
p(2, &[&[1]]),
p(3, &[&[2]]),
p(4, &[&[1, 4]]),
p(5, &[&[1]]),
],
),
(
b"c".to_vec(),
vec![
p(1, &[&[2]]),
p(2, &[&[3]]),
p(3, &[&[1]]),
p(4, &[&[2, 5]]),
],
),
];
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=5).map(|r| (r, vec![8])).collect();
let seg = build_segment(&terms, 5, &[40], &doc_sizes, 1000, 0);
assert_eq!(
lookup_phrase_rowids_k(&seg.data, &[b"a", b"b", b"c"]),
Some(vec![1, 4])
);
assert_eq!(
lookup_phrase_rowids_k(&seg.data, &[b"a", b"b"]),
Some(vec![1, 2, 4, 5])
);
assert_eq!(
lookup_phrase_rowids_k(&seg.data, &[b"a", b"b", b"zzz"]),
Some(Vec::new())
);
}
#[test]
fn lookup_phrase_k_repeated_word() {
let terms = vec![(
b"a".to_vec(),
vec![p(1, &[&[0, 1, 2]]), p(2, &[&[0, 1]]), p(3, &[&[0, 2, 4]])],
)];
let doc_sizes = [(1, vec![3]), (2, vec![2]), (3, vec![5])];
let seg = build_segment(&terms, 3, &[10], &doc_sizes, 1000, 0);
assert_eq!(
lookup_phrase_rowids_k(&seg.data, &[b"a", b"a"]),
Some(vec![1, 2])
);
assert_eq!(
lookup_phrase_rowids_k(&seg.data, &[b"a", b"a", b"a"]),
Some(vec![1])
);
}
#[test]
fn lookup_phrase_k_column_boundary() {
let terms = vec![
(
b"a".to_vec(),
vec![p(1, &[&[0], &[]]), p(2, &[&[0], &[]]), p(3, &[&[], &[3]])],
),
(
b"b".to_vec(),
vec![p(1, &[&[1], &[]]), p(2, &[&[1], &[]]), p(3, &[&[], &[4]])],
),
(
b"c".to_vec(),
vec![p(1, &[&[2], &[]]), p(2, &[&[], &[0]]), p(3, &[&[], &[5]])],
),
];
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=3).map(|r| (r, vec![8, 8])).collect();
let seg = build_segment(&terms, 3, &[40, 40], &doc_sizes, 1000, 0);
assert_eq!(
lookup_phrase_rowids_k(&seg.data, &[b"a", b"b", b"c"]),
Some(vec![1, 3])
);
assert_eq!(
lookup_phrase_rowids_in_column_k(&seg.data, &[b"a", b"b", b"c"], 0),
Some(vec![1])
);
assert_eq!(
lookup_phrase_rowids_in_column_k(&seg.data, &[b"a", b"b", b"c"], 1),
Some(vec![3])
);
}
#[test]
fn lookup_near_rowids_distance_boundary() {
let terms = vec![
(
b"a".to_vec(),
vec![
p(1, &[&[0]]),
p(2, &[&[0]]),
p(3, &[&[0]]),
p(4, &[&[0]]),
p(5, &[&[0]]),
p(6, &[&[0]]),
],
),
(
b"b".to_vec(),
vec![
p(1, &[&[1]]),
p(2, &[&[2]]),
p(3, &[&[3]]),
p(4, &[&[4]]),
p(6, &[&[1]]),
],
),
];
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=6).map(|r| (r, vec![8])).collect();
let seg = build_segment(&terms, 6, &[40], &doc_sizes, 1000, 0);
assert_eq!(
lookup_near_rowids(&seg.data, b"a", b"b", 0),
Some(vec![1, 6])
);
assert_eq!(
lookup_near_rowids(&seg.data, b"a", b"b", 1),
Some(vec![1, 2, 6])
);
assert_eq!(
lookup_near_rowids(&seg.data, b"a", b"b", 2),
Some(vec![1, 2, 3, 6])
);
assert_eq!(
lookup_near_rowids(&seg.data, b"a", b"b", 3),
Some(vec![1, 2, 3, 4, 6])
);
assert_eq!(
lookup_near_rowids(&seg.data, b"b", b"a", 0),
Some(vec![1, 6])
);
assert_eq!(
lookup_near_rowids(&seg.data, b"a", b"zzz", 10),
Some(Vec::new())
);
}
#[test]
fn lookup_near_rowids_requires_same_column() {
let terms = vec![
(
b"a".to_vec(),
vec![
p(1, &[&[0], &[]]),
p(2, &[&[0], &[]]),
p(3, &[&[], &[5]]),
p(4, &[&[0], &[]]),
],
),
(
b"b".to_vec(),
vec![
p(1, &[&[2], &[]]),
p(2, &[&[], &[0]]),
p(3, &[&[], &[7]]),
p(4, &[&[], &[9]]),
],
),
];
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=4).map(|r| (r, vec![8, 12])).collect();
let seg = build_segment(&terms, 4, &[40, 60], &doc_sizes, 1000, 0);
assert_eq!(
lookup_near_rowids(&seg.data, b"a", b"b", 1),
Some(vec![1, 3])
);
}
#[test]
fn lookup_near_rowids_multi_leaf() {
let n = 40i64;
let a_post: Vec<Posting> = (1..=n).map(|r| p(r, &[&[0]])).collect();
let b_post: Vec<Posting> = (1..=n)
.map(|r| {
if r % 2 == 0 {
p(r, &[&[1]])
} else {
p(r, &[&[5]])
}
})
.collect();
let terms = vec![(b"a".to_vec(), a_post), (b"b".to_vec(), b_post)];
let doc_sizes: Vec<(i64, Vec<u64>)> = (1..=n).map(|r| (r, vec![8])).collect();
let seg = build_segment(&terms, n as u64, &[8 * n as u64], &doc_sizes, 64, 0);
assert!(leaf_count(&seg) > 1, "pgsz 64 must span the doclists");
let even: Vec<i64> = (1..=n).filter(|r| r % 2 == 0).collect();
assert_eq!(lookup_near_rowids(&seg.data, b"a", b"b", 0), Some(even));
assert_eq!(
lookup_near_rowids(&seg.data, b"a", b"b", 4),
Some((1..=n).collect::<Vec<_>>())
);
}
#[test]
fn lookup_near_empty_index_falls_back() {
let seg = build_segment(&[], 0, &[0], &[], 1000, 0);
assert_eq!(lookup_near_rowids(&seg.data, b"a", b"b", 10), None);
}
type SegSpec = (i64, Vec<(Vec<u8>, Vec<Posting>)>);
fn multiseg_data(specs: &[SegSpec]) -> Vec<(i64, Vec<u8>)> {
let mut data: Vec<(i64, Vec<u8>)> = Vec::new();
let mut struct_body: Vec<u8> = 0u32.to_be_bytes().to_vec(); put_varint(&mut struct_body, 1); put_varint(&mut struct_body, specs.len() as u64); put_varint(&mut struct_body, 0); put_varint(&mut struct_body, 0); put_varint(&mut struct_body, specs.len() as u64); for (segid, terms) in specs {
let n_docs = terms.iter().flat_map(|(_, ps)| ps.iter()).count() as u64;
let seg = build_segment(terms, n_docs.max(1), &[64], &[], 4096, 0);
let mut pgno = 1i64;
let mut n_leaves = 0i64;
loop {
let rid = segment_leaf_rowid(1, pgno);
match seg.data.iter().find(|(id, _)| *id == rid) {
Some((_, blob)) => {
data.push((segment_leaf_rowid(*segid, pgno), blob.clone()));
n_leaves += 1;
pgno += 1;
}
None => break,
}
}
put_varint(&mut struct_body, *segid as u64);
put_varint(&mut struct_body, 1);
put_varint(&mut struct_body, n_leaves as u64);
}
data.push((STRUCTURE_ROWID, struct_body));
data
}
#[test]
fn multiseg_bare_term_unions_across_segments_and_routes() {
let specs = vec![
(
1i64,
vec![(b"cat".to_vec(), vec![p(1, &[&[0]]), p(4, &[&[0]])])],
),
(
2i64,
vec![(b"cat".to_vec(), vec![p(2, &[&[0]]), p(7, &[&[0]])])],
),
(
3i64,
vec![
(b"cat".to_vec(), vec![p(5, &[&[0]])]),
(b"dog".to_vec(), vec![p(9, &[&[0]])]),
],
),
];
let data = multiseg_data(&specs);
let before = INDEX_ROUTE_HITS.load(core::sync::atomic::Ordering::Relaxed);
assert_eq!(lookup_term_rowids(&data, b"cat"), Some(vec![1, 2, 4, 5, 7]));
assert!(
INDEX_ROUTE_HITS.load(core::sync::atomic::Ordering::Relaxed) > before,
"the multi-segment bare term must take the index route"
);
assert_eq!(lookup_term_rowids(&data, b"dog"), Some(vec![9]));
assert_eq!(lookup_term_rowids(&data, b"zzz"), Some(Vec::new()));
}
#[test]
fn multiseg_overlapping_docid_bails_to_scan() {
let specs = vec![
(1i64, vec![(b"cat".to_vec(), vec![p(4, &[&[0]])])]),
(2i64, vec![(b"cat".to_vec(), vec![p(4, &[&[1]])])]),
];
let data = multiseg_data(&specs);
assert_eq!(lookup_term_rowids(&data, b"cat"), None);
}
#[test]
fn multiseg_boolean_tree_merges_across_segments() {
use crate::vtab::{Fts5BoolOp, Fts5BoolTree};
use alloc::boxed::Box;
let specs = vec![
(
1i64,
vec![(
b"a".to_vec(),
vec![p(1, &[&[0]]), p(2, &[&[0]]), p(3, &[&[0]])],
)],
),
(
2i64,
vec![(
b"b".to_vec(),
vec![p(2, &[&[0]]), p(3, &[&[0]]), p(4, &[&[0]])],
)],
),
(
3i64,
vec![(b"c".to_vec(), vec![p(3, &[&[0]]), p(5, &[&[0]])])],
),
];
let data = multiseg_data(&specs);
let leaf = |t: &[u8]| Fts5BoolTree::Leaf(t.to_vec());
let op = |o, l, r| Fts5BoolTree::Op(o, Box::new(l), Box::new(r));
let t = op(Fts5BoolOp::And, leaf(b"a"), leaf(b"b"));
assert_eq!(lookup_bool_tree_rowids(&data, &t), Some(vec![2, 3]));
let t = op(Fts5BoolOp::Or, leaf(b"a"), leaf(b"c"));
assert_eq!(lookup_bool_tree_rowids(&data, &t), Some(vec![1, 2, 3, 5]));
let t = op(
Fts5BoolOp::Not,
op(Fts5BoolOp::Or, leaf(b"a"), leaf(b"b")),
leaf(b"c"),
);
assert_eq!(lookup_bool_tree_rowids(&data, &t), Some(vec![1, 2, 4]));
}
#[test]
fn multiseg_boolean_tree_bails_when_a_leaf_overlaps() {
use crate::vtab::{Fts5BoolOp, Fts5BoolTree};
use alloc::boxed::Box;
let specs = vec![
(1i64, vec![(b"a".to_vec(), vec![p(1, &[&[0]])])]),
(2i64, vec![(b"b".to_vec(), vec![p(1, &[&[0]])])]),
(3i64, vec![(b"b".to_vec(), vec![p(1, &[&[1]])])]),
];
let data = multiseg_data(&specs);
let t = Fts5BoolTree::Op(
Fts5BoolOp::And,
Box::new(Fts5BoolTree::Leaf(b"a".to_vec())),
Box::new(Fts5BoolTree::Leaf(b"b".to_vec())),
);
assert_eq!(lookup_bool_tree_rowids(&data, &t), None);
}
#[test]
fn multiseg_prefix_unions_across_segments() {
let specs = vec![
(1i64, vec![(b"apple".to_vec(), vec![p(1, &[&[0]])])]),
(2i64, vec![(b"apply".to_vec(), vec![p(2, &[&[0]])])]),
(
3i64,
vec![
(b"apex".to_vec(), vec![p(3, &[&[0]])]),
(b"banana".to_vec(), vec![p(4, &[&[0]])]),
],
),
];
let data = multiseg_data(&specs);
assert_eq!(lookup_prefix_rowids(&data, b"ap"), Some(vec![1, 2, 3]));
assert_eq!(lookup_prefix_rowids(&data, b"appl"), Some(vec![1, 2]));
assert_eq!(lookup_prefix_rowids(&data, b"ban"), Some(vec![4]));
assert_eq!(lookup_prefix_rowids(&data, b"zzz"), Some(Vec::new()));
}
#[test]
fn multiseg_prefix_same_doc_two_terms_one_segment_is_not_overlap() {
let specs = vec![
(
1i64,
vec![
(b"apple".to_vec(), vec![p(1, &[&[0]])]),
(b"apply".to_vec(), vec![p(1, &[&[1]])]),
],
),
(2i64, vec![(b"apex".to_vec(), vec![p(2, &[&[0]])])]),
];
let data = multiseg_data(&specs);
assert_eq!(lookup_prefix_rowids(&data, b"ap"), Some(vec![1, 2]));
}
#[test]
fn delete_tombstone_in_poslist_bails() {
let mut body: Vec<u8> = Vec::new();
let key = term_key(b"cat");
put_varint(&mut body, key.len() as u64);
body.extend_from_slice(&key);
let term_off = 4; put_varint(&mut body, 1); put_varint(&mut body, 1); let footer_off = 4 + body.len();
let mut leaf: Vec<u8> = Vec::new();
leaf.extend_from_slice(&0u16.to_be_bytes()); leaf.extend_from_slice(&(footer_off as u16).to_be_bytes());
leaf.extend_from_slice(&body);
put_varint(&mut leaf, term_off as u64);
let mut struct_body: Vec<u8> = 0u32.to_be_bytes().to_vec();
for v in [1u64, 1, 0, 0, 1, 1, 1, 1] {
put_varint(&mut struct_body, v);
}
let data = vec![
(segment_leaf_rowid(1, 1), leaf),
(STRUCTURE_ROWID, struct_body),
];
assert_eq!(lookup_term_rowids(&data, b"cat"), None);
}
#[test]
fn multiseg_two_term_phrase_unions_across_segments_and_routes() {
let specs = vec![
(
1i64,
vec![
(b"a".to_vec(), vec![p(1, &[&[0]]), p(2, &[&[0]])]),
(b"b".to_vec(), vec![p(1, &[&[1]]), p(2, &[&[2]])]),
],
),
(
2i64,
vec![
(b"a".to_vec(), vec![p(4, &[&[0]]), p(5, &[&[1]])]),
(b"b".to_vec(), vec![p(4, &[&[1]]), p(5, &[&[0]])]),
],
),
(
3i64,
vec![
(b"a".to_vec(), vec![p(7, &[&[0, 3]]), p(9, &[&[5]])]),
(b"b".to_vec(), vec![p(7, &[&[1]])]),
],
),
];
let data = multiseg_data(&specs);
let before = INDEX_ROUTE_HITS.load(core::sync::atomic::Ordering::Relaxed);
assert_eq!(
lookup_phrase_rowids_k(&data, &[b"a", b"b"]),
Some(vec![1, 4, 7])
);
assert!(
INDEX_ROUTE_HITS.load(core::sync::atomic::Ordering::Relaxed) > before,
"the multi-segment phrase must take the index route"
);
assert_eq!(lookup_phrase_rowids_k(&data, &[b"b", b"a"]), Some(vec![5]));
assert_eq!(
lookup_phrase_rowids_k(&data, &[b"a", b"zzz"]),
Some(Vec::new())
);
}
#[test]
fn multiseg_three_term_phrase_across_segments() {
let specs = vec![
(
1i64,
vec![
(b"a".to_vec(), vec![p(1, &[&[0]]), p(2, &[&[0]])]),
(b"b".to_vec(), vec![p(1, &[&[1]]), p(2, &[&[1]])]),
(b"c".to_vec(), vec![p(1, &[&[2]]), p(2, &[&[3]])]),
],
),
(
2i64,
vec![
(b"a".to_vec(), vec![p(3, &[&[5]]), p(4, &[&[0]])]),
(b"b".to_vec(), vec![p(3, &[&[6]]), p(4, &[&[2]])]),
(b"c".to_vec(), vec![p(3, &[&[7]]), p(4, &[&[3]])]),
],
),
];
let data = multiseg_data(&specs);
assert_eq!(
lookup_phrase_rowids_k(&data, &[b"a", b"b", b"c"]),
Some(vec![1, 3])
);
}
#[test]
fn multiseg_phrase_in_column_across_segments() {
let specs = vec![
(
1i64,
vec![
(b"a".to_vec(), vec![p(1, &[&[0], &[]]), p(2, &[&[0], &[]])]),
(b"b".to_vec(), vec![p(1, &[&[1], &[]]), p(2, &[&[], &[1]])]),
],
),
(
2i64,
vec![
(b"a".to_vec(), vec![p(3, &[&[], &[2]]), p(4, &[&[0], &[5]])]),
(b"b".to_vec(), vec![p(3, &[&[], &[3]]), p(4, &[&[1], &[6]])]),
],
),
];
let data = multiseg_data(&specs);
assert_eq!(
lookup_phrase_rowids_in_column_k(&data, &[b"a", b"b"], 0),
Some(vec![1, 4])
);
assert_eq!(
lookup_phrase_rowids_in_column_k(&data, &[b"a", b"b"], 1),
Some(vec![3, 4])
);
}
#[test]
fn multiseg_near_unions_across_segments_and_routes() {
let specs = vec![
(
1i64,
vec![
(b"a".to_vec(), vec![p(1, &[&[0]]), p(2, &[&[0]])]),
(b"b".to_vec(), vec![p(1, &[&[1]]), p(2, &[&[5]])]),
],
),
(
2i64,
vec![
(b"a".to_vec(), vec![p(3, &[&[4]]), p(4, &[&[0]])]),
(b"b".to_vec(), vec![p(3, &[&[2]]), p(4, &[&[9]])]),
],
),
];
let data = multiseg_data(&specs);
let before = INDEX_ROUTE_HITS.load(core::sync::atomic::Ordering::Relaxed);
assert_eq!(lookup_near_rowids(&data, b"a", b"b", 1), Some(vec![1, 3]));
assert!(
INDEX_ROUTE_HITS.load(core::sync::atomic::Ordering::Relaxed) > before,
"the multi-segment NEAR must take the index route"
);
assert_eq!(
lookup_near_rowids(&data, b"a", b"b", 4),
Some(vec![1, 2, 3])
);
assert_eq!(
lookup_near_rowids(&data, b"a", b"zzz", 10),
Some(Vec::new())
);
}
#[test]
fn multiseg_phrase_overlapping_docid_bails_to_scan() {
let specs = vec![
(
1i64,
vec![
(b"a".to_vec(), vec![p(4, &[&[0]])]),
(b"b".to_vec(), vec![p(4, &[&[1]])]),
],
),
(
2i64,
vec![
(b"a".to_vec(), vec![p(4, &[&[0]])]),
(b"b".to_vec(), vec![p(4, &[&[1]])]),
],
),
];
let data = multiseg_data(&specs);
assert_eq!(lookup_phrase_rowids_k(&data, &[b"a", b"b"]), None);
assert_eq!(lookup_near_rowids(&data, b"a", b"b", 5), None);
}
#[test]
fn multiseg_phrase_cross_term_overlap_bails_to_scan() {
let specs = vec![
(1i64, vec![(b"a".to_vec(), vec![p(4, &[&[0]])])]),
(2i64, vec![(b"b".to_vec(), vec![p(4, &[&[1]])])]),
];
let data = multiseg_data(&specs);
assert_eq!(lookup_phrase_rowids_k(&data, &[b"a", b"b"]), None);
assert_eq!(lookup_near_rowids(&data, b"a", b"b", 5), None);
}
#[test]
fn multiseg_phrase_tombstone_bails_to_scan() {
let mut tomb_leaf: Vec<u8> = Vec::new();
let mut body: Vec<u8> = Vec::new();
let key = term_key(b"a");
put_varint(&mut body, key.len() as u64);
body.extend_from_slice(&key);
let term_off = 4;
put_varint(&mut body, 1); put_varint(&mut body, 1); let footer_off = 4 + body.len();
tomb_leaf.extend_from_slice(&0u16.to_be_bytes());
tomb_leaf.extend_from_slice(&(footer_off as u16).to_be_bytes());
tomb_leaf.extend_from_slice(&body);
put_varint(&mut tomb_leaf, term_off as u64);
let seg1 = build_segment(
&[
(b"a".to_vec(), vec![p(3, &[&[0]])]),
(b"b".to_vec(), vec![p(3, &[&[1]])]),
],
1,
&[64],
&[],
4096,
0,
);
let leaf1 = seg1
.data
.iter()
.find(|(id, _)| *id == segment_leaf_rowid(1, 1))
.unwrap()
.1
.clone();
let mut struct_body: Vec<u8> = 0u32.to_be_bytes().to_vec();
for v in [1u64, 2, 0, 0, 2, 1, 1, 1, 2, 1, 1] {
put_varint(&mut struct_body, v);
}
let data = vec![
(segment_leaf_rowid(1, 1), leaf1),
(segment_leaf_rowid(2, 1), tomb_leaf),
(STRUCTURE_ROWID, struct_body),
];
assert_eq!(lookup_phrase_rowids_k(&data, &[b"a", b"b"]), None);
assert_eq!(lookup_near_rowids(&data, b"a", b"b", 5), None);
}
fn tomb(rowid: i64) -> Posting {
Posting {
rowid,
cols: vec![Vec::new()],
del: true,
}
}
fn seg_leaves_of(terms: &[(Vec<u8>, Vec<Posting>)], segid: i64) -> Vec<Vec<u8>> {
let block = build_segment_block(terms, &[], 4050, segid, &[]);
block
.data
.into_iter()
.filter(|(id, _)| (*id & (1 << 36)) == 0) .map(|(_, b)| b)
.collect()
}
#[test]
fn read_segment_postings_roundtrips_with_tombstone() {
let terms = vec![
(b"apple".to_vec(), vec![p(1, &[&[0, 3]]), tomb(4)]),
(b"banana".to_vec(), vec![p(2, &[&[1]]), p(5, &[&[0]])]),
];
let leaves = seg_leaves_of(&terms, 7);
let refs: Vec<&[u8]> = leaves.iter().map(|l| l.as_slice()).collect();
let got = read_segment_postings(&refs).expect("servable");
assert_eq!(got, terms);
}
#[test]
fn merge_precedence_and_annihilation() {
let old = vec![(b"x".to_vec(), vec![p(1, &[&[0]]), p(2, &[&[0]])])];
let new = vec![(b"x".to_vec(), vec![tomb(1), p(3, &[&[0]])])];
let old_l = seg_leaves_of(&old, 1);
let new_l = seg_leaves_of(&new, 2);
let segs: Vec<Vec<&[u8]>> = vec![
old_l.iter().map(|l| l.as_slice()).collect(),
new_l.iter().map(|l| l.as_slice()).collect(),
];
let merged = merge_segments_keepdel(&segs, false).expect("servable");
assert_eq!(
merged,
vec![(b"x".to_vec(), vec![tomb(1), p(2, &[&[0]]), p(3, &[&[0]])])]
);
let merged_oldest = merge_segments_keepdel(&segs, true).expect("servable");
assert_eq!(
merged_oldest,
vec![(b"x".to_vec(), vec![p(2, &[&[0]]), p(3, &[&[0]])])]
);
}
}