Skip to main content

steeldb/
jsonl.rs

1//! Streaming JSONL loader. Each line is one situation `{ "tokens": [...], "num": {...} }`; the line
2//! index (0-based) is the situation id. Builds `token -> ascending sids` in a single pass — the sids
3//! land in id order for free, so the index build hits the fast sorted path.
4
5use serde::Deserialize;
6use std::collections::HashMap;
7use std::fs::File;
8use std::io::{BufRead, BufReader};
9use std::path::Path;
10
11#[derive(Deserialize)]
12struct Row {
13    #[serde(default)]
14    tokens: Vec<String>,
15}
16
17/// Returns (token -> sids, n_situations).
18pub fn load(path: &Path, max_lines: Option<usize>) -> std::io::Result<(HashMap<String, Vec<u32>>, u32)> {
19    let file = File::open(path)?;
20    let reader = BufReader::with_capacity(1 << 20, file);
21    let mut by_token: HashMap<String, Vec<u32>> = HashMap::new();
22    let mut sid: u32 = 0;
23
24    for line in reader.lines() {
25        let line = line?;
26        if line.trim().is_empty() {
27            continue;
28        }
29        if let Some(max) = max_lines {
30            if sid as usize >= max {
31                break;
32            }
33        }
34        let row: Row = match serde_json::from_str(&line) {
35            Ok(r) => r,
36            Err(_) => {
37                sid += 1;
38                continue;
39            }
40        };
41        for tok in row.tokens {
42            by_token.entry(tok).or_default().push(sid);
43        }
44        sid += 1;
45    }
46    Ok((by_token, sid))
47}