1use 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
17pub 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}