use serde::Deserialize;
use std::collections::HashMap;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::path::Path;
#[derive(Deserialize)]
struct Row {
#[serde(default)]
tokens: Vec<String>,
}
pub fn load(path: &Path, max_lines: Option<usize>) -> std::io::Result<(HashMap<String, Vec<u32>>, u32)> {
let file = File::open(path)?;
let reader = BufReader::with_capacity(1 << 20, file);
let mut by_token: HashMap<String, Vec<u32>> = HashMap::new();
let mut sid: u32 = 0;
for line in reader.lines() {
let line = line?;
if line.trim().is_empty() {
continue;
}
if let Some(max) = max_lines {
if sid as usize >= max {
break;
}
}
let row: Row = match serde_json::from_str(&line) {
Ok(r) => r,
Err(_) => {
sid += 1;
continue;
}
};
for tok in row.tokens {
by_token.entry(tok).or_default().push(sid);
}
sid += 1;
}
Ok((by_token, sid))
}