use std::{
collections::{BinaryHeap, HashMap},
ops::Range,
sync::Arc,
};
use bytes::Bytes;
use super::{
cursor::{TermCursor, TermMeta},
filter::ExcludeFilter,
metadata::{ColumnMeta, FtsColumnConfig, NormTable, OpenOptions},
phrase::{AnyCursor, PhraseCursor},
sink::{TopKEntry, drain_top_k_desc},
work::{term_cursor_bytes, term_cursor_ranges},
};
use crate::superfile::{
ReadError,
error::FtsError,
format::{
self, FST_SEPARATOR,
checksum::crc32c,
fts::{
HEADER_SIZE_V1_LEGACY as FTS_HEADER_SIZE, MAGIC_BYTES, U32_BYTES, U64_BYTES, hdr,
term_meta,
},
},
fts::{
builder::{DOC_LENGTHS_ENTRY_SIZE, TERM_META_SIZE},
dict::{DictReader, make_key},
fst_value::FstValue,
positions::decode_run,
posting::{self, BLOCK_LEN, ENCODING_BITSET, decode_block_doc_ids},
tokenize::{Tokenizer, tokenizer_for_name},
},
lazy_source::{LazyByteSource, PrefetchedSource, RangeCoalescePlan, Source},
};
const TERM_RANGE_COALESCE_MAX_GAP: usize = 64 * 1024;
const TERM_RANGE_COALESCE_MAX_OVERFETCH: usize = 512 * 1024;
pub(crate) type GlobalTermIdf = std::collections::HashMap<String, f32>;
#[derive(Default)]
pub(crate) struct ClauseLists<'a> {
pub musts: &'a [&'a str],
pub shoulds: &'a [&'a str],
pub negatives: &'a [&'a str],
pub must_phrases: &'a [Vec<String>],
pub should_phrases: &'a [Vec<String>],
pub negative_phrases: &'a [Vec<String>],
pub global_idf: Option<&'a GlobalTermIdf>,
}
impl ClauseLists<'_> {
pub(super) fn has_phrases(&self) -> bool {
!self.must_phrases.is_empty()
|| !self.should_phrases.is_empty()
|| !self.negative_phrases.is_empty()
}
pub(super) fn no_positive_atoms(&self) -> bool {
self.musts.is_empty()
&& self.shoulds.is_empty()
&& self.must_phrases.is_empty()
&& self.should_phrases.is_empty()
}
pub(super) fn no_negative_atoms(&self) -> bool {
self.negatives.is_empty() && self.negative_phrases.is_empty()
}
}
pub(crate) enum PreparedClauses {
Done {
hits: Vec<(u32, f32)>,
postings_bytes: u64,
planned_ranges: u64,
kernel_cpu_ns: u64,
},
Must {
column_id: u32,
must_cursors: Vec<TermCursor>,
filter: Option<ExcludeFilter>,
k: usize,
floor_eff: f32,
dict_ranges: u64,
},
MustShould {
column_id: u32,
must_cursors: Vec<TermCursor>,
should_cursors: Vec<TermCursor>,
filter: Option<ExcludeFilter>,
k: usize,
floor_eff: f32,
dict_ranges: u64,
},
Or {
column_id: u32,
cursors: Vec<TermCursor>,
filter: Option<ExcludeFilter>,
k: usize,
floor_eff: f32,
dict_ranges: u64,
},
}
impl PreparedClauses {
pub(crate) fn posting_mass(&self) -> u64 {
match self {
PreparedClauses::Done { .. } => 0,
PreparedClauses::Must { must_cursors, .. } => {
must_cursors.iter().map(|c| c.df).min().unwrap_or(0)
}
PreparedClauses::MustShould { must_cursors, .. } => {
must_cursors.iter().map(|c| c.df).min().unwrap_or(0)
}
PreparedClauses::Or { cursors, .. } => cursors.iter().map(|c| c.df).sum(),
}
}
pub(crate) fn postings_bytes(&self) -> u64 {
let filter_bytes =
|filter: &Option<ExcludeFilter>| filter.as_ref().map_or(0, |f| f.postings_bytes());
match self {
PreparedClauses::Done { postings_bytes, .. } => *postings_bytes,
PreparedClauses::Must {
must_cursors,
filter,
..
} => term_cursor_bytes(must_cursors) + filter_bytes(filter),
PreparedClauses::MustShould {
must_cursors,
should_cursors,
filter,
..
} => {
term_cursor_bytes(must_cursors)
+ term_cursor_bytes(should_cursors)
+ filter_bytes(filter)
}
PreparedClauses::Or {
cursors, filter, ..
} => term_cursor_bytes(cursors) + filter_bytes(filter),
}
}
pub(crate) fn inline_kernel_cpu_ns(&self) -> u64 {
match self {
PreparedClauses::Done { kernel_cpu_ns, .. } => *kernel_cpu_ns,
_ => 0,
}
}
pub(crate) fn planned_ranges(&self) -> u64 {
let filter_ranges = |filter: &Option<ExcludeFilter>| {
filter.as_ref().map_or(0, ExcludeFilter::planned_ranges)
};
match self {
PreparedClauses::Done { planned_ranges, .. } => *planned_ranges,
PreparedClauses::Must {
must_cursors,
filter,
dict_ranges,
..
} => term_cursor_ranges(must_cursors) + filter_ranges(filter) + dict_ranges,
PreparedClauses::MustShould {
must_cursors,
should_cursors,
filter,
dict_ranges,
..
} => {
term_cursor_ranges(must_cursors)
+ term_cursor_ranges(should_cursors)
+ filter_ranges(filter)
+ dict_ranges
}
PreparedClauses::Or {
cursors,
filter,
dict_ranges,
..
} => term_cursor_ranges(cursors) + filter_ranges(filter) + dict_ranges,
}
}
}
#[doc(hidden)]
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum OrAlgo {
Bmm,
WandBmw,
Exhaustive,
Windowed,
}
pub(super) const OR_WINDOW: u32 = 4096;
pub(super) const OR_WINDOW_WORDS: usize = (OR_WINDOW as usize).div_ceil(64);
pub(super) const OR_COUNT_ANCHOR_DOMINANCE: u64 = 8;
pub(super) const OR_COUNT_BITSET_DENSITY_DIVISOR: u64 = 16;
pub(crate) const OR_WINDOW_MIN_TERMS: usize = 3;
pub(super) const WAND_BMW_2TERM_MAX_K: usize = 128;
pub(super) const WAND_BMW_2TERM_DF_RATIO: u64 = 16;
pub(super) const OR_WINDOW_DOMINANCE_MULT: f32 = 1.5;
pub(super) const OR_WINDOWED_UNIFORM_MAX_PRUNING_K: usize = 32;
pub(super) fn no_dominant_term_ub(cursors: &[TermCursor]) -> bool {
let total: f32 = cursors.iter().map(|c| c.term_max_bm25).sum();
if total <= 0.0 {
return false;
}
let max = cursors
.iter()
.map(|c| c.term_max_bm25)
.fold(0.0f32, f32::max);
let avg = total / cursors.len() as f32;
max <= OR_WINDOW_DOMINANCE_MULT * avg
}
pub(super) fn route_or_to_windowed(cursors: &[TermCursor], k: usize) -> bool {
or_topk_pruning_ineffective(cursors, k)
|| (k > OR_WINDOWED_UNIFORM_MAX_PRUNING_K
&& cursors.len() >= OR_WINDOW_MIN_TERMS
&& no_dominant_term_ub(cursors))
}
pub(super) const OR_WINDOWED_MIN_DOMINANT_DF: u64 = 100_000;
pub(super) fn or_topk_pruning_ineffective(cursors: &[TermCursor], k: usize) -> bool {
let max_df = cursors.iter().map(|c| c.df).max().unwrap_or(0);
let total_df: u64 = cursors.iter().map(|c| c.df).sum();
or_reroute_by_df(max_df, total_df, cursors.len(), k)
}
pub(super) fn or_reroute_by_df(max_df: u64, total_df: u64, n_terms: usize, k: usize) -> bool {
if n_terms < 2 || max_df < OR_WINDOWED_MIN_DOMINANT_DF {
return false;
}
let rest_df = total_df.saturating_sub(max_df);
k as u64 >= rest_df
}
pub(crate) fn top_k_initial_capacity(k: usize, n_docs: u64, range: Option<(u32, u32)>) -> usize {
let docs_in_scope = match range {
Some((start, end)) => (end.saturating_sub(start) as usize).min(n_docs as usize),
None => n_docs as usize,
};
k.min(docs_in_scope).max(1)
}
pub(super) fn two_term_has_rare_anchor(cursors: &[TermCursor]) -> bool {
if cursors.len() != 2 {
return false;
}
let lo = cursors[0].df.min(cursors[1].df);
let hi = cursors[0].df.max(cursors[1].df);
lo > 0 && hi >= lo.saturating_mul(WAND_BMW_2TERM_DF_RATIO)
}
#[derive(Debug)]
pub struct FtsReader {
pub(super) source: Source,
pub(super) n_docs: u32,
pub(super) n_terms_total: u32,
pub(super) fst_range: Range<usize>,
pub(super) postings_range: Range<usize>,
pub(super) positions_range: Option<Range<usize>>,
pub(super) has_position_subindex: bool,
pub(super) has_bitset_blocks: bool,
pub(super) columns: Vec<ColumnMeta>,
pub(super) column_id_by_name: HashMap<String, u32>,
}
impl FtsReader {
pub fn open(blob: Bytes, columns_json: &str) -> Result<Self, FtsError> {
Self::open_with(blob, columns_json, OpenOptions::default())
}
pub fn open_with(blob: Bytes, columns_json: &str, opts: OpenOptions) -> Result<Self, FtsError> {
Self::open_with_source(Source::InMemory(blob), columns_json, opts)
}
pub async fn open_lazy(
source: Arc<dyn LazyByteSource>,
columns_json: &str,
opts: OpenOptions,
) -> Result<Self, FtsError> {
let fts_blob_len = source.size() as usize;
let header_fetch = format::fts::HEADER_SIZE_V2.min(fts_blob_len);
let header = fetch_lazy_range(source.as_ref(), 0..header_fetch, "fts header").await?;
if header.len() < FTS_HEADER_SIZE {
return Err(FtsError::Read(ReadError::MissingKv("fts header")));
}
if &header[0..MAGIC_BYTES] != format::fts::MAGIC {
return Err(FtsError::Read(ReadError::BadMagic {
section: "fts",
expected: format::fts::MAGIC,
actual: header[0..MAGIC_BYTES].to_vec(),
}));
}
let version = read_u32_le(&header[hdr::VERSION_OFF..hdr::VERSION_OFF + U32_BYTES]);
if version != format::fts::VERSION_V1_LEGACY
&& version != format::fts::VERSION_V2
&& version != format::fts::VERSION_V3
&& version != format::fts::VERSION_V4
{
return Err(FtsError::Read(ReadError::UnsupportedVersion(format!(
"fts section version {version}"
))));
}
let header_size = match version {
v if v == format::fts::VERSION_V2
|| v == format::fts::VERSION_V3
|| v == format::fts::VERSION_V4 =>
{
format::fts::HEADER_SIZE_V2
}
_ => FTS_HEADER_SIZE,
};
if header.len() < header_size {
return Err(FtsError::Read(ReadError::MissingKv("fts header")));
}
let postings_offset =
read_u64_le(&header[hdr::POSTINGS_OFFSET_OFF..hdr::POSTINGS_OFFSET_OFF + U64_BYTES])
as usize;
let doc_lengths_table_offset =
read_u64_le(&header[hdr::DOC_LENGTHS_DIR_OFF..hdr::DOC_LENGTHS_DIR_OFF + U64_BYTES])
as usize;
let (fst_region, doc_lengths_tail) = futures::try_join!(
fetch_lazy_range(source.as_ref(), header_size..postings_offset, "fts/dict"),
fetch_lazy_range(
source.as_ref(),
doc_lengths_table_offset..fts_blob_len,
"fts/doc_lengths_tail",
),
)?;
let mut overlay = PrefetchedSource::new(source);
overlay.install(0, header);
overlay.install(header_size as u64, fst_region);
overlay.install(doc_lengths_table_offset as u64, doc_lengths_tail);
Self::open_with_source(Source::Lazy(Arc::new(overlay)), columns_json, opts)
}
pub(crate) fn open_with_source(
source: Source,
columns_json: &str,
opts: OpenOptions,
) -> Result<Self, FtsError> {
let source_len = source.len();
if source_len < FTS_HEADER_SIZE {
return Err(FtsError::Read(ReadError::MissingKv("fts header")));
}
let header = fetch_source_range(&source, 0..FTS_HEADER_SIZE, "fts header")?;
if &header[0..MAGIC_BYTES] != format::fts::MAGIC {
return Err(FtsError::Read(ReadError::BadMagic {
section: "fts",
expected: format::fts::MAGIC,
actual: header[0..MAGIC_BYTES].to_vec(),
}));
}
let version = read_u32_le(&header[hdr::VERSION_OFF..hdr::VERSION_OFF + U32_BYTES]);
let positional_blob = match version {
v if v == format::fts::VERSION_V1_LEGACY => false,
v if v == format::fts::VERSION_V2 => true,
v if v == format::fts::VERSION_V3 => true,
v if v == format::fts::VERSION_V4 => true,
_ => {
return Err(FtsError::Read(ReadError::UnsupportedVersion(format!(
"fts section version {version}"
))));
}
};
let has_position_subindex =
version == format::fts::VERSION_V3 || version == format::fts::VERSION_V4;
let has_bitset_blocks = version == format::fts::VERSION_V4;
let header_size = match positional_blob {
true => format::fts::HEADER_SIZE_V2,
false => FTS_HEADER_SIZE,
};
if source_len < header_size {
return Err(FtsError::Read(ReadError::MissingKv("fts header")));
}
let n_columns =
read_u32_le(&header[hdr::N_COLUMNS_OFF..hdr::N_COLUMNS_OFF + U32_BYTES]) as usize;
let n_docs = read_u32_le(&header[hdr::N_DOCS_OFF..hdr::N_DOCS_OFF + U32_BYTES]);
let n_terms_total = read_u32_le(&header[hdr::N_TERMS_OFF..hdr::N_TERMS_OFF + U32_BYTES]);
let fst_offset =
read_u64_le(&header[hdr::FST_OFFSET_OFF..hdr::FST_OFFSET_OFF + U64_BYTES]) as usize;
let postings_offset =
read_u64_le(&header[hdr::POSTINGS_OFFSET_OFF..hdr::POSTINGS_OFFSET_OFF + U64_BYTES])
as usize;
let doc_lengths_table_offset =
read_u64_le(&header[hdr::DOC_LENGTHS_DIR_OFF..hdr::DOC_LENGTHS_DIR_OFF + U64_BYTES])
as usize;
let positions_offset: Option<usize> = match positional_blob {
true => {
let ext = fetch_source_range(
&source,
FTS_HEADER_SIZE..format::fts::HEADER_SIZE_V2,
"fts header ext",
)?;
Some(read_u64_le(&ext[0..U64_BYTES]) as usize)
}
false => None,
};
let postings_end = positions_offset.unwrap_or(doc_lengths_table_offset);
if fst_offset < header_size
|| postings_offset < fst_offset + 4
|| postings_end < postings_offset + 4
|| doc_lengths_table_offset < postings_end
|| doc_lengths_table_offset > source_len
|| positions_offset.is_some_and(|po| doc_lengths_table_offset < po + 4)
{
return Err(FtsError::Read(ReadError::MalformedVersion(format!(
"fts header offsets out of range: fst={fst_offset}, postings={postings_offset}, \
positions={positions_offset:?}, doc_lengths={doc_lengths_table_offset}, \
blob_len={}",
source_len
))));
}
let fst_range = fst_offset..postings_offset.saturating_sub(4); let postings_range = postings_offset..postings_end.saturating_sub(4); let positions_range: Option<Range<usize>> =
positions_offset.map(|po| po..doc_lengths_table_offset.saturating_sub(4));
if opts.verify_crc {
let fst_crc_bytes = fetch_source_range(
&source,
postings_offset.saturating_sub(4)..postings_offset,
"fts/dict crc",
)?;
let fst_crc_expected = read_u32_le(&fst_crc_bytes);
let fst_bytes = fetch_source_range(&source, fst_range.clone(), "fts/dict")?;
let fst_crc_actual = crc32c(&fst_bytes);
if fst_crc_expected != fst_crc_actual {
return Err(FtsError::Read(ReadError::ChecksumMismatch {
section: "fts/dict",
column: String::new(),
}));
}
}
if opts.verify_crc {
let postings_crc_pos = postings_end.saturating_sub(4);
let postings_crc_bytes =
fetch_source_range(&source, postings_crc_pos..postings_end, "fts/postings crc")?;
let postings_crc_expected = read_u32_le(&postings_crc_bytes);
let postings_bytes =
fetch_source_range(&source, postings_range.clone(), "fts/postings")?;
let postings_crc_actual = crc32c(&postings_bytes);
if postings_crc_expected != postings_crc_actual {
return Err(FtsError::Read(ReadError::ChecksumMismatch {
section: "fts/postings",
column: String::new(),
}));
}
}
if opts.verify_crc
&& let Some(pos_range) = &positions_range
{
let crc_pos = doc_lengths_table_offset.saturating_sub(4);
let crc_bytes = fetch_source_range(
&source,
crc_pos..doc_lengths_table_offset,
"fts/positions crc",
)?;
let crc_expected = read_u32_le(&crc_bytes);
let pos_bytes = fetch_source_range(&source, pos_range.clone(), "fts/positions")?;
let crc_actual = crc32c(&pos_bytes);
if crc_expected != crc_actual {
return Err(FtsError::Read(ReadError::ChecksumMismatch {
section: "fts/positions",
column: String::new(),
}));
}
}
let cols: Vec<FtsColumnConfig> = serde_json::from_str(columns_json).map_err(|e| {
FtsError::Read(ReadError::MalformedVersion(format!(
"inf.fts.columns JSON: {e}"
)))
})?;
if cols.len() != n_columns {
return Err(FtsError::Read(ReadError::MalformedVersion(format!(
"inf.fts.columns has {} entries, header says {}",
cols.len(),
n_columns
))));
}
let dir_size = n_columns * DOC_LENGTHS_ENTRY_SIZE;
let dir_end = doc_lengths_table_offset + dir_size;
if dir_end + 4 > source_len {
return Err(FtsError::Read(ReadError::MalformedVersion(
"doc-lengths directory runs past blob end".into(),
)));
}
let dir_region = fetch_source_range(
&source,
doc_lengths_table_offset..dir_end + 4,
"fts/doc_lengths_dir",
)?;
let dir_bytes = &dir_region[..dir_size];
if opts.verify_crc {
let dir_crc_expected = read_u32_le(&dir_region[dir_size..dir_size + 4]);
let dir_crc_actual = crc32c(dir_bytes);
if dir_crc_expected != dir_crc_actual {
return Err(FtsError::Read(ReadError::ChecksumMismatch {
section: "fts/doc_lengths_dir",
column: String::new(),
}));
}
}
let mut columns = Vec::with_capacity(n_columns);
let mut column_id_by_name = HashMap::with_capacity(n_columns);
for (i, col_cfg) in cols.iter().enumerate() {
let entry_off = i * DOC_LENGTHS_ENTRY_SIZE;
let column_id = u32::from_le_bytes([
dir_bytes[entry_off],
dir_bytes[entry_off + 1],
dir_bytes[entry_off + 2],
dir_bytes[entry_off + 3],
]);
let doc_lengths_offset =
read_u64_le(&dir_bytes[entry_off + 4..entry_off + 12]) as usize;
let avgdl_x1000 = read_u32_le(&dir_bytes[entry_off + 12..entry_off + 16]) as u64;
if column_id != i as u32 {
return Err(FtsError::Read(ReadError::MalformedVersion(format!(
"doc-lengths directory entry {i} has column_id {column_id}"
))));
}
let array_byte_len = 4 * n_docs as usize;
let array_end = doc_lengths_offset + array_byte_len;
if array_end + 4 > source_len {
return Err(FtsError::Read(ReadError::MalformedVersion(format!(
"doc-lengths array {i} runs past blob end"
))));
}
let array_region = fetch_source_range(
&source,
doc_lengths_offset..array_end + 4,
"fts/doc_lengths_array",
)?;
if opts.verify_crc {
let array_crc_expected =
read_u32_le(&array_region[array_byte_len..array_byte_len + 4]);
let array_crc_actual = crc32c(&array_region[..array_byte_len]);
if array_crc_expected != array_crc_actual {
return Err(FtsError::Read(ReadError::ChecksumMismatch {
section: "fts/doc_lengths_array",
column: format!(" (column '{}')", col_cfg.name),
}));
}
}
let avgdl = (avgdl_x1000 as f32) / format::fts::AVGDL_FIXED_POINT_SCALE;
let n = n_docs as usize;
let dl_norm_k1 = NormTable::new(
(0..n).map(|d| read_u32_le(&array_region[d * 4..d * 4 + 4])),
n,
avgdl,
);
let tokenizer = tokenizer_for_name(&col_cfg.tokenizer).ok_or_else(|| {
FtsError::Read(ReadError::MalformedVersion(format!(
"inf.fts.columns: unknown tokenizer {:?} for column {:?}",
col_cfg.tokenizer, col_cfg.name
)))
})?;
columns.push(ColumnMeta {
name: col_cfg.name.clone(),
doc_lengths_range: doc_lengths_offset..array_end,
avgdl,
dl_norm_k1,
positions: col_cfg.positions,
tokenizer,
});
column_id_by_name.insert(col_cfg.name.clone(), i as u32);
}
Ok(FtsReader {
source,
n_docs,
n_terms_total,
fst_range,
postings_range,
positions_range,
has_position_subindex,
has_bitset_blocks,
columns,
column_id_by_name,
})
}
pub fn n_docs(&self) -> u32 {
self.n_docs
}
pub fn n_terms(&self) -> u32 {
self.n_terms_total
}
pub fn fts_columns(&self) -> impl Iterator<Item = &str> {
self.columns.iter().map(|c| c.name.as_str())
}
pub fn fts_columns_config(&self) -> impl Iterator<Item = &ColumnMeta> {
self.columns.iter()
}
pub fn column_tokenizer(&self, column: &str) -> Result<Arc<dyn Tokenizer>, FtsError> {
let id = self.resolve_column_id(column)?;
Ok(Arc::clone(&self.columns[id as usize].tokenizer))
}
fn dict_bytes(&self) -> Result<Bytes, FtsError> {
fetch_source_range(&self.source, self.fst_range.clone(), "fts/dict")
}
pub(super) async fn dict_bytes_async(&self) -> Result<Bytes, FtsError> {
self.source
.range_async(self.fst_range.clone())
.await
.map_err(|e| {
FtsError::Read(ReadError::MalformedVersion(format!(
"fts/dict range fetch failed: {e}"
)))
})
}
pub(super) async fn fetch_term_postings(
&self,
terms: &[(usize, Option<usize>)],
) -> Result<Vec<Bytes>, FtsError> {
if terms.is_empty() {
return Ok(Vec::new());
}
let probe_ranges: Vec<(usize, usize)> = terms
.iter()
.filter(|(_, len)| len.is_none())
.map(|&(metadata_offset, _)| (metadata_offset, TERM_META_SIZE))
.collect();
let probed = self.fetch_ranges(&probe_ranges).await?;
let mut resolved: Vec<(usize, usize)> = Vec::with_capacity(terms.len());
let mut next_probe = 0usize;
for &(metadata_offset, slot_length) in terms {
let postings_length = match slot_length {
Some(length) => length,
None => {
let header = probed.get(next_probe).ok_or_else(|| {
FtsError::Read(ReadError::MalformedVersion(
"fetched fewer term metadata headers than probed".into(),
))
})?;
next_probe += 1;
header_postings_length(header.as_ref())?
}
};
resolved.push((metadata_offset, postings_length));
}
self.fetch_ranges(&resolved).await
}
async fn fetch_ranges(&self, terms: &[(usize, usize)]) -> Result<Vec<Bytes>, FtsError> {
if terms.is_empty() {
return Ok(Vec::new());
}
let base = self.postings_range.start;
let region_len = self.postings_range.len();
let mut ranges: Vec<Range<usize>> = Vec::with_capacity(terms.len());
for &(m, postings_length) in terms {
if postings_length < TERM_META_SIZE || m + postings_length > region_len {
return Err(FtsError::Read(ReadError::MalformedVersion(
"term postings range runs past postings region".into(),
)));
}
ranges.push(base + m..base + m + postings_length);
}
let plan = RangeCoalescePlan::new(
&ranges,
TERM_RANGE_COALESCE_MAX_GAP,
TERM_RANGE_COALESCE_MAX_OVERFETCH,
);
let fetched = self
.source
.get_ranges_parallel_async(plan.fetch_ranges())
.await
.map_err(|e| {
FtsError::Read(ReadError::MalformedVersion(format!(
"fts/postings term body range fetch failed: {e}"
)))
})?;
Ok(plan.restore(&fetched))
}
async fn fetch_term_positions(&self, terms: &[(u64, u32)]) -> Result<Vec<Bytes>, FtsError> {
if terms.iter().all(|&(_, len)| len == 0) {
return Ok(vec![Bytes::new(); terms.len()]);
}
let region = self.positions_range.as_ref().ok_or_else(|| {
FtsError::Read(ReadError::MalformedVersion(
"positional term in a blob with no positions region".into(),
))
})?;
let base = region.start;
let region_len = region.len();
let mut ranges: Vec<Range<usize>> = Vec::with_capacity(terms.len());
for &(off, len) in terms {
let off = off as usize;
let len = len as usize;
if off + len > region_len {
return Err(FtsError::Read(ReadError::MalformedVersion(
"term positions range runs past positions region".into(),
)));
}
ranges.push(base + off..base + off + len);
}
self.source
.get_ranges_parallel_async(&ranges)
.await
.map_err(|e| {
FtsError::Read(ReadError::MalformedVersion(format!(
"fts/positions term range fetch failed: {e}"
)))
})
}
pub(super) async fn build_atom_cursors(
&self,
column_id: u32,
terms: &[&str],
phrases: &[Vec<String>],
global_idf: Option<&GlobalTermIdf>,
) -> Result<(Vec<Option<AnyCursor>>, u64), FtsError> {
let col_meta = &self.columns[column_id as usize];
if !phrases.is_empty() && !col_meta.positions {
return Err(FtsError::PositionsUnavailable {
column: col_meta.name.clone(),
});
}
let mut dict_ranges = 0u64;
let mut out: Vec<Option<AnyCursor>> = Vec::with_capacity(terms.len() + phrases.len());
for term in terms {
let mut cursors = self
.build_term_cursors(column_id, &[term], global_idf, false)
.await?;
dict_ranges += 1;
out.push(cursors.pop().map(AnyCursor::Term));
}
for phrase in phrases {
let member_refs: Vec<&str> = phrase.iter().map(|t| t.as_str()).collect();
let cursors = self
.build_term_cursors(column_id, &member_refs, global_idf, false)
.await?;
dict_ranges += 1;
if cursors.len() != member_refs.len() {
out.push(None);
continue;
}
let mut positional: Vec<(Option<TermMeta>, Option<u32>)> =
Vec::with_capacity(cursors.len());
for (cursor, term) in cursors.iter().zip(&member_refs) {
match cursor.bytes.is_empty() {
false => {
let term_meta = TermMeta::parse(
cursor.bytes.as_ref(),
0,
true,
self.has_position_subindex,
)?;
positional.push((Some(term_meta), None));
}
true => {
dict_ranges += 1;
let fst_bytes = self.dict_bytes_async().await?;
let dict = DictReader::open(&fst_bytes).map_err(|e| {
FtsError::Read(ReadError::MalformedVersion(format!(
"FST parse failed: {e}"
)))
})?;
let key = make_key(&col_meta.name, term);
let packed = dict
.lookup(&key)
.expect("inline member cursor was built from this dict");
let position = match FstValue::unpack(packed) {
FstValue::Inline { tf: slot, .. } => slot,
FstValue::Pfor { .. } => {
unreachable!("inline cursor from a PFOR FST value")
}
};
positional.push((None, Some(position)));
}
}
}
let pos_ranges: Vec<(u64, u32)> = positional
.iter()
.map(|(term_meta, _)| {
term_meta
.map(|tm| (tm.positions_offset, tm.positions_length))
.unwrap_or((0, 0))
})
.collect();
let positions = self.fetch_term_positions(&pos_ranges).await?;
out.push(Some(AnyCursor::Phrase(PhraseCursor::new(
cursors, positions, positional,
)?)));
}
Ok((out, dict_ranges))
}
pub fn iter_column_terms(&self, column: &str) -> Result<Vec<Vec<u8>>, FtsError> {
self.iter_terms_with_prefix(column, b"")
}
pub(crate) fn for_each_term_posting(
&self,
column_id: u32,
mut emit: impl FnMut(&[u8], u32, u32, &[u32]) -> Result<(), FtsError>,
) -> Result<(), FtsError> {
let col_meta = &self.columns[column_id as usize];
let positional = col_meta.positions;
let n_docs = u64::from(self.n_docs);
let column_name = col_meta.name.clone();
let region_base = self.postings_range.start;
let positions_region = self.positions_range.clone();
let fst_bytes = self.dict_bytes()?;
let dict = DictReader::open(&fst_bytes).map_err(|e| {
FtsError::Read(ReadError::MalformedVersion(format!(
"FST parse failed: {e}"
)))
})?;
let mut column_prefix = column_name.as_bytes().to_vec();
column_prefix.push(FST_SEPARATOR);
let prefix_len = column_prefix.len();
let mut positions_buf: Vec<u32> = Vec::new();
for (key, packed) in dict.iter_prefix(&column_prefix) {
let term = &key[prefix_len..];
match FstValue::unpack(packed) {
FstValue::Inline { doc_id, tf } => {
if positional {
emit(term, doc_id, 1, &[tf])?;
} else {
emit(term, doc_id, tf, &[])?;
}
}
FstValue::Pfor {
metadata_offset,
postings_length_hint,
} => {
let start = region_base + metadata_offset as usize;
let postings_length = match postings_length_hint {
Some(len) => len as usize,
None => {
let header = fetch_source_range(
&self.source,
start..start + TERM_META_SIZE,
"fts/merge header",
)?;
header_postings_length(header.as_ref())?
}
};
let term_bytes = fetch_source_range(
&self.source,
start..start + postings_length,
"fts/merge postings",
)?;
let position_bytes = if positional {
let meta = TermMeta::parse(term_bytes.as_ref(), 0, true, false)?;
let region = positions_region.as_ref().ok_or_else(|| {
FtsError::Read(ReadError::MalformedVersion(
"positional column missing a positions region".into(),
))
})?;
let pstart = region.start + meta.positions_offset as usize;
let pend = pstart + meta.positions_length as usize;
Some(fetch_source_range(
&self.source,
pstart..pend,
"fts/merge positions",
)?)
} else {
None
};
let mut pos_at = 0usize;
let mut cursor =
TermCursor::new(term_bytes, n_docs, positional, None, false, false)?;
while !cursor.is_exhausted() {
while cursor.pos < cursor.block_n {
let doc_id = cursor.block_doc_ids[cursor.pos];
let tf = cursor.block_tfs[cursor.pos];
let positions: &[u32] = match &position_bytes {
Some(bytes) => {
positions_buf.clear();
decode_run(bytes.as_ref(), &mut pos_at, tf, &mut positions_buf)
.ok_or_else(|| {
FtsError::Read(ReadError::MalformedVersion(
"truncated position run in merge read".into(),
))
})?;
&positions_buf
}
None => &[],
};
emit(term, doc_id, tf, positions)?;
cursor.pos += 1;
}
cursor.next();
}
}
}
}
Ok(())
}
pub(crate) fn read_doc_lengths(&self, column_id: u32) -> Result<Vec<u32>, FtsError> {
let n = self.n_docs as usize;
let range = self.columns[column_id as usize].doc_lengths_range.clone();
let bytes = fetch_source_range(&self.source, range, "fts/merge doc_lengths")?;
let region = bytes.as_ref();
if region.len() < n * U32_BYTES {
return Err(FtsError::Read(ReadError::MalformedVersion(
"doc-lengths region shorter than n_docs entries".into(),
)));
}
Ok((0..n)
.map(|d| read_u32_le(®ion[d * U32_BYTES..d * U32_BYTES + U32_BYTES]))
.collect())
}
pub fn iter_terms_with_prefix(
&self,
column: &str,
term_prefix: &[u8],
) -> Result<Vec<Vec<u8>>, FtsError> {
if !self.column_id_by_name.contains_key(column) {
return Ok(Vec::new());
}
let mut full_prefix = column.as_bytes().to_vec();
full_prefix.push(FST_SEPARATOR);
let column_prefix_len = full_prefix.len();
full_prefix.extend_from_slice(term_prefix);
let fst_bytes = self
.dict_bytes()
.expect("FST bytes must be available for term iteration");
let dict = DictReader::open(&fst_bytes).map_err(|e| {
FtsError::Read(ReadError::MalformedVersion(format!(
"FST parse failed: {e}"
)))
})?;
let pairs = dict.iter_prefix(&full_prefix);
Ok(pairs
.into_iter()
.map(|(key, _)| key[column_prefix_len..].to_vec())
.collect())
}
}
pub(crate) struct OrCursorSet {
pub(super) column_id: u32,
pub(super) cursors: Vec<TermCursor>,
}
impl OrCursorSet {
pub(crate) fn len(&self) -> usize {
self.cursors.len()
}
pub(crate) fn postings_bytes(&self) -> u64 {
term_cursor_bytes(&self.cursors)
}
pub(crate) fn planned_ranges(&self) -> u64 {
term_cursor_ranges(&self.cursors)
}
}
pub(super) fn top_k(scores: HashMap<u32, f32>, k: usize) -> Vec<(u32, f32)> {
let mut sorted: Vec<(u32, f32)> = scores.into_iter().collect();
sorted.sort_unstable_by_key(|(d, _)| *d);
let mut heap: BinaryHeap<TopKEntry> = BinaryHeap::with_capacity(k.min(sorted.len()).max(1));
for (doc_id, score) in sorted {
if heap.len() < k {
heap.push(TopKEntry(score, doc_id));
} else if let Some(TopKEntry(top_score, _)) = heap.peek()
&& score > *top_score
{
heap.pop();
heap.push(TopKEntry(score, doc_id));
}
}
drain_top_k_desc(heap)
}
fn fetch_source_range(source: &Source, range: Range<usize>, what: &str) -> Result<Bytes, FtsError> {
source.get_range(range).map_err(|e| {
FtsError::Read(ReadError::MalformedVersion(format!(
"{what} lazy source range fetch failed: {e}"
)))
})
}
async fn fetch_lazy_range(
source: &dyn LazyByteSource,
range: Range<usize>,
what: &str,
) -> Result<Bytes, FtsError> {
source
.range(range.start as u64, range.len() as u64)
.await
.map_err(|e| {
FtsError::Read(ReadError::MalformedVersion(format!(
"{what} lazy source range fetch failed: {e}"
)))
})
}
#[inline]
pub(super) fn read_u32_le(b: &[u8]) -> u32 {
u32::from_le_bytes([b[0], b[1], b[2], b[3]])
}
#[inline]
pub(super) fn read_u64_le(b: &[u8]) -> u64 {
let mut buf = [0u8; 8];
buf.copy_from_slice(&b[0..8]);
u64::from_le_bytes(buf)
}
pub(super) fn or_walk_unranked(mut cursors: Vec<TermCursor>, mut emit: impl FnMut(u32)) {
loop {
let min_doc = cursors
.iter()
.filter(|c| !c.is_exhausted())
.map(TermCursor::current_doc_id)
.min();
let Some(min_doc) = min_doc else { break };
emit(min_doc);
for c in cursors.iter_mut() {
if !c.is_exhausted() && c.current_doc_id() == min_doc {
c.next();
}
}
}
}
pub(super) fn or_merge_unranked(cursors: Vec<TermCursor>) -> Vec<u32> {
let mut out = Vec::new();
or_walk_unranked(cursors, |doc| out.push(doc));
out
}
pub(super) fn or_count_unranked(mut cursors: Vec<TermCursor>) -> u64 {
if let Some(anchor) = dominant_anchor_index(&cursors) {
return or_count_anchored(cursors, anchor);
}
let total_df: u64 = cursors.iter().map(|c| c.df).sum();
let max_doc = cursors
.iter()
.filter_map(|c| c.blocks.last())
.map(|b| b.last_doc_id)
.max()
.unwrap_or(0);
if total_df.saturating_mul(OR_COUNT_BITSET_DENSITY_DIVISOR) >= u64::from(max_doc) {
return or_count_bitset(cursors, max_doc);
}
let mut present = [0u64; OR_WINDOW_WORDS];
let mut n = 0u64;
loop {
let mut min_doc = u32::MAX;
for c in &cursors {
if !c.is_exhausted() {
min_doc = min_doc.min(c.current_doc_id());
}
}
if min_doc == u32::MAX {
break;
}
let base = min_doc & !(OR_WINDOW - 1);
let window_end = base.saturating_add(OR_WINDOW);
for c in &mut cursors {
while !c.is_exhausted() {
let d = c.current_doc_id();
if d >= window_end {
break;
}
let local = (d - base) as usize;
present[local >> 6] |= 1u64 << (local & 63);
c.next();
}
}
for word in present.iter_mut() {
n += word.count_ones() as u64;
*word = 0;
}
}
n
}
fn or_count_bitset(cursors: Vec<TermCursor>, max_doc: u32) -> u64 {
let words = max_doc as usize / 64 + 1;
let mut union = vec![0u64; words];
let mut scratch = [0u32; BLOCK_LEN];
for c in &cursors {
or_cursor_into_bitset(&mut union, c, &mut scratch);
}
union.iter().map(|w| w.count_ones() as u64).sum()
}
pub(super) fn or_cursor_into_bitset(
dest: &mut [u64],
c: &TermCursor,
scratch: &mut [u32; BLOCK_LEN],
) {
if c.bytes.is_empty() {
for &d in &c.block_doc_ids[..c.block_n] {
dest[(d >> 6) as usize] |= 1u64 << (d & 63);
}
return;
}
for block in c.blocks.iter() {
let bytes = c.bytes.slice(block.block_byte_offset..block.block_byte_end);
let bytes = bytes.as_ref();
if bytes[posting::ENCODING_OFF] == ENCODING_BITSET {
let base_word = read_u32_le(&bytes[4..8]) as usize / 64;
let tf_bits = bytes[2] as usize;
let tfs_size = BLOCK_LEN * tf_bits / 8;
let presence = &bytes[posting::HEADER_SIZE..bytes.len() - tfs_size];
for (i, chunk) in presence.chunks_exact(8).enumerate() {
dest[base_word + i] |= u64::from_le_bytes(chunk.try_into().expect("8 bytes"));
}
} else {
let n = decode_block_doc_ids(bytes, scratch);
for &d in &scratch[..n] {
dest[(d >> 6) as usize] |= 1u64 << (d & 63);
}
}
}
}
fn dominant_anchor_index(cursors: &[TermCursor]) -> Option<usize> {
dominant_anchor_of_dfs(cursors.iter().map(|c| c.df))
}
fn dominant_anchor_of_dfs(dfs: impl IntoIterator<Item = u64>) -> Option<usize> {
let dfs: Vec<u64> = dfs.into_iter().collect();
if dfs.len() < 2 {
return None;
}
let (max_idx, &max_df) = dfs.iter().enumerate().max_by_key(|&(_, df)| *df)?;
let others_df: u64 = dfs
.iter()
.enumerate()
.filter(|&(i, _)| i != max_idx)
.map(|(_, df)| *df)
.sum();
match max_df >= others_df.saturating_mul(OR_COUNT_ANCHOR_DOMINANCE) {
true => Some(max_idx),
false => None,
}
}
fn or_count_anchored(mut cursors: Vec<TermCursor>, anchor_idx: usize) -> u64 {
let mut anchor = cursors.swap_remove(anchor_idx);
let mut n = anchor.df;
loop {
let mut min_doc = u32::MAX;
for c in &cursors {
if !c.is_exhausted() {
min_doc = min_doc.min(c.current_doc_id());
}
}
if min_doc == u32::MAX {
break;
}
if !anchor.contains(min_doc) {
n += 1;
}
for c in cursors.iter_mut() {
if !c.is_exhausted() && c.current_doc_id() == min_doc {
c.next();
}
}
}
n
}
fn header_postings_length(header: &[u8]) -> Result<usize, FtsError> {
let field_end = term_meta::POSTINGS_LENGTH_OFF + U32_BYTES;
if header.len() < field_end {
return Err(FtsError::Read(ReadError::MalformedVersion(
"term metadata header shorter than its postings_length field".into(),
)));
}
Ok(read_u32_le(&header[term_meta::POSTINGS_LENGTH_OFF..field_end]) as usize)
}
#[cfg(test)]
mod tests {
use std::collections::HashSet;
use super::{super::test_util::*, *};
use crate::superfile::{
BytesLazyByteSource,
fts::{builder::FtsBuilder, reader::BoolMode, tokenize::AsciiLowerTokenizer},
};
#[test]
fn open_accepts_valid_blob() {
let (blob, json) = build_blob();
let r = FtsReader::open(blob, &json).expect("open should succeed");
assert_eq!(r.n_docs(), 3);
assert!(r.n_terms() > 0);
assert_eq!(r.fts_columns().collect::<Vec<_>>(), vec!["body"]);
}
#[test]
fn for_each_term_posting_round_trips_doc_ids_and_tfs() {
use std::collections::BTreeMap;
let (blob, json) = build_blob();
let r = FtsReader::open(blob, &json).expect("open");
let mut got: BTreeMap<Vec<u8>, Vec<(u32, u32)>> = BTreeMap::new();
r.for_each_term_posting(0, |term, doc_id, tf, positions| {
assert!(
positions.is_empty(),
"non-positional column yields no positions"
);
got.entry(term.to_vec()).or_default().push((doc_id, tf));
Ok(())
})
.expect("stream postings");
for postings in got.values() {
assert!(
postings.windows(2).all(|w| w[0].0 < w[1].0),
"doc_ids must be ascending"
);
}
let t = |s: &str| s.as_bytes().to_vec();
assert_eq!(
got.get(&t("rust")).expect("term streamed").as_slice(),
&[(0, 1), (1, 1)]
);
assert_eq!(
got.get(&t("runtime")).expect("term streamed").as_slice(),
&[(0, 1), (1, 1)]
);
assert_eq!(
got.get(&t("async")).expect("term streamed").as_slice(),
&[(0, 1)]
);
assert_eq!(
got.get(&t("tokio")).expect("term streamed").as_slice(),
&[(1, 1)]
);
assert_eq!(
got.get(&t("java")).expect("term streamed").as_slice(),
&[(2, 1)]
);
assert_eq!(
got.get(&t("boot")).expect("term streamed").as_slice(),
&[(2, 1)]
);
assert_eq!(got.len() as u32, r.n_terms());
}
#[test]
fn for_each_term_posting_round_trips_positions() {
use std::collections::BTreeMap;
let tok = Arc::new(AsciiLowerTokenizer);
let mut b = FtsBuilder::new(tok);
b.register_column("body".into(), true)
.expect("register positional column");
b.add_doc(0, 0, "a b a").expect("add doc 0");
b.add_doc(0, 1, "b a c").expect("add doc 1");
let bytes = b.finish().expect("finish");
let json = r#"[{"name":"body","tokenizer":"ascii_lower","positions":true}]"#;
let r = FtsReader::open(Bytes::from(bytes), json).expect("open");
let mut got: BTreeMap<Vec<u8>, Vec<(u32, u32, Vec<u32>)>> = BTreeMap::new();
r.for_each_term_posting(0, |term, doc_id, tf, positions| {
got.entry(term.to_vec())
.or_default()
.push((doc_id, tf, positions.to_vec()));
Ok(())
})
.expect("stream positional postings");
let t = |s: &str| s.as_bytes().to_vec();
assert_eq!(
got.get(&t("a")).expect("term streamed").as_slice(),
&[(0, 2, vec![0, 2]), (1, 1, vec![1])]
);
assert_eq!(
got.get(&t("b")).expect("term streamed").as_slice(),
&[(0, 1, vec![1]), (1, 1, vec![0])]
);
assert_eq!(
got.get(&t("c")).expect("term streamed").as_slice(),
&[(1, 1, vec![2])]
);
}
#[test]
fn add_prebuilt_term_posting_round_trips_read_to_write() {
use std::collections::BTreeMap;
let json = r#"[{"name":"body","tokenizer":"ascii_lower","positions":true}]"#;
let tok = Arc::new(AsciiLowerTokenizer);
let mut a = FtsBuilder::new(tok.clone());
a.register_column("body".into(), true).expect("register a");
a.add_doc(0, 0, "a b a").expect("a doc 0");
a.add_doc(0, 1, "b a c").expect("a doc 1");
let ra = FtsReader::open(Bytes::from(a.finish().expect("finish a")), json).expect("open a");
let mut b = FtsBuilder::new(tok);
b.register_column("body".into(), true).expect("register b");
ra.for_each_term_posting(0, |term, doc_id, tf, positions| {
let term_str = std::str::from_utf8(term).expect("utf8 term");
b.add_prebuilt_term_posting(0, term_str, doc_id, tf, positions)
.expect("prebuilt push");
Ok(())
})
.expect("feed prebuilt postings");
b.set_prebuilt_doc_lengths(0, ra.read_doc_lengths(0).expect("doc lengths"));
let rb = FtsReader::open(Bytes::from(b.finish().expect("finish b")), json).expect("open b");
let collect = |r: &FtsReader| {
let mut m: BTreeMap<Vec<u8>, Vec<(u32, u32, Vec<u32>)>> = BTreeMap::new();
r.for_each_term_posting(0, |t, d, tf, p| {
m.entry(t.to_vec()).or_default().push((d, tf, p.to_vec()));
Ok(())
})
.expect("collect");
m
};
assert_eq!(
collect(&ra),
collect(&rb),
"prebuilt-fed postings must match"
);
assert_eq!(rb.n_docs(), 2);
assert_eq!(rb.n_terms(), ra.n_terms());
}
#[test]
fn add_prebuilt_term_posting_spilled_round_trips() {
use std::collections::BTreeMap;
let json = r#"[{"name":"body","tokenizer":"ascii_lower","positions":true}]"#;
let tok = Arc::new(AsciiLowerTokenizer);
let mut a = FtsBuilder::new(tok.clone());
a.register_column("body".into(), true).expect("register a");
a.add_doc(0, 0, "a b c a").expect("a doc 0");
a.add_doc(0, 1, "b c d").expect("a doc 1");
a.add_doc(0, 2, "a d e").expect("a doc 2");
let ra = FtsReader::open(Bytes::from(a.finish().expect("finish a")), json).expect("open a");
let mut b = FtsBuilder::new(tok);
b.register_column("body".into(), true).expect("register b");
b.set_spill_threshold_bytes(1);
ra.for_each_term_posting(0, |term, doc_id, tf, positions| {
let term_str = std::str::from_utf8(term).expect("utf8 term");
b.add_prebuilt_term_posting(0, term_str, doc_id, tf, positions)
.expect("prebuilt push (spilled)");
Ok(())
})
.expect("feed prebuilt postings");
b.set_prebuilt_doc_lengths(0, ra.read_doc_lengths(0).expect("doc lengths"));
let rb = FtsReader::open(Bytes::from(b.finish().expect("finish b")), json).expect("open b");
let collect = |r: &FtsReader| {
let mut m: BTreeMap<Vec<u8>, Vec<(u32, u32, Vec<u32>)>> = BTreeMap::new();
r.for_each_term_posting(0, |t, d, tf, p| {
m.entry(t.to_vec()).or_default().push((d, tf, p.to_vec()));
Ok(())
})
.expect("collect");
m
};
assert_eq!(
collect(&ra),
collect(&rb),
"spilled prebuilt-fed postings must match a fresh build"
);
assert_eq!(rb.n_docs(), 3);
assert_eq!(rb.n_terms(), ra.n_terms());
}
#[test]
fn read_doc_lengths_returns_token_counts() {
let tok = Arc::new(AsciiLowerTokenizer);
let mut b = FtsBuilder::new(tok);
b.register_column("body".into(), false).expect("register");
b.add_doc(0, 0, "a b a").expect("doc 0"); b.add_doc(0, 1, "b a c d").expect("doc 1"); let json = r#"[{"name":"body","tokenizer":"ascii_lower"}]"#;
let r = FtsReader::open(Bytes::from(b.finish().expect("finish")), json).expect("open");
assert_eq!(r.read_doc_lengths(0).expect("doc lengths"), vec![3, 4]);
}
#[test]
fn open_rejects_bad_magic() {
let (mut blob_vec, json) = build_blob();
let mut bytes = blob_vec.to_vec();
bytes[0] = b'X';
blob_vec = Bytes::from(bytes);
let err = FtsReader::open(blob_vec, &json).expect_err("expected error");
assert!(matches!(err, FtsError::Read(ReadError::BadMagic { .. })));
}
#[test]
fn open_rejects_short_blob() {
let err = FtsReader::open(Bytes::from(vec![0u8; 8]), "[]").expect_err("expected error");
assert!(matches!(err, FtsError::Read(_)));
}
#[test]
fn open_rejects_columns_json_mismatch() {
let (blob, _) = build_blob();
let bad_json = r#"[{"name":"body","tokenizer":"ascii_lower"},{"name":"title","tokenizer":"ascii_lower"}]"#;
let err = FtsReader::open(blob, bad_json).expect_err("expected error");
assert!(matches!(
err,
FtsError::Read(ReadError::MalformedVersion(_))
));
}
#[test]
fn df1_inline_form_flag_set_on_fst_value() {
let (blob, _json) = build_mixed_df_blob();
let header_size = 48usize;
let fst_off =
u64::from_le_bytes(blob[24..32].try_into().expect("fst_off slice is 8 bytes")) as usize;
let postings_off = u64::from_le_bytes(
blob[32..40]
.try_into()
.expect("postings_off slice is 8 bytes"),
) as usize;
let fst_bytes = &blob[fst_off..postings_off - 4];
let dict = DictReader::open(fst_bytes).expect("open dict");
assert_eq!(header_size, 48);
let val_common = dict.lookup(b"body\x1Fcommon").expect("common in FST");
let val_rust = dict.lookup(b"body\x1Frust").expect("rust in FST");
let val_uniq_d0 = dict.lookup(b"body\x1Funiqzero").expect("uniqzero in FST");
let val_uniq_d2 = dict.lookup(b"body\x1Funiqtwo").expect("uniqtwo in FST");
assert_eq!(val_common & 1, 0, "df=3 common term must use PFOR form");
assert_eq!(val_rust & 1, 0, "df=2 rust term must use PFOR form");
assert_eq!(val_uniq_d0 & 1, 1, "df=1 uniqzero must use inline form");
assert_eq!(val_uniq_d2 & 1, 1, "df=1 uniqtwo must use inline form");
match FstValue::unpack(val_uniq_d0) {
FstValue::Inline { doc_id, tf } => {
assert_eq!(doc_id, 0);
assert_eq!(tf, 1);
}
FstValue::Pfor { .. } => panic!("expected inline form"),
}
match FstValue::unpack(val_uniq_d2) {
FstValue::Inline { doc_id, tf } => {
assert_eq!(doc_id, 2);
assert_eq!(tf, 1);
}
FstValue::Pfor { .. } => panic!("expected inline form"),
}
}
#[test]
fn df1_inline_path_skips_postings_region_writes() {
let tok = Arc::new(AsciiLowerTokenizer);
let mut b_inline = FtsBuilder::new(tok.clone());
b_inline
.register_column("body".into(), false)
.expect("register column");
for i in 0..20 {
b_inline
.add_doc(0, i, &format!("uniq{i:03}"))
.expect("add doc");
}
let blob_inline = b_inline.finish().expect("finish inline");
let mut b_pfor = FtsBuilder::new(tok);
b_pfor
.register_column("body".into(), false)
.expect("register column");
for i in 0..20 {
let text = (0..20)
.map(|j| format!("uniq{j:03}"))
.collect::<Vec<_>>()
.join(" ");
b_pfor.add_doc(0, i, &text).expect("add doc");
}
let blob_pfor = b_pfor.finish().expect("finish pfor");
let postings_off_i = u64::from_le_bytes(
blob_inline[32..40]
.try_into()
.expect("postings_off_i slice is 8 bytes"),
) as usize;
let positions_off_i = u64::from_le_bytes(
blob_inline[48..56]
.try_into()
.expect("positions_off_i slice is 8 bytes"),
) as usize;
let postings_size_inline = positions_off_i - postings_off_i;
let postings_off_p = u64::from_le_bytes(
blob_pfor[32..40]
.try_into()
.expect("postings_off_p slice is 8 bytes"),
) as usize;
let positions_off_p = u64::from_le_bytes(
blob_pfor[48..56]
.try_into()
.expect("positions_off_p slice is 8 bytes"),
) as usize;
let postings_size_pfor = positions_off_p - postings_off_p;
assert_eq!(
postings_size_inline, 4,
"all-df=1 postings region should hold only the trailing CRC32; \
got {postings_size_inline} bytes"
);
assert!(
postings_size_pfor > 20 * 36,
"PFOR postings region should be hundreds of bytes; got {postings_size_pfor}"
);
}
#[test]
fn iter_column_terms_lists_every_term_in_lex_order() {
let (blob, json) = build_blob();
let r = FtsReader::open(blob, &json).expect("open");
let terms: Vec<String> = r
.iter_column_terms("body")
.expect("iter terms")
.into_iter()
.map(|b| String::from_utf8(b).expect("utf8"))
.collect();
let mut sorted = terms.clone();
sorted.sort();
assert_eq!(terms, sorted, "terms must be in lex order");
for expected in [
"rust", "async", "runtime", "tokio", "java", "spring", "boot",
] {
assert!(terms.contains(&expected.to_string()), "missing {expected}");
}
}
#[test]
fn iter_column_terms_unknown_column_is_empty() {
let (blob, json) = build_blob();
let r = FtsReader::open(blob, &json).expect("open");
assert!(r.iter_column_terms("nope").expect("ok").is_empty());
}
#[test]
fn iter_terms_with_prefix_bounds_the_walk() {
let (blob, json) = build_blob();
let r = FtsReader::open(blob, &json).expect("open");
let terms: Vec<String> = r
.iter_terms_with_prefix("body", b"run")
.expect("prefix walk")
.into_iter()
.map(|b| String::from_utf8(b).expect("utf8"))
.collect();
assert_eq!(terms, vec!["runtime".to_string()]);
assert!(
r.iter_terms_with_prefix("body", b"zzz")
.expect("prefix walk")
.is_empty()
);
}
#[test]
fn top_k_capacity_is_scoped_to_the_range_the_scan_visits() {
const N_DOCS: u64 = 1_000_000;
const BIG_K: usize = N_DOCS as usize;
assert_eq!(top_k_initial_capacity(BIG_K, N_DOCS, None), N_DOCS as usize);
let eighth = (N_DOCS / 8) as u32;
assert_eq!(
top_k_initial_capacity(BIG_K, N_DOCS, Some((0, eighth))),
eighth as usize
);
assert_eq!(
top_k_initial_capacity(BIG_K, N_DOCS, Some((eighth, 2 * eighth))),
eighth as usize
);
assert_eq!(
top_k_initial_capacity(BIG_K, N_DOCS, Some((0, u32::MAX))),
N_DOCS as usize
);
assert_eq!(top_k_initial_capacity(10, N_DOCS, Some((0, eighth))), 10);
assert_eq!(top_k_initial_capacity(0, N_DOCS, None), 1);
assert_eq!(top_k_initial_capacity(BIG_K, N_DOCS, Some((5, 5))), 1);
assert_eq!(
top_k_initial_capacity(usize::MAX, N_DOCS, None),
N_DOCS as usize
);
}
#[test]
fn read_u32_le_and_u64_le_decode_little_endian() {
let b32 = [0x78, 0x56, 0x34, 0x12];
assert_eq!(read_u32_le(&b32), 0x1234_5678);
let b64 = [0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
assert_eq!(read_u64_le(&b64), 1);
}
#[test]
fn dominant_anchor_routes_at_the_dominance_boundary() {
let k = OR_COUNT_ANCHOR_DOMINANCE;
assert_eq!(dominant_anchor_of_dfs([100 * k, 100]), Some(0));
assert_eq!(dominant_anchor_of_dfs([100 * k - 1, 100]), None);
assert_eq!(dominant_anchor_of_dfs([100, 100 * k]), Some(1));
assert_eq!(dominant_anchor_of_dfs([100 * k]), None);
assert_eq!(dominant_anchor_of_dfs([500, 500]), None);
assert_eq!(dominant_anchor_of_dfs([1500, 100, 100]), None);
assert_eq!(dominant_anchor_of_dfs([100 * k, 60, 40]), Some(0));
assert_eq!(dominant_anchor_of_dfs([0u64; 0]), None);
}
#[test]
fn top_k_keeps_highest_scores_with_doc_id_tiebreak() {
let mut scores: HashMap<u32, f32> = HashMap::new();
scores.insert(0, 1.0);
scores.insert(1, 3.0);
scores.insert(2, 2.0);
scores.insert(3, 3.0); let out = top_k(scores, 2);
assert_eq!(out, vec![(1, 3.0), (3, 3.0)]);
}
#[test]
fn top_k_smaller_than_k_returns_all_sorted() {
let mut scores: HashMap<u32, f32> = HashMap::new();
scores.insert(5, 2.0);
scores.insert(9, 5.0);
let out = top_k(scores, 10);
assert_eq!(out, vec![(9, 5.0), (5, 2.0)]);
}
#[tokio::test]
async fn open_lazy_round_trips_a_search() {
let (blob, json) = build_blob();
let src: Arc<dyn LazyByteSource> = Arc::new(BytesLazyByteSource::new(blob));
let r = FtsReader::open_lazy(src, &json, OpenOptions::for_object_store())
.await
.expect("open_lazy");
assert_eq!(r.n_docs(), 3);
let hits = r
.search("body", &["rust"], 10, BoolMode::Or)
.await
.expect("search over lazy reader");
let ids: HashSet<u32> = hits.iter().map(|(d, _)| *d).collect();
assert!(ids.contains(&0) && ids.contains(&1));
}
}