#![cfg(feature = "std")]
#![cfg(feature = "fts5")]
use std::process::Command;
fn put_varint(out: &mut Vec<u8>, v: u64) {
if v >> 56 != 0 {
for i in (0..8).rev() {
out.push((((v >> (i * 7 + 8)) & 0x7f) as u8) | 0x80);
}
out.push((v & 0xff) as u8);
return;
}
let mut groups = [0u8; 9];
let mut n = 0;
let mut w = v;
loop {
groups[n] = (w & 0x7f) as u8;
n += 1;
w >>= 7;
if w == 0 {
break;
}
}
for i in (0..n).rev() {
let byte = groups[i] | if i != 0 { 0x80 } else { 0 };
out.push(byte);
}
}
#[derive(Clone)]
struct Posting {
rowid: i64,
positions: Vec<u32>,
}
fn encode_poslist(p: &Posting) -> Vec<u8> {
let mut collist = Vec::new();
let mut prev = 0u32;
for (i, &pos) in p.positions.iter().enumerate() {
let delta = if i == 0 { pos } else { pos - prev };
put_varint(&mut collist, (delta as u64) + 2);
prev = pos;
}
let mut out = Vec::new();
put_varint(&mut out, (collist.len() as u64) * 2);
out.extend_from_slice(&collist);
out
}
fn encode_doclist(postings: &[Posting]) -> Vec<u8> {
let mut out = Vec::new();
let mut prev = 0i64;
for (i, p) in postings.iter().enumerate() {
if i == 0 {
put_varint(&mut out, p.rowid as u64);
} else {
put_varint(&mut out, (p.rowid - prev) as u64);
}
out.extend_from_slice(&encode_poslist(p));
prev = p.rowid;
}
out
}
fn encode_leaf(terms: &[(String, Vec<Posting>)]) -> Vec<u8> {
let mut body = Vec::new();
let mut offsets = Vec::new();
let mut prev_key: Vec<u8> = Vec::new();
for (term, postings) in terms {
offsets.push(4 + body.len()); let mut key = Vec::with_capacity(term.len() + 1);
key.push(b'0'); key.extend_from_slice(term.as_bytes());
if prev_key.is_empty() {
put_varint(&mut body, key.len() as u64);
body.extend_from_slice(&key);
} else {
let n_common = key
.iter()
.zip(prev_key.iter())
.take_while(|(a, b)| a == b)
.count();
put_varint(&mut body, n_common as u64);
put_varint(&mut body, (key.len() - n_common) as u64);
body.extend_from_slice(&key[n_common..]);
}
body.extend_from_slice(&encode_doclist(postings));
prev_key = key;
}
let footer_off = 4 + body.len();
let mut pgidx = Vec::new();
let mut prev_off = 0usize;
for (i, &off) in offsets.iter().enumerate() {
put_varint(
&mut pgidx,
(if i == 0 { off } else { off - prev_off }) as u64,
);
prev_off = off;
}
let mut leaf = 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);
leaf.extend_from_slice(&pgidx);
leaf
}
fn encode_averages(n_row: u64, total_tokens_col0: u64) -> Vec<u8> {
let mut out = Vec::new();
put_varint(&mut out, n_row);
put_varint(&mut out, total_tokens_col0);
out
}
fn encode_structure(n_leaves: u64, cookie: u32) -> Vec<u8> {
let mut out = cookie.to_be_bytes().to_vec(); for v in [
1, 1, n_leaves, 0, 1, 1, 1, n_leaves,
] {
put_varint(&mut out, v);
}
out
}
struct Leaf {
bytes: Vec<u8>,
first_key: Vec<u8>,
last_key: Vec<u8>,
}
fn pack_leaves(terms: &[(String, Vec<Posting>)], pgsz: usize) -> Vec<Leaf> {
let mut leaves = Vec::new();
let mut cur: Vec<(String, Vec<Posting>)> = Vec::new();
let flush = |cur: &[(String, Vec<Posting>)]| Leaf {
bytes: encode_leaf(cur),
first_key: term_key(&cur[0].0),
last_key: term_key(&cur[cur.len() - 1].0),
};
for (term, postings) in terms {
if !cur.is_empty() {
let mut trial = cur.clone();
trial.push((term.clone(), postings.clone()));
if encode_leaf(&trial).len() >= pgsz {
leaves.push(flush(&cur));
cur.clear();
}
}
cur.push((term.clone(), postings.clone()));
}
if !cur.is_empty() {
leaves.push(flush(&cur));
}
leaves
}
fn term_key(term: &str) -> Vec<u8> {
let mut key = vec![b'0'];
key.extend_from_slice(term.as_bytes());
key
}
fn idx_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()
}
fn build_pgidx(offsets: &[usize]) -> Vec<u8> {
let mut pgidx = Vec::new();
let mut prev = 0usize;
for (i, &off) in offsets.iter().enumerate() {
put_varint(&mut pgidx, (if i == 0 { off } else { off - prev }) as u64);
prev = off;
}
pgidx
}
fn finish_leaf(body: &[u8], term_offsets: &[usize], first_rowid_off: usize) -> Vec<u8> {
let footer_off = 4 + body.len();
let mut leaf = Vec::new();
leaf.extend_from_slice(&(first_rowid_off as u16).to_be_bytes());
leaf.extend_from_slice(&(footer_off as u16).to_be_bytes());
leaf.extend_from_slice(body);
leaf.extend_from_slice(&build_pgidx(term_offsets));
leaf
}
fn encode_single_term_segment(term: &str, postings: &[Posting], pgsz: usize) -> Vec<Vec<u8>> {
let key = term_key(term);
let mut leaves: Vec<Vec<u8>> = Vec::new();
let mut body: Vec<u8> = Vec::new();
let mut term_offsets: Vec<usize> = Vec::new();
let mut first_rowid_off = 0usize; let mut prev_rowid = 0i64; let pgidx_len = |offs: &[usize]| build_pgidx(offs).len();
term_offsets.push(4);
put_varint(&mut body, key.len() as u64);
body.extend_from_slice(&key);
for p in postings {
if 4 + body.len() + pgidx_len(&term_offsets) > pgsz && !term_offsets.is_empty() {
leaves.push(finish_leaf(&body, &term_offsets, first_rowid_off));
body.clear();
term_offsets.clear();
first_rowid_off = 0;
prev_rowid = 0;
}
if term_offsets.is_empty() && first_rowid_off == 0 {
first_rowid_off = 4 + body.len(); }
let poslist = encode_poslist(p);
let size_len = get_varint(&poslist).1;
put_varint(&mut body, (p.rowid - prev_rowid) as u64);
body.extend_from_slice(&poslist[..size_len]);
prev_rowid = p.rowid;
for &b in &poslist[size_len..] {
if 4 + body.len() + pgidx_len(&term_offsets) >= pgsz {
leaves.push(finish_leaf(&body, &term_offsets, first_rowid_off));
body.clear();
term_offsets.clear();
first_rowid_off = 0;
prev_rowid = 0;
}
body.push(b);
}
}
leaves.push(finish_leaf(&body, &term_offsets, first_rowid_off));
leaves
}
fn get_varint(buf: &[u8]) -> (u64, usize) {
let mut v = 0u64;
for (i, &b) in buf.iter().enumerate().take(9) {
if i == 8 {
return (v << 8 | b as u64, 9);
}
v = (v << 7) | (b & 0x7f) as u64;
if b & 0x80 == 0 {
return (v, i + 1);
}
}
(v, buf.len())
}
struct SegWriter {
pgsz: usize,
leaves: Vec<Vec<u8>>,
idx: Vec<(i64, Vec<u8>)>,
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,
}
impl SegWriter {
fn new(pgsz: usize) -> Self {
SegWriter {
pgsz,
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,
}
}
fn leaf_size(&self) -> usize {
4 + self.body.len() + build_pgidx(&self.term_offsets).len()
}
fn flush(&mut self) {
self.leaves.push(finish_leaf(
&self.body,
&self.term_offsets,
self.first_rowid_off,
));
if let Some(ft) = self.leaf_first_term.take() {
let sep = match &self.prev_leaf_last_term {
Some(p) => idx_separator(p, &ft),
None => Vec::new(),
};
self.idx.push((self.pgno << 1, sep));
}
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;
}
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_with(&self) -> usize {
let mut probe = self.term_offsets.clone();
probe.push(4 + self.body.len());
build_pgidx(&probe).len()
}
fn add_term(&mut self, term: &str, postings: &[Posting]) {
let key = term_key(term);
let doclist = encode_doclist(postings); let rec = self.term_record(&key);
if !self.body.is_empty()
&& 4 + self.body.len() + rec.len() + doclist.len() + self.pgidx_with() >= self.pgsz
{
self.flush();
}
let rec = self.term_record(&key); let fits_whole =
4 + self.body.len() + rec.len() + doclist.len() + self.pgidx_with() <= self.pgsz;
self.term_offsets.push(4 + self.body.len());
if self.leaf_first_term.is_none() {
self.leaf_first_term = Some(key.clone());
}
self.leaf_last_term = Some(key.clone());
self.body.extend_from_slice(&rec);
self.prev_term_key = Some(key);
if fits_whole {
self.body.extend_from_slice(&doclist);
return;
}
self.prev_rowid = 0;
for p in postings {
if self.leaf_size() > self.pgsz && !self.body.is_empty() {
self.flush();
}
if self.term_offsets.is_empty() && self.first_rowid_off == 0 {
self.first_rowid_off = 4 + self.body.len();
}
let poslist = encode_poslist(p);
let size_len = get_varint(&poslist).1;
put_varint(&mut self.body, (p.rowid - self.prev_rowid) as u64);
self.body.extend_from_slice(&poslist[..size_len]);
self.prev_rowid = p.rowid;
for &b in &poslist[size_len..] {
if self.leaf_size() >= self.pgsz {
self.flush();
}
self.body.push(b);
}
}
}
#[allow(clippy::type_complexity)]
fn finish(mut self) -> (Vec<Vec<u8>>, Vec<(i64, Vec<u8>)>) {
self.flush();
(self.leaves, self.idx)
}
}
#[allow(clippy::type_complexity)]
fn encode_segment(
terms: &[(String, Vec<Posting>)],
pgsz: usize,
) -> (Vec<Vec<u8>>, Vec<(i64, Vec<u8>)>) {
let mut w = SegWriter::new(pgsz);
for (term, postings) in terms {
w.add_term(term, postings);
}
w.finish()
}
const SEG_LEAF_ROWID: i64 = (1i64 << 37) | 1;
fn parse_hex(s: &str) -> Vec<u8> {
let s = s.trim();
let inner = s
.strip_prefix("X'")
.or_else(|| s.strip_prefix("x'"))
.and_then(|r| r.strip_suffix('\''))
.unwrap_or("");
(0..inner.len())
.step_by(2)
.map(|i| u8::from_str_radix(&inner[i..i + 2], 16).unwrap())
.collect()
}
fn sqlite_data_rows(docs: &[(i64, &str)]) -> (Vec<u8>, Vec<u8>, Vec<u8>) {
use std::sync::atomic::{AtomicU64, Ordering};
static SEQ: AtomicU64 = AtomicU64::new(0);
let path = std::env::temp_dir().join(format!(
"gsql-fts5seg-{}-{}.db",
std::process::id(),
SEQ.fetch_add(1, Ordering::Relaxed)
));
let path = path.to_string_lossy().into_owned();
let _ = std::fs::remove_file(&path);
let values: Vec<String> = docs.iter().map(|(r, b)| format!("({r},'{b}')")).collect();
let script = format!(
"CREATE VIRTUAL TABLE t USING fts5(body);\
INSERT INTO t(rowid,body) VALUES {};",
values.join(",")
);
let o = Command::new("sqlite3")
.arg(&path)
.arg(&script)
.output()
.unwrap();
assert!(o.status.success(), "sqlite build failed: {o:?}");
let fetch = |id: i64| {
let o = Command::new("sqlite3")
.arg(&path)
.arg(format!("SELECT quote(block) FROM t_data WHERE id={id};"))
.output()
.unwrap();
parse_hex(&String::from_utf8_lossy(&o.stdout))
};
let r = (fetch(1), fetch(10), fetch(SEG_LEAF_ROWID));
let _ = std::fs::remove_file(&path);
r
}
fn build_index(docs: &[(i64, &str)]) -> (Vec<(String, Vec<Posting>)>, u64, u64) {
use std::collections::BTreeMap;
let mut idx: BTreeMap<String, BTreeMap<i64, Vec<u32>>> = BTreeMap::new();
let mut total_tokens = 0u64;
for (rowid, body) in docs {
for (pos, tok) in body.split_whitespace().enumerate() {
total_tokens += 1;
idx.entry(tok.to_string())
.or_default()
.entry(*rowid)
.or_default()
.push(pos as u32);
}
}
let terms = idx
.into_iter()
.map(|(term, docs)| {
let postings = docs
.into_iter()
.map(|(rowid, positions)| Posting { rowid, positions })
.collect();
(term, postings)
})
.collect();
(terms, docs.len() as u64, total_tokens)
}
fn check(docs: &[(i64, &str)]) {
if Command::new("sqlite3").arg("--version").output().is_err() {
eprintln!("sqlite3 not found; skipping");
return;
}
let (avg, structure, leaf) = sqlite_data_rows(docs);
let (terms, n_row, total_tokens) = build_index(docs);
assert_eq!(
encode_averages(n_row, total_tokens),
avg,
"averages mismatch"
);
assert_eq!(encode_structure(1, 0), structure, "structure mismatch");
assert_eq!(encode_leaf(&terms), leaf, "leaf mismatch for {docs:?}");
}
#[allow(clippy::type_complexity)]
fn sqlite_segment(pgsz: usize, body: &str) -> (Vec<Vec<u8>>, Vec<u8>, Vec<(i64, Vec<u8>)>) {
use std::sync::atomic::{AtomicU64, Ordering};
static SEQ: AtomicU64 = AtomicU64::new(0);
let path = std::env::temp_dir().join(format!(
"gsql-fts5ml-{}-{}.db",
std::process::id(),
SEQ.fetch_add(1, Ordering::Relaxed)
));
let path = path.to_string_lossy().into_owned();
let _ = std::fs::remove_file(&path);
let script = format!(
"CREATE VIRTUAL TABLE t USING fts5(body);\
INSERT INTO t(t,rank) VALUES('pgsz',{pgsz});\
INSERT INTO t(rowid,body) VALUES (1,'{body}');"
);
let o = Command::new("sqlite3")
.arg(&path)
.arg(&script)
.output()
.unwrap();
assert!(o.status.success(), "sqlite build failed: {o:?}");
let run = |q: &str| {
String::from_utf8_lossy(
&Command::new("sqlite3")
.arg(&path)
.arg(q)
.output()
.unwrap()
.stdout,
)
.into_owned()
};
let leaves: Vec<Vec<u8>> = run("SELECT quote(block) FROM t_data WHERE id>100 ORDER BY id;")
.lines()
.map(parse_hex)
.collect();
let structure = parse_hex(&run("SELECT quote(block) FROM t_data WHERE id=10;"));
let idx: Vec<(i64, Vec<u8>)> =
run("SELECT pgno || '|' || quote(term) FROM t_idx ORDER BY pgno;")
.lines()
.map(|l| {
let (p, t) = l.split_once('|').unwrap();
(p.parse().unwrap(), parse_hex(t))
})
.collect();
let _ = std::fs::remove_file(&path);
(leaves, structure, idx)
}
#[allow(clippy::type_complexity)]
fn sqlite_segment_rows(
pgsz: usize,
docs: &[(i64, &str)],
) -> (Vec<Vec<u8>>, Vec<u8>, Vec<(i64, Vec<u8>)>) {
use std::sync::atomic::{AtomicU64, Ordering};
static SEQ: AtomicU64 = AtomicU64::new(0);
let path = std::env::temp_dir().join(format!(
"gsql-fts5dl-{}-{}.db",
std::process::id(),
SEQ.fetch_add(1, Ordering::Relaxed)
));
let path = path.to_string_lossy().into_owned();
let _ = std::fs::remove_file(&path);
let values: Vec<String> = docs.iter().map(|(r, b)| format!("({r},'{b}')")).collect();
let script = format!(
"CREATE VIRTUAL TABLE t USING fts5(body);\
INSERT INTO t(t,rank) VALUES('pgsz',{pgsz});\
INSERT INTO t(rowid,body) VALUES {};",
values.join(",")
);
let o = Command::new("sqlite3")
.arg(&path)
.arg(&script)
.output()
.unwrap();
assert!(o.status.success(), "sqlite build failed: {o:?}");
let run = |q: &str| {
String::from_utf8_lossy(
&Command::new("sqlite3")
.arg(&path)
.arg(q)
.output()
.unwrap()
.stdout,
)
.into_owned()
};
let leaves: Vec<Vec<u8>> = run("SELECT quote(block) FROM t_data WHERE id>100 ORDER BY id;")
.lines()
.map(parse_hex)
.collect();
let structure = parse_hex(&run("SELECT quote(block) FROM t_data WHERE id=10;"));
let idx: Vec<(i64, Vec<u8>)> =
run("SELECT pgno || '|' || quote(term) FROM t_idx ORDER BY pgno;")
.lines()
.map(|l| {
let (p, t) = l.split_once('|').unwrap();
(p.parse().unwrap(), parse_hex(t))
})
.collect();
let _ = std::fs::remove_file(&path);
(leaves, structure, idx)
}
fn check_segment(pgsz: usize, body: &str) {
if Command::new("sqlite3").arg("--version").output().is_err() {
eprintln!("sqlite3 not found; skipping");
return;
}
let (sq_leaves, sq_structure, sq_idx) = sqlite_segment(pgsz, body);
let (terms, _, _) = build_index(&[(1, body)]);
let leaves = pack_leaves(&terms, pgsz);
assert!(leaves.len() > 1, "test did not produce multiple leaves");
assert_eq!(leaves.len(), sq_leaves.len(), "leaf count");
for (i, (leaf, sq)) in leaves.iter().zip(&sq_leaves).enumerate() {
assert_eq!(&leaf.bytes, sq, "leaf {} (pgno {}) bytes", i, i + 1);
}
assert_eq!(
encode_structure(leaves.len() as u64, 1),
sq_structure,
"structure"
);
let want_idx: Vec<(i64, Vec<u8>)> = leaves
.iter()
.enumerate()
.map(|(i, leaf)| {
let pgno = (i as i64) + 1;
let term = if i == 0 {
Vec::new()
} else {
idx_separator(&leaves[i - 1].last_key, &leaf.first_key)
};
(pgno << 1, term)
})
.collect();
assert_eq!(want_idx, sq_idx, "%_idx rows");
}
#[test]
fn single_term_single_doc() {
check(&[(1, "a")]);
check(&[(5, "cat")]);
}
#[test]
fn two_terms_prefix_compressed() {
check(&[(1, "a b")]);
check(&[(1, "apple apply")]); }
#[test]
fn multi_doc_rowid_deltas() {
check(&[(1, "hello"), (2, "hello")]);
check(&[(1, "x"), (3, "x")]); }
#[test]
fn multi_position_collist() {
check(&[(1, "the the")]); }
#[test]
fn mixed_three_terms() {
check(&[(1, "red green"), (2, "green blue")]);
}
#[test]
fn multi_leaf_pagination() {
let body: String = (1..=40)
.map(|i| format!("term{i:03}"))
.collect::<Vec<_>>()
.join(" ");
check_segment(64, &body);
}
#[test]
fn multi_leaf_varied_pgsz() {
let body: String = (1..=50)
.map(|i| format!("term{i:03}"))
.collect::<Vec<_>>()
.join(" ");
check_segment(80, &body);
check_segment(128, &body);
}
fn check_doclist_spans(n_docs: i64, pgsz: usize) {
if Command::new("sqlite3").arg("--version").output().is_err() {
eprintln!("sqlite3 not found; skipping");
return;
}
let body = vec!["x"; n_docs as usize].join(" ");
let docs: Vec<(i64, &str)> = (1..=n_docs).map(|r| (r, "x")).collect();
let (sq_leaves, sq_structure, sq_idx) = sqlite_segment_rows(pgsz, &docs);
let (terms, _, _) = build_index(&docs);
assert_eq!(terms.len(), 1, "expected a single term");
let leaves = encode_single_term_segment(&terms[0].0, &terms[0].1, pgsz);
assert!(
leaves.len() > 1,
"doclist did not span leaves (body={body:?})"
);
assert_eq!(leaves.len(), sq_leaves.len(), "leaf count");
for (i, (mine, sq)) in leaves.iter().zip(&sq_leaves).enumerate() {
assert_eq!(mine, sq, "leaf {} (pgno {})", i, i + 1);
}
assert_eq!(
encode_structure(leaves.len() as u64, 1),
sq_structure,
"structure"
);
assert_eq!(sq_idx, vec![(2i64, Vec::new())], "%_idx rows");
}
#[test]
fn doclist_spans_leaves() {
check_doclist_spans(30, 64);
check_doclist_spans(50, 80);
check_doclist_spans(100, 128);
}
fn check_unified(pgsz: usize, docs: &[(i64, &str)]) {
if Command::new("sqlite3").arg("--version").output().is_err() {
eprintln!("sqlite3 not found; skipping");
return;
}
let (sq_leaves, sq_structure, sq_idx) = sqlite_segment_rows(pgsz, docs);
let (terms, _, _) = build_index(docs);
let (leaves, idx) = encode_segment(&terms, pgsz);
assert_eq!(leaves.len(), sq_leaves.len(), "leaf count for {docs:?}");
for (i, (mine, sq)) in leaves.iter().zip(&sq_leaves).enumerate() {
assert_eq!(mine, sq, "leaf {} (pgno {})", i, i + 1);
}
assert_eq!(
encode_structure(leaves.len() as u64, 1),
sq_structure,
"structure"
);
assert_eq!(idx, sq_idx, "%_idx rows");
}
#[test]
fn unified_reproduces_term_pagination() {
let docs: Vec<(i64, String)> = vec![(
1,
(1..=40)
.map(|i| format!("term{i:03}"))
.collect::<Vec<_>>()
.join(" "),
)];
let refs: Vec<(i64, &str)> = docs.iter().map(|(r, b)| (*r, b.as_str())).collect();
check_unified(64, &refs);
}
#[test]
fn unified_varying_rowids() {
let docs: Vec<(i64, String)> = (1..=25).map(|i| (i, format!("word{i:03}"))).collect();
let refs: Vec<(i64, &str)> = docs.iter().map(|(r, b)| (*r, b.as_str())).collect();
check_unified(80, &refs);
check_unified(128, &refs);
}
#[test]
fn unified_reproduces_doclist_spanning() {
let docs: Vec<(i64, &str)> = (1..=40).map(|r| (r, "x")).collect();
check_unified(64, &docs);
}
struct PostingMc {
rowid: i64,
cols: Vec<Vec<u32>>,
}
fn encode_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 encode_poslist_mc(p: &PostingMc) -> 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(&encode_collist(positions));
}
let mut out = Vec::new();
put_varint(&mut out, (content.len() as u64) * 2);
out.extend_from_slice(&content);
out
}
fn encode_doclist_mc(postings: &[PostingMc]) -> Vec<u8> {
let mut out = Vec::new();
let mut prev = 0i64;
for (i, p) in postings.iter().enumerate() {
put_varint(
&mut out,
(if i == 0 { p.rowid } else { p.rowid - prev }) as u64,
);
out.extend_from_slice(&encode_poslist_mc(p));
prev = p.rowid;
}
out
}
fn encode_leaf_mc(terms: &[(String, Vec<PostingMc>)]) -> Vec<u8> {
let mut body = Vec::new();
let mut offsets = Vec::new();
let mut prev_key: Vec<u8> = Vec::new();
for (term, postings) in terms {
offsets.push(4 + body.len());
let key = term_key(term);
if prev_key.is_empty() {
put_varint(&mut body, key.len() as u64);
body.extend_from_slice(&key);
} else {
let n_common = key
.iter()
.zip(prev_key.iter())
.take_while(|(a, b)| a == b)
.count();
put_varint(&mut body, n_common as u64);
put_varint(&mut body, (key.len() - n_common) as u64);
body.extend_from_slice(&key[n_common..]);
}
body.extend_from_slice(&encode_doclist_mc(postings));
prev_key = key;
}
finish_leaf(&body, &offsets, 0)
}
#[allow(clippy::type_complexity)]
fn build_index_mc(
docs: &[(i64, Vec<&str>)],
ncols: usize,
) -> (Vec<(String, Vec<PostingMc>)>, Vec<u64>) {
use std::collections::BTreeMap;
let mut idx: BTreeMap<String, BTreeMap<i64, Vec<Vec<u32>>>> = BTreeMap::new();
let mut totals = vec![0u64; ncols];
for (rowid, cols) in docs {
for (c, text) in cols.iter().enumerate() {
for (pos, tok) in text.split_whitespace().enumerate() {
totals[c] += 1;
idx.entry(tok.to_string())
.or_default()
.entry(*rowid)
.or_insert_with(|| vec![Vec::new(); ncols])[c]
.push(pos as u32);
}
}
}
let terms = idx
.into_iter()
.map(|(term, docs)| {
let postings = docs
.into_iter()
.map(|(rowid, cols)| PostingMc { rowid, cols })
.collect();
(term, postings)
})
.collect();
(terms, totals)
}
fn check_mc(docs: &[(i64, Vec<&str>)], ncols: usize) {
if Command::new("sqlite3").arg("--version").output().is_err() {
eprintln!("sqlite3 not found; skipping");
return;
}
use std::sync::atomic::{AtomicU64, Ordering};
static SEQ: AtomicU64 = AtomicU64::new(0);
let path = std::env::temp_dir().join(format!(
"gsql-fts5mc-{}-{}.db",
std::process::id(),
SEQ.fetch_add(1, Ordering::Relaxed)
));
let path = path.to_string_lossy().into_owned();
let _ = std::fs::remove_file(&path);
let coldefs: Vec<String> = (0..ncols).map(|c| format!("c{c}")).collect();
let collist = coldefs.join(",");
let values: Vec<String> = docs
.iter()
.map(|(r, cols)| {
let vs: Vec<String> = cols.iter().map(|t| format!("'{t}'")).collect();
format!("({r},{})", vs.join(","))
})
.collect();
let script = format!(
"CREATE VIRTUAL TABLE t USING fts5({collist});\
INSERT INTO t(rowid,{collist}) VALUES {};",
values.join(",")
);
let o = Command::new("sqlite3")
.arg(&path)
.arg(&script)
.output()
.unwrap();
assert!(o.status.success(), "sqlite build failed: {o:?}");
let q = |sql: &str| {
parse_hex(&String::from_utf8_lossy(
&Command::new("sqlite3")
.arg(&path)
.arg(sql)
.output()
.unwrap()
.stdout,
))
};
let sq_leaf = q(&format!(
"SELECT quote(block) FROM t_data WHERE id={SEG_LEAF_ROWID};"
));
let sq_avg = q("SELECT quote(block) FROM t_data WHERE id=1;");
let (terms, totals) = build_index_mc(docs, ncols);
assert_eq!(
encode_leaf_mc(&terms),
sq_leaf,
"multi-col leaf for {docs:?}"
);
let mut avg = Vec::new();
put_varint(&mut avg, docs.len() as u64);
for t in &totals {
put_varint(&mut avg, *t);
}
assert_eq!(avg, sq_avg, "averages");
let _ = std::fs::remove_file(&path);
}
#[test]
fn multi_column_poslists() {
check_mc(&[(1, vec!["hello world", "hello there"])], 2);
check_mc(&[(1, vec!["a b", "b c", "c a"])], 3);
check_mc(&[(1, vec!["x y", "y"]), (2, vec!["x", "x y"])], 2);
}