use crate::bitmap::{Postings, RoarPostings};
use crate::index::InfonIndex;
use crate::programs;
use crate::projector::{CorpusKind, Projector};
use crate::projectors::{CsvProjector, JsonProjector, JsonlProjector};
use crate::text::Gazetteer;
use crate::tokenql::evaluate;
use serde::Serialize;
use serde_json::Value;
use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::OnceLock;
#[cfg(not(target_arch = "wasm32"))]
use std::time::Instant;
const STOPW: &[&str] = &[
"the", "a", "an", "of", "to", "in", "on", "for", "and", "or", "is", "are", "was", "were", "be",
"do", "how", "what", "which", "who", "why", "when", "with", "without", "into", "over", "under",
"from", "your", "our", "their", "this", "that", "these", "those", "can", "may", "will", "would",
"should", "could", "not", "you", "use", "used", "using", "best", "common", "about", "across",
"based", "provide", "provides", "support", "supports", "need", "needs", "want", "wants", "able",
];
#[derive(Serialize, Default, Debug)]
pub struct FolderReport {
pub situations: u32,
pub ingested: Vec<(String, usize)>,
pub skipped: Vec<(String, String)>,
}
fn collect_files(dir: &Path) -> Vec<PathBuf> {
let mut out = Vec::new();
let mut stack = vec![dir.to_path_buf()];
while let Some(d) = stack.pop() {
let rd = match std::fs::read_dir(&d) {
Ok(r) => r,
Err(_) => continue,
};
for e in rd.flatten() {
let name = e.file_name().to_string_lossy().to_string();
if name.starts_with('.') || name == "node_modules" || name == "target" {
continue;
}
let p = e.path();
if p.is_dir() {
stack.push(p);
} else {
out.push(p);
}
}
}
out
}
#[cfg(feature = "onnx")]
fn default_text_engine() -> Option<crate::projectors::TextEngine> {
let ml = crate::paths::model_dir("step0_bundle_ml", "STEELDB_ML_BUNDLE", "spo.onnx")?;
let splade = crate::paths::model_dir("splade", "STEELDB_SPLADE_DIR", "splade.onnx");
crate::projectors::TextEngine::load(&ml, splade.as_deref()).ok()
}
pub struct Corpus {
ix: InfonIndex<RoarPostings>,
rows: Vec<Vec<String>>,
columns: Vec<String>,
source: String,
kind: CorpusKind,
noise: OnceLock<HashSet<String>>,
gaz: OnceLock<Option<Gazetteer>>,
gaz_overlay: Option<PathBuf>,
}
#[derive(Serialize)]
pub struct Hit {
pub sid: u32,
pub cells: Vec<String>,
}
#[derive(Serialize)]
pub struct QueryOut {
pub count: usize,
pub micros: f64,
pub columns: Vec<String>,
pub hits: Vec<Hit>,
}
#[derive(Serialize)]
pub struct Stats {
pub source: String,
pub kind: CorpusKind,
pub situations: u32,
pub vocab: usize,
pub columns: Vec<String>,
pub facets: Vec<(String, usize)>,
pub numeric_fields: Vec<String>,
}
impl Corpus {
pub fn from_projector(p: Box<dyn Projector>) -> std::io::Result<Corpus> {
let columns = p.columns();
let kind = p.kind();
let source = p.source();
let mut by_token: HashMap<String, Vec<u32>> = HashMap::new();
let mut rows: Vec<Vec<String>> = Vec::new();
let mut numbers_raw: Vec<(u32, String, f64)> = Vec::new();
let mut sid: u32 = 0;
p.project(&mut |s| {
for tok in s.tokens {
by_token.entry(tok).or_default().push(sid);
}
for (f, v) in s.numbers {
numbers_raw.push((sid, f, v));
}
rows.push(s.display);
sid += 1;
})?;
let mut ix = InfonIndex::from_postings(by_token, sid);
for (sid, f, v) in numbers_raw {
ix.add_number(sid, &f, v);
}
Ok(Corpus { ix, rows, columns, source, kind, noise: OnceLock::new(), gaz: OnceLock::new(), gaz_overlay: None })
}
pub fn from_csv(path: &Path) -> std::io::Result<Corpus> {
Corpus::from_projector(Box::new(CsvProjector::open(path)?))
}
pub fn from_json(path: &Path) -> std::io::Result<Corpus> {
Corpus::from_projector(Box::new(JsonProjector::open(path)))
}
pub fn from_jsonl(path: &Path, max_lines: Option<usize>) -> std::io::Result<Corpus> {
Corpus::from_projector(Box::new(JsonlProjector::open(path, max_lines)))
}
pub fn from_folder(dir: &Path) -> std::io::Result<(Corpus, FolderReport)> {
let mut files = collect_files(dir);
files.sort();
let mut by_token: HashMap<String, Vec<u32>> = HashMap::new();
let mut rows: Vec<Vec<String>> = Vec::new();
let mut numbers_raw: Vec<(u32, String, f64)> = Vec::new();
let mut sid: u32 = 0;
let mut report = FolderReport::default();
#[cfg(feature = "onnx")]
let mut text_engine = default_text_engine();
for path in &files {
let rel = path.strip_prefix(dir).unwrap_or(path).to_string_lossy().to_string();
let src_tok = format!("src/{}", crate::projector::slug(&rel));
let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("").to_lowercase();
let push = |s: crate::projector::Situation, by_token: &mut HashMap<String, Vec<u32>>, rows: &mut Vec<Vec<String>>, numbers_raw: &mut Vec<(u32, String, f64)>, sid: &mut u32| {
for t in s.tokens {
by_token.entry(t).or_default().push(*sid);
}
for (f, v) in s.numbers {
numbers_raw.push((*sid, f, v));
}
by_token.entry(src_tok.clone()).or_default().push(*sid);
rows.push(vec![rel.clone(), s.display.join(" · ")]);
*sid += 1;
};
#[cfg(feature = "docs")]
if crate::docs::is_doc_ext(&ext) {
if let Some(eng) = text_engine.as_mut() {
match crate::docs::extract_text(path) {
Ok(Some(text)) => {
let before = sid;
eng.project_text(&text, &mut |s| push(s, &mut by_token, &mut rows, &mut numbers_raw, &mut sid));
report.ingested.push((rel.clone(), (sid - before) as usize));
}
Ok(None) => {}
Err(e) => report.skipped.push((rel.clone(), format!("extract failed: {e}"))),
}
} else {
report.skipped.push((rel.clone(), "text models unavailable".into()));
}
continue;
}
let projector: Option<Box<dyn Projector>> = match ext.as_str() {
"csv" | "tsv" => CsvProjector::open(path).ok().map(|p| Box::new(p) as Box<dyn Projector>),
"json" | "ndjson" | "jsonl" => Some(Box::new(JsonProjector::open(path))),
"txt" | "md" | "text" => {
#[cfg(feature = "onnx")]
{
if let Some(eng) = text_engine.as_mut() {
if let Ok(text) = std::fs::read_to_string(path) {
let before = sid;
eng.project_text(&text, &mut |s| push(s, &mut by_token, &mut rows, &mut numbers_raw, &mut sid));
report.ingested.push((rel.clone(), (sid - before) as usize));
}
} else {
report.skipped.push((rel.clone(), "text models unavailable".into()));
}
None
}
#[cfg(not(feature = "onnx"))]
{
report.skipped.push((rel.clone(), "built without onnx feature".into()));
None
}
}
other if other.is_empty() => None,
other => {
report.skipped.push((rel.clone(), format!("no projector for .{other}")));
None
}
};
if let Some(p) = projector {
let before = sid;
let _ = p.project(&mut |s| push(s, &mut by_token, &mut rows, &mut numbers_raw, &mut sid));
report.ingested.push((rel.clone(), (sid - before) as usize));
}
}
let mut ix = InfonIndex::from_postings(by_token, sid);
for (sid, f, v) in numbers_raw {
ix.add_number(sid, &f, v);
}
report.situations = sid;
Ok((
Corpus {
ix,
rows,
columns: vec!["file".into(), "record".into()],
source: dir.display().to_string(),
kind: CorpusKind::Csv,
noise: OnceLock::new(),
gaz: OnceLock::new(), gaz_overlay: None,
},
report,
))
}
#[cfg(feature = "docs")]
pub fn from_document(path: &Path) -> std::io::Result<Corpus> {
let text = crate::docs::extract_text(path)?
.ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, "no extractable text"))?;
let mut eng = default_text_engine()
.ok_or_else(|| std::io::Error::new(std::io::ErrorKind::Other, "text models unavailable (set STEELDB_ML_BUNDLE)"))?;
let mut by_token: HashMap<String, Vec<u32>> = HashMap::new();
let mut rows: Vec<Vec<String>> = Vec::new();
let mut numbers_raw: Vec<(u32, String, f64)> = Vec::new();
let mut sid: u32 = 0;
eng.project_text(&text, &mut |s| {
for t in s.tokens {
by_token.entry(t).or_default().push(sid);
}
for (f, v) in s.numbers {
numbers_raw.push((sid, f, v));
}
rows.push(s.display);
sid += 1;
});
let mut ix = InfonIndex::from_postings(by_token, sid);
for (sid, f, v) in numbers_raw {
ix.add_number(sid, &f, v);
}
Ok(Corpus {
ix,
rows,
columns: vec!["sentence".into()],
source: path.display().to_string(),
kind: CorpusKind::Text,
noise: OnceLock::new(),
gaz: OnceLock::new(), gaz_overlay: None,
})
}
#[cfg(feature = "onnx")]
pub fn from_text(
path: &Path,
ml_bundle: &Path,
splade_dir: Option<&Path>,
) -> Result<Corpus, Box<dyn std::error::Error + Send + Sync>> {
let p = crate::projectors::TextProjector::open(path, ml_bundle, splade_dir)?;
Ok(Corpus::from_projector(Box::new(p))?)
}
pub fn new_incremental(source: impl Into<String>, columns: Vec<String>, kind: CorpusKind) -> Corpus {
Corpus {
ix: InfonIndex::from_postings(HashMap::new(), 0),
rows: Vec::new(),
columns,
source: source.into(),
kind,
noise: OnceLock::new(),
gaz: OnceLock::new(), gaz_overlay: None,
}
}
pub fn set_gazetteer_overlay(&mut self, path: impl Into<PathBuf>) {
self.gaz_overlay = Some(path.into());
}
pub fn add_situation(&mut self, tokens: Vec<String>, display: Vec<String>) -> u32 {
self.add_situation_num(tokens, display, Vec::new())
}
pub fn add_situation_num(&mut self, tokens: Vec<String>, display: Vec<String>, numbers: Vec<(String, f64)>) -> u32 {
self.add_situation_polar(tokens, display, numbers, Vec::new())
}
pub fn add_situation_polar(
&mut self,
mut tokens: Vec<String>,
display: Vec<String>,
numbers: Vec<(String, f64)>,
beliefs: Vec<(String, f32)>,
) -> u32 {
tokens.sort();
tokens.dedup();
let sid = self.ix.add(&tokens);
for (tok, level) in &beliefs {
self.ix.add_infon_polar(sid, tok, *level);
}
for (f, v) in numbers {
self.ix.add_number(sid, &f, v);
}
self.rows.push(display);
self.noise.take(); sid
}
pub fn index(&self) -> &InfonIndex<RoarPostings> {
&self.ix
}
pub fn query(&self, ikl: &str, limit: usize) -> QueryOut {
#[cfg(not(target_arch = "wasm32"))]
let t = Instant::now();
let result = evaluate(&self.ix, ikl);
#[cfg(not(target_arch = "wasm32"))]
let micros = t.elapsed().as_secs_f64() * 1e6;
#[cfg(target_arch = "wasm32")]
let micros = 0.0;
let sids = result.to_sorted();
let hits = sids
.iter()
.take(limit)
.map(|&sid| Hit {
sid,
cells: self.rows.get(sid as usize).cloned().unwrap_or_default(),
})
.collect();
QueryOut { count: sids.len(), micros, columns: self.columns.clone(), hits }
}
pub fn facet_tokens(&self, facet: &str, limit: usize) -> Vec<(String, usize)> {
self.ix.tokens_in_facet(facet, limit)
}
pub fn facet_names(&self) -> Vec<String> {
let mut set: HashSet<&str> = HashSet::new();
for t in self.ix.tokens() {
set.insert(t.split('/').next().unwrap_or(t));
}
let mut v: Vec<String> = set.into_iter().map(String::from).collect();
v.sort();
v
}
pub fn has_token(&self, token: &str) -> bool {
self.ix.post_len(token) > 0
}
pub fn belief_interval(&self, token: &str) -> (f64, f64) {
let universe = crate::tokenql::TokenStore::universe(&self.ix);
self.ix.belief_interval(token, &universe)
}
pub fn signed_mass(&self, token: &str) -> f64 {
let universe = crate::tokenql::TokenStore::universe(&self.ix);
self.ix.signed_mass(token, &universe)
}
pub fn linter(&self) -> crate::linter::Linter {
crate::linter::Linter::from_tokens(self.ix.tokens().cloned())
.with_numeric_fields(self.ix.numeric_fields().cloned())
}
pub fn top_token_leaves(&self, limit: usize) -> Vec<String> {
let mut v: Vec<(&String, usize)> = self.ix.tokens().map(|t| (t, self.ix.post_len(t))).collect();
v.sort_by(|a, b| b.1.cmp(&a.1));
let mut seen = HashSet::new();
let mut out = Vec::new();
for (t, _) in v {
let leaf = t.split('/').nth(1).unwrap_or(t).to_string();
if leaf.len() >= 3 && seen.insert(leaf.clone()) {
out.push(leaf);
if out.len() >= limit {
break;
}
}
}
out
}
pub fn noise_tokens(&self) -> &HashSet<String> {
self.noise.get_or_init(|| {
let thresh = 0.12 * self.ix.situations() as f64;
self.ix
.tokens()
.filter(|t| {
let f = t.split('/').next().unwrap_or("");
!matches!(f, "time" | "geo" | "qty" | "dur") && self.ix.post_len(t) as f64 > thresh
})
.cloned()
.collect()
})
}
pub fn search_ranked(&self, tokens: &[String], limit: usize) -> Vec<(u32, usize, Vec<String>)> {
if tokens.is_empty() {
return Vec::new();
}
let posts: Vec<RoarPostings> = tokens.iter().map(|t| self.ix.post(t)).collect();
let mut union = RoarPostings::empty();
for p in &posts {
union.or_inplace(p);
}
let mut scored: Vec<(u32, usize)> = union
.to_sorted()
.into_iter()
.map(|sid| (sid, posts.iter().filter(|p| p.contains(sid)).count()))
.collect();
scored.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0)));
scored.truncate(limit);
scored
.into_iter()
.map(|(sid, cov)| (sid, cov, self.rows.get(sid as usize).cloned().unwrap_or_default()))
.collect()
}
fn gazetteer(&self) -> Option<&Gazetteer> {
self.gaz
.get_or_init(|| {
let p = std::env::var("STEELDB_GAZETTEER")
.map(PathBuf::from)
.ok()
.filter(|p| p.exists())
.or_else(|| {
crate::paths::model_dir("splade", "STEELDB_SPLADE_DIR", "gazetteer.json")
.map(|d| d.join("gazetteer.json"))
});
let mut g = p.and_then(|p| Gazetteer::load(&p).ok());
if let Some(ov) = &self.gaz_overlay {
if ov.exists() {
let mut base = g.take().unwrap_or_else(Gazetteer::empty);
base.merge_overlay(ov);
g = Some(base);
}
}
g
})
.as_ref()
}
pub fn entity_link(&self, question: &str) -> Vec<String> {
let noise = self.noise_tokens();
let vocab: HashSet<&str> = self.ix.tokens().map(|s| s.as_str()).collect();
let mut leaf_idx: HashMap<String, Vec<String>> = HashMap::new();
let put = |k: &str, t: &str, idx: &mut HashMap<String, Vec<String>>| {
if k.len() >= 3 {
let e = idx.entry(k.to_string()).or_default();
if !e.iter().any(|x| x == t) {
e.push(t.to_string());
}
}
};
for t in self.ix.tokens() {
let leaf = match t.find('/') {
Some(i) => &t[i + 1..],
None => t.as_str(),
};
put(leaf, t, &mut leaf_idx);
for part in leaf.split(['-', '/']) {
put(part, t, &mut leaf_idx);
}
}
let mut out: Vec<String> = Vec::new();
let mut seen: HashSet<String> = HashSet::new();
let add = |key: &str, out: &mut Vec<String>, seen: &mut HashSet<String>| {
if let Some(hits) = leaf_idx.get(key) {
for t in hits.iter().take(5) {
if !noise.contains(t) && seen.insert(t.clone()) {
out.push(t.clone());
}
}
}
};
for w in question.split(|c: char| !c.is_alphanumeric()) {
if (2..=6).contains(&w.chars().count()) && w.chars().all(|c| c.is_ascii_uppercase()) {
add(&w.to_lowercase(), &mut out, &mut seen);
}
}
let words: Vec<String> = question
.to_lowercase()
.split(|c: char| !c.is_alphanumeric())
.filter(|w| w.len() >= 3 && !STOPW.contains(w))
.map(|w| w.to_string())
.collect();
for w in &words {
add(w, &mut out, &mut seen);
if w.len() >= 5 {
let mut prefix_hits: Vec<String> = Vec::new();
for (leaf, toks) in &leaf_idx {
if leaf.len() >= 4 && w.starts_with(leaf.as_str()) {
for t in toks.iter().take(3) {
prefix_hits.push(t.clone());
}
}
}
for t in prefix_hits {
if !noise.contains(&t) && seen.insert(t.clone()) {
out.push(t);
}
}
}
}
for pair in words.windows(2) {
add(&format!("{}-{}", pair[0], pair[1]), &mut out, &mut seen);
}
if let Some(gaz) = self.gazetteer() {
for h in gaz.extract(question) {
if vocab.contains(h.token.as_str()) && !noise.contains(&h.token) && seen.insert(h.token.clone()) {
out.push(h.token);
}
}
}
out.truncate(50);
out
}
pub fn breakdown(&self, anchor: &str, facet: &str, k: usize) -> Value {
programs::breakdown(&self.ix, anchor, facet, k)
}
pub fn crosstab(&self, anchor: &str, facet_a: &str, facet_b: &str, k: usize) -> Value {
programs::crosstab(&self.ix, anchor, facet_a, facet_b, k)
}
pub fn rank(&self, facet: &str, k: usize) -> Value {
programs::rank(&self.ix, facet, k, self.noise_tokens())
}
pub fn cooccurs(&self, token: &str, k: usize) -> Value {
programs::cooccurs(&self.ix, token, k, self.noise_tokens())
}
pub fn s_path(&self, a: &str, b: &str, s: usize) -> Value {
programs::s_path(&self.ix, a, b, s, self.noise_tokens())
}
pub fn s_clusters(&self, s: usize, k: usize) -> Value {
programs::s_clusters(&self.ix, s, k, self.noise_tokens())
}
pub fn narrow(&self, scope: &[String], filters: &[String]) -> Value {
programs::narrow(&self.ix, scope, filters)
}
pub fn stats(&self) -> Stats {
let mut facet_tokens: HashMap<String, usize> = HashMap::new();
for tok in self.ix.tokens() {
let facet = tok.split('/').next().unwrap_or(tok).to_string();
*facet_tokens.entry(facet).or_default() += 1;
}
let mut facets: Vec<(String, usize)> = facet_tokens.into_iter().collect();
facets.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0)));
let mut numeric_fields: Vec<String> = self.ix.numeric_fields().cloned().collect();
numeric_fields.sort();
Stats {
source: self.source.clone(),
kind: self.kind,
situations: self.ix.situations(),
vocab: self.ix.vocab_size(),
columns: self.columns.clone(),
facets,
numeric_fields,
}
}
}