1use crate::bitmap::{Postings, RoarPostings};
8use crate::index::InfonIndex;
9use crate::programs;
10use crate::projector::{CorpusKind, Projector};
11use crate::projectors::{CsvProjector, JsonProjector, JsonlProjector};
12use crate::text::Gazetteer;
13use crate::tokenql::evaluate;
14use serde::Serialize;
15use serde_json::Value;
16use std::collections::{HashMap, HashSet};
17use std::path::{Path, PathBuf};
18use std::sync::OnceLock;
19#[cfg(not(target_arch = "wasm32"))]
23use std::time::Instant;
24
25const STOPW: &[&str] = &[
27 "the", "a", "an", "of", "to", "in", "on", "for", "and", "or", "is", "are", "was", "were", "be",
28 "do", "how", "what", "which", "who", "why", "when", "with", "without", "into", "over", "under",
29 "from", "your", "our", "their", "this", "that", "these", "those", "can", "may", "will", "would",
30 "should", "could", "not", "you", "use", "used", "using", "best", "common", "about", "across",
31 "based", "provide", "provides", "support", "supports", "need", "needs", "want", "wants", "able",
32];
33
34#[derive(Serialize, Default, Debug)]
35pub struct FolderReport {
36 pub situations: u32,
37 pub ingested: Vec<(String, usize)>,
39 pub skipped: Vec<(String, String)>,
41}
42
43fn collect_files(dir: &Path) -> Vec<PathBuf> {
45 let mut out = Vec::new();
46 let mut stack = vec![dir.to_path_buf()];
47 while let Some(d) = stack.pop() {
48 let rd = match std::fs::read_dir(&d) {
49 Ok(r) => r,
50 Err(_) => continue,
51 };
52 for e in rd.flatten() {
53 let name = e.file_name().to_string_lossy().to_string();
54 if name.starts_with('.') || name == "node_modules" || name == "target" {
55 continue;
56 }
57 let p = e.path();
58 if p.is_dir() {
59 stack.push(p);
60 } else {
61 out.push(p);
62 }
63 }
64 }
65 out
66}
67
68#[cfg(feature = "onnx")]
71fn default_text_engine() -> Option<crate::projectors::TextEngine> {
72 let ml = crate::paths::model_dir("step0_bundle_ml", "STEELDB_ML_BUNDLE", "spo.onnx")?;
76 let splade = crate::paths::model_dir("splade", "STEELDB_SPLADE_DIR", "splade.onnx");
77 crate::projectors::TextEngine::load(&ml, splade.as_deref()).ok()
78}
79
80pub struct Corpus {
81 ix: InfonIndex<RoarPostings>,
82 rows: Vec<Vec<String>>,
83 columns: Vec<String>,
84 source: String,
85 kind: CorpusKind,
86 noise: OnceLock<HashSet<String>>,
88 gaz: OnceLock<Option<Gazetteer>>,
90 gaz_overlay: Option<PathBuf>,
92}
93
94#[derive(Serialize)]
95pub struct Hit {
96 pub sid: u32,
97 pub cells: Vec<String>,
98}
99
100#[derive(Serialize)]
101pub struct QueryOut {
102 pub count: usize,
103 pub micros: f64,
104 pub columns: Vec<String>,
105 pub hits: Vec<Hit>,
106}
107
108#[derive(Serialize)]
109pub struct Stats {
110 pub source: String,
111 pub kind: CorpusKind,
112 pub situations: u32,
113 pub vocab: usize,
114 pub columns: Vec<String>,
115 pub facets: Vec<(String, usize)>,
117 pub numeric_fields: Vec<String>,
119}
120
121impl Corpus {
122 pub fn from_projector(p: Box<dyn Projector>) -> std::io::Result<Corpus> {
124 let columns = p.columns();
125 let kind = p.kind();
126 let source = p.source();
127 let mut by_token: HashMap<String, Vec<u32>> = HashMap::new();
128 let mut rows: Vec<Vec<String>> = Vec::new();
129 let mut numbers_raw: Vec<(u32, String, f64)> = Vec::new();
130 let mut sid: u32 = 0;
131 p.project(&mut |s| {
132 for tok in s.tokens {
133 by_token.entry(tok).or_default().push(sid);
134 }
135 for (f, v) in s.numbers {
136 numbers_raw.push((sid, f, v));
137 }
138 rows.push(s.display);
139 sid += 1;
140 })?;
141 let mut ix = InfonIndex::from_postings(by_token, sid);
143 for (sid, f, v) in numbers_raw {
144 ix.add_number(sid, &f, v);
145 }
146 Ok(Corpus { ix, rows, columns, source, kind, noise: OnceLock::new(), gaz: OnceLock::new(), gaz_overlay: None })
147 }
148
149 pub fn from_csv(path: &Path) -> std::io::Result<Corpus> {
150 Corpus::from_projector(Box::new(CsvProjector::open(path)?))
151 }
152 pub fn from_json(path: &Path) -> std::io::Result<Corpus> {
153 Corpus::from_projector(Box::new(JsonProjector::open(path)))
154 }
155 pub fn from_jsonl(path: &Path, max_lines: Option<usize>) -> std::io::Result<Corpus> {
156 Corpus::from_projector(Box::new(JsonlProjector::open(path, max_lines)))
157 }
158
159 pub fn from_folder(dir: &Path) -> std::io::Result<(Corpus, FolderReport)> {
164 let mut files = collect_files(dir);
165 files.sort();
166 let mut by_token: HashMap<String, Vec<u32>> = HashMap::new();
167 let mut rows: Vec<Vec<String>> = Vec::new();
168 let mut numbers_raw: Vec<(u32, String, f64)> = Vec::new();
169 let mut sid: u32 = 0;
170 let mut report = FolderReport::default();
171
172 #[cfg(feature = "onnx")]
174 let mut text_engine = default_text_engine();
175
176 for path in &files {
177 let rel = path.strip_prefix(dir).unwrap_or(path).to_string_lossy().to_string();
178 let src_tok = format!("src/{}", crate::projector::slug(&rel));
179 let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("").to_lowercase();
180
181 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| {
182 for t in s.tokens {
183 by_token.entry(t).or_default().push(*sid);
184 }
185 for (f, v) in s.numbers {
186 numbers_raw.push((*sid, f, v));
187 }
188 by_token.entry(src_tok.clone()).or_default().push(*sid);
189 rows.push(vec![rel.clone(), s.display.join(" · ")]);
190 *sid += 1;
191 };
192
193 #[cfg(feature = "docs")]
195 if crate::docs::is_doc_ext(&ext) {
196 if let Some(eng) = text_engine.as_mut() {
197 match crate::docs::extract_text(path) {
198 Ok(Some(text)) => {
199 let before = sid;
200 eng.project_text(&text, &mut |s| push(s, &mut by_token, &mut rows, &mut numbers_raw, &mut sid));
201 report.ingested.push((rel.clone(), (sid - before) as usize));
202 }
203 Ok(None) => {}
204 Err(e) => report.skipped.push((rel.clone(), format!("extract failed: {e}"))),
205 }
206 } else {
207 report.skipped.push((rel.clone(), "text models unavailable".into()));
208 }
209 continue;
210 }
211
212 let projector: Option<Box<dyn Projector>> = match ext.as_str() {
213 "csv" | "tsv" => CsvProjector::open(path).ok().map(|p| Box::new(p) as Box<dyn Projector>),
214 "json" | "ndjson" | "jsonl" => Some(Box::new(JsonProjector::open(path))),
215 "txt" | "md" | "text" => {
216 #[cfg(feature = "onnx")]
217 {
218 if let Some(eng) = text_engine.as_mut() {
219 if let Ok(text) = std::fs::read_to_string(path) {
220 let before = sid;
221 eng.project_text(&text, &mut |s| push(s, &mut by_token, &mut rows, &mut numbers_raw, &mut sid));
222 report.ingested.push((rel.clone(), (sid - before) as usize));
223 }
224 } else {
225 report.skipped.push((rel.clone(), "text models unavailable".into()));
226 }
227 None
228 }
229 #[cfg(not(feature = "onnx"))]
230 {
231 report.skipped.push((rel.clone(), "built without onnx feature".into()));
232 None
233 }
234 }
235 other if other.is_empty() => None,
236 other => {
237 report.skipped.push((rel.clone(), format!("no projector for .{other}")));
238 None
239 }
240 };
241
242 if let Some(p) = projector {
243 let before = sid;
244 let _ = p.project(&mut |s| push(s, &mut by_token, &mut rows, &mut numbers_raw, &mut sid));
245 report.ingested.push((rel.clone(), (sid - before) as usize));
246 }
247 }
248
249 let mut ix = InfonIndex::from_postings(by_token, sid);
250 for (sid, f, v) in numbers_raw {
251 ix.add_number(sid, &f, v);
252 }
253 report.situations = sid;
254 Ok((
255 Corpus {
256 ix,
257 rows,
258 columns: vec!["file".into(), "record".into()],
259 source: dir.display().to_string(),
260 kind: CorpusKind::Csv,
261 noise: OnceLock::new(),
262 gaz: OnceLock::new(), gaz_overlay: None,
263 },
264 report,
265 ))
266 }
267
268 #[cfg(feature = "docs")]
271 pub fn from_document(path: &Path) -> std::io::Result<Corpus> {
272 let text = crate::docs::extract_text(path)?
273 .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, "no extractable text"))?;
274 let mut eng = default_text_engine()
275 .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::Other, "text models unavailable (set STEELDB_ML_BUNDLE)"))?;
276 let mut by_token: HashMap<String, Vec<u32>> = HashMap::new();
277 let mut rows: Vec<Vec<String>> = Vec::new();
278 let mut numbers_raw: Vec<(u32, String, f64)> = Vec::new();
279 let mut sid: u32 = 0;
280 eng.project_text(&text, &mut |s| {
281 for t in s.tokens {
282 by_token.entry(t).or_default().push(sid);
283 }
284 for (f, v) in s.numbers {
285 numbers_raw.push((sid, f, v));
286 }
287 rows.push(s.display);
288 sid += 1;
289 });
290 let mut ix = InfonIndex::from_postings(by_token, sid);
291 for (sid, f, v) in numbers_raw {
292 ix.add_number(sid, &f, v);
293 }
294 Ok(Corpus {
295 ix,
296 rows,
297 columns: vec!["sentence".into()],
298 source: path.display().to_string(),
299 kind: CorpusKind::Text,
300 noise: OnceLock::new(),
301 gaz: OnceLock::new(), gaz_overlay: None,
302 })
303 }
304
305 #[cfg(feature = "onnx")]
307 pub fn from_text(
308 path: &Path,
309 ml_bundle: &Path,
310 splade_dir: Option<&Path>,
311 ) -> Result<Corpus, Box<dyn std::error::Error + Send + Sync>> {
312 let p = crate::projectors::TextProjector::open(path, ml_bundle, splade_dir)?;
313 Ok(Corpus::from_projector(Box::new(p))?)
314 }
315
316 pub fn new_incremental(source: impl Into<String>, columns: Vec<String>, kind: CorpusKind) -> Corpus {
319 Corpus {
320 ix: InfonIndex::from_postings(HashMap::new(), 0),
321 rows: Vec::new(),
322 columns,
323 source: source.into(),
324 kind,
325 noise: OnceLock::new(),
326 gaz: OnceLock::new(), gaz_overlay: None,
327 }
328 }
329
330 pub fn set_gazetteer_overlay(&mut self, path: impl Into<PathBuf>) {
333 self.gaz_overlay = Some(path.into());
334 }
335
336 pub fn add_situation(&mut self, tokens: Vec<String>, display: Vec<String>) -> u32 {
339 self.add_situation_num(tokens, display, Vec::new())
340 }
341
342 pub fn add_situation_num(&mut self, tokens: Vec<String>, display: Vec<String>, numbers: Vec<(String, f64)>) -> u32 {
344 self.add_situation_polar(tokens, display, numbers, Vec::new())
345 }
346
347 pub fn add_situation_polar(
351 &mut self,
352 mut tokens: Vec<String>,
353 display: Vec<String>,
354 numbers: Vec<(String, f64)>,
355 beliefs: Vec<(String, f32)>,
356 ) -> u32 {
357 tokens.sort();
358 tokens.dedup();
359 let sid = self.ix.add(&tokens);
360 for (tok, level) in &beliefs {
361 self.ix.add_infon_polar(sid, tok, *level);
362 }
363 for (f, v) in numbers {
364 self.ix.add_number(sid, &f, v);
365 }
366 self.rows.push(display);
367 self.noise.take(); sid
369 }
370
371 pub fn index(&self) -> &InfonIndex<RoarPostings> {
374 &self.ix
375 }
376
377 pub fn query(&self, ikl: &str, limit: usize) -> QueryOut {
378 #[cfg(not(target_arch = "wasm32"))]
379 let t = Instant::now();
380 let result = evaluate(&self.ix, ikl);
381 #[cfg(not(target_arch = "wasm32"))]
382 let micros = t.elapsed().as_secs_f64() * 1e6;
383 #[cfg(target_arch = "wasm32")]
384 let micros = 0.0;
385 let sids = result.to_sorted();
386 let hits = sids
387 .iter()
388 .take(limit)
389 .map(|&sid| Hit {
390 sid,
391 cells: self.rows.get(sid as usize).cloned().unwrap_or_default(),
392 })
393 .collect();
394 QueryOut { count: sids.len(), micros, columns: self.columns.clone(), hits }
395 }
396
397 pub fn facet_tokens(&self, facet: &str, limit: usize) -> Vec<(String, usize)> {
399 self.ix.tokens_in_facet(facet, limit)
400 }
401
402 pub fn facet_names(&self) -> Vec<String> {
405 let mut set: HashSet<&str> = HashSet::new();
406 for t in self.ix.tokens() {
407 set.insert(t.split('/').next().unwrap_or(t));
408 }
409 let mut v: Vec<String> = set.into_iter().map(String::from).collect();
410 v.sort();
411 v
412 }
413
414 pub fn has_token(&self, token: &str) -> bool {
416 self.ix.post_len(token) > 0
417 }
418
419 pub fn belief_interval(&self, token: &str) -> (f64, f64) {
421 let universe = crate::tokenql::TokenStore::universe(&self.ix);
422 self.ix.belief_interval(token, &universe)
423 }
424
425 pub fn signed_mass(&self, token: &str) -> f64 {
427 let universe = crate::tokenql::TokenStore::universe(&self.ix);
428 self.ix.signed_mass(token, &universe)
429 }
430
431 pub fn linter(&self) -> crate::linter::Linter {
434 crate::linter::Linter::from_tokens(self.ix.tokens().cloned())
435 .with_numeric_fields(self.ix.numeric_fields().cloned())
436 }
437
438 pub fn top_token_leaves(&self, limit: usize) -> Vec<String> {
441 let mut v: Vec<(&String, usize)> = self.ix.tokens().map(|t| (t, self.ix.post_len(t))).collect();
442 v.sort_by(|a, b| b.1.cmp(&a.1));
443 let mut seen = HashSet::new();
444 let mut out = Vec::new();
445 for (t, _) in v {
446 let leaf = t.split('/').nth(1).unwrap_or(t).to_string();
447 if leaf.len() >= 3 && seen.insert(leaf.clone()) {
448 out.push(leaf);
449 if out.len() >= limit {
450 break;
451 }
452 }
453 }
454 out
455 }
456
457 pub fn noise_tokens(&self) -> &HashSet<String> {
462 self.noise.get_or_init(|| {
463 let thresh = 0.12 * self.ix.situations() as f64;
464 self.ix
465 .tokens()
466 .filter(|t| {
467 let f = t.split('/').next().unwrap_or("");
468 !matches!(f, "time" | "geo" | "qty" | "dur") && self.ix.post_len(t) as f64 > thresh
469 })
470 .cloned()
471 .collect()
472 })
473 }
474
475 pub fn search_ranked(&self, tokens: &[String], limit: usize) -> Vec<(u32, usize, Vec<String>)> {
479 if tokens.is_empty() {
480 return Vec::new();
481 }
482 let posts: Vec<RoarPostings> = tokens.iter().map(|t| self.ix.post(t)).collect();
483 let mut union = RoarPostings::empty();
484 for p in &posts {
485 union.or_inplace(p);
486 }
487 let mut scored: Vec<(u32, usize)> = union
488 .to_sorted()
489 .into_iter()
490 .map(|sid| (sid, posts.iter().filter(|p| p.contains(sid)).count()))
491 .collect();
492 scored.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0)));
493 scored.truncate(limit);
494 scored
495 .into_iter()
496 .map(|(sid, cov)| (sid, cov, self.rows.get(sid as usize).cloned().unwrap_or_default()))
497 .collect()
498 }
499
500 fn gazetteer(&self) -> Option<&Gazetteer> {
502 self.gaz
503 .get_or_init(|| {
504 let p = std::env::var("STEELDB_GAZETTEER")
505 .map(PathBuf::from)
506 .ok()
507 .filter(|p| p.exists())
508 .or_else(|| {
509 crate::paths::model_dir("splade", "STEELDB_SPLADE_DIR", "gazetteer.json")
510 .map(|d| d.join("gazetteer.json"))
511 });
512 let mut g = p.and_then(|p| Gazetteer::load(&p).ok());
513 if let Some(ov) = &self.gaz_overlay {
516 if ov.exists() {
517 let mut base = g.take().unwrap_or_else(Gazetteer::empty);
518 base.merge_overlay(ov);
519 g = Some(base);
520 }
521 }
522 g
523 })
524 .as_ref()
525 }
526
527 pub fn entity_link(&self, question: &str) -> Vec<String> {
532 let noise = self.noise_tokens();
533 let vocab: HashSet<&str> = self.ix.tokens().map(|s| s.as_str()).collect();
534
535 let mut leaf_idx: HashMap<String, Vec<String>> = HashMap::new();
537 let put = |k: &str, t: &str, idx: &mut HashMap<String, Vec<String>>| {
538 if k.len() >= 3 {
539 let e = idx.entry(k.to_string()).or_default();
540 if !e.iter().any(|x| x == t) {
541 e.push(t.to_string());
542 }
543 }
544 };
545 for t in self.ix.tokens() {
546 let leaf = match t.find('/') {
547 Some(i) => &t[i + 1..],
548 None => t.as_str(),
549 };
550 put(leaf, t, &mut leaf_idx);
551 for part in leaf.split(['-', '/']) {
552 put(part, t, &mut leaf_idx);
553 }
554 }
555
556 let mut out: Vec<String> = Vec::new();
557 let mut seen: HashSet<String> = HashSet::new();
558 let add = |key: &str, out: &mut Vec<String>, seen: &mut HashSet<String>| {
559 if let Some(hits) = leaf_idx.get(key) {
560 for t in hits.iter().take(5) {
561 if !noise.contains(t) && seen.insert(t.clone()) {
562 out.push(t.clone());
563 }
564 }
565 }
566 };
567
568 for w in question.split(|c: char| !c.is_alphanumeric()) {
570 if (2..=6).contains(&w.chars().count()) && w.chars().all(|c| c.is_ascii_uppercase()) {
571 add(&w.to_lowercase(), &mut out, &mut seen);
572 }
573 }
574 let words: Vec<String> = question
575 .to_lowercase()
576 .split(|c: char| !c.is_alphanumeric())
577 .filter(|w| w.len() >= 3 && !STOPW.contains(w))
578 .map(|w| w.to_string())
579 .collect();
580 for w in &words {
581 add(w, &mut out, &mut seen);
582 if w.len() >= 5 {
584 let mut prefix_hits: Vec<String> = Vec::new();
585 for (leaf, toks) in &leaf_idx {
586 if leaf.len() >= 4 && w.starts_with(leaf.as_str()) {
587 for t in toks.iter().take(3) {
588 prefix_hits.push(t.clone());
589 }
590 }
591 }
592 for t in prefix_hits {
593 if !noise.contains(&t) && seen.insert(t.clone()) {
594 out.push(t);
595 }
596 }
597 }
598 }
599 for pair in words.windows(2) {
601 add(&format!("{}-{}", pair[0], pair[1]), &mut out, &mut seen);
602 }
603 if let Some(gaz) = self.gazetteer() {
605 for h in gaz.extract(question) {
606 if vocab.contains(h.token.as_str()) && !noise.contains(&h.token) && seen.insert(h.token.clone()) {
607 out.push(h.token);
608 }
609 }
610 }
611 out.truncate(50);
612 out
613 }
614
615 pub fn breakdown(&self, anchor: &str, facet: &str, k: usize) -> Value {
617 programs::breakdown(&self.ix, anchor, facet, k)
618 }
619 pub fn crosstab(&self, anchor: &str, facet_a: &str, facet_b: &str, k: usize) -> Value {
620 programs::crosstab(&self.ix, anchor, facet_a, facet_b, k)
621 }
622 pub fn rank(&self, facet: &str, k: usize) -> Value {
623 programs::rank(&self.ix, facet, k, self.noise_tokens())
624 }
625 pub fn cooccurs(&self, token: &str, k: usize) -> Value {
626 programs::cooccurs(&self.ix, token, k, self.noise_tokens())
627 }
628 pub fn s_path(&self, a: &str, b: &str, s: usize) -> Value {
629 programs::s_path(&self.ix, a, b, s, self.noise_tokens())
630 }
631 pub fn s_clusters(&self, s: usize, k: usize) -> Value {
632 programs::s_clusters(&self.ix, s, k, self.noise_tokens())
633 }
634 pub fn narrow(&self, scope: &[String], filters: &[String]) -> Value {
635 programs::narrow(&self.ix, scope, filters)
636 }
637
638 pub fn stats(&self) -> Stats {
639 let mut facet_tokens: HashMap<String, usize> = HashMap::new();
640 for tok in self.ix.tokens() {
641 let facet = tok.split('/').next().unwrap_or(tok).to_string();
642 *facet_tokens.entry(facet).or_default() += 1;
643 }
644 let mut facets: Vec<(String, usize)> = facet_tokens.into_iter().collect();
645 facets.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0)));
646 let mut numeric_fields: Vec<String> = self.ix.numeric_fields().cloned().collect();
647 numeric_fields.sort();
648 Stats {
649 source: self.source.clone(),
650 kind: self.kind,
651 situations: self.ix.situations(),
652 vocab: self.ix.vocab_size(),
653 columns: self.columns.clone(),
654 facets,
655 numeric_fields,
656 }
657 }
658}