engramdb_core/
count_index.rs

1//! I2 频率索引：rowid → count（排序键值对文件）与热集选择。
2//!
3//! 构建：流式 sum（计数原子性由调用方保证，去重由输入保证），内存 map → 排序 → 落盘。
4//! 400M 级表全量应用时换外部排序（M1.5）；当前语料规模（千万级唯一）内存版足够。
5
6use std::collections::HashMap;
7use std::io::{BufWriter, Read, Write};
8use std::path::Path;
9
10#[derive(Debug, Default)]
11pub struct CountIndex {
12    counts: HashMap<u64, u32>,
13}
14
15impl CountIndex {
16    /// 从二进制流构建：每行 u64 (nat endian) rowid 序列（重复即累加）。
17    pub fn build_from_bin_stream(mut src: impl Read) -> std::io::Result<Self> {
18        let mut counts = HashMap::<u64, u32>::new();
19        let mut buf = [0u8; 8];
20        loop {
21            let n = src.read(&mut buf)?;
22            if n == 0 {
23                break;
24            }
25            if n != 8 {
26                return Err(std::io::Error::new(
27                    std::io::ErrorKind::UnexpectedEof,
28                    "partial rowid",
29                ));
30            }
31            let rowid = u64::from_le_bytes(buf);
32            *counts.entry(rowid).or_insert(0) += 1;
33        }
34        Ok(Self { counts })
35    }
36
37    pub fn count(&self, rowid: u64) -> u64 {
38        self.counts.get(&rowid).copied().unwrap_or(0) as u64
39    }
40
41    pub fn iter(&self) -> impl Iterator<Item = (u64, u32)> {
42        let mut v: Vec<_> = self.counts.iter().map(|(&a, &b)| (a, b)).collect();
43        v.sort_unstable_by_key(|&(a, _)| a);
44        v.into_iter()
45    }
46
47    /// 落盘：binary (u64 rowid, u32 count) 排序流。
48    pub fn write_bin(&self, path: &Path) -> std::io::Result<()> {
49        let mut w = BufWriter::new(std::fs::File::create(path)?);
50        for (rowid, count) in self.iter() {
51            w.write_all(&rowid.to_le_bytes())?;
52            w.write_all(&count.to_le_bytes())?;
53        }
54        w.flush()
55    }
56
57    pub fn write_dump(&self, path: &Path) -> std::io::Result<()> {
58        let mut w = BufWriter::new(std::fs::File::create(path)?);
59        for (rowid, count) in self.iter() {
60            writeln!(w, "{} {}", rowid, count)?;
61        }
62        w.flush()
63    }
64
65    pub fn from_bin(path: &Path) -> std::io::Result<Self> {
66        let data = std::fs::read(path)?;
67        let mut counts = HashMap::new();
68        for chunk in data.chunks(12) {
69            if chunk.len() != 12 {
70                break;
71            }
72            let rowid = u64::from_le_bytes(chunk[..8].try_into().unwrap());
73            let c = u32::from_le_bytes(chunk[8..12].try_into().unwrap());
74            counts.insert(rowid, c);
75        }
76        Ok(Self { counts })
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn build_count_write_read() {
86        let data: Vec<u8> = [5u64, 3, 5, 7, 5]
87            .iter()
88            .flat_map(|x| x.to_le_bytes())
89            .collect();
90        let idx = CountIndex::build_from_bin_stream(&data[..]).unwrap();
91        assert_eq!(idx.count(5), 3);
92        assert_eq!(idx.count(3), 1);
93        let p = std::env::temp_dir().join("ct_test.bin");
94        idx.write_bin(&p).unwrap();
95        let idx2 = CountIndex::from_bin(&p).unwrap();
96        assert_eq!(idx2.count(7), 1);
97    }
98}
engramdb_core/count_index.rs

engramdb_core/
count_index.rs
