Rust 实现毫秒级百万词汇纠错系统

我们上线 Hacker News 搜索和 RAG 引擎时，使用的还是一套半成品拼写纠错系统。最初版本处理拼写正确的查询需要 30ms 以上，速度慢到我们不得不默认关闭它。最新版本的速度提升了 100 倍：拼写正确的查询只需 300μs，拼写错误的查询约为每个单词 5ms。本文将介绍我们是如何做到的！

可供尝试的查询示例

点击链接，即可在 hn.trieve.ai 上亲自体验这套拼写纠错系统。

Cnva devloper platfirm

创建单词及词频字典

对于小型数据集，这项任务很简单。使用一个 worker 和基础的单词切分逻辑，可以在 10 秒内滚动处理约 1,000 个 Hacker News 帖子大小的文本块。然而，当数据规模增长到 Hacker News Demo 的体量（超过 3,800 万篇帖子）时，就必须以分布式方式处理任务。

最终，我们决定使用两种不同的 worker 来构建字典：

• Cronjob：滚动遍历每位用户搜索索引中的所有文档，每次从数据库中取出 500 个 chunk id，并将其加入 Redis 队列。

• Word worker：从队列中取出任务，每次处理 500 个 chunk。它会拉取每个 chunk 的文本，将文本切分成单词，然后把每个单词写入 ClickHouse。

我们选择用 ClickHouse 存储字典，是因为随着 worker 数量增加，使用 Postgres 写入时遇到了死锁和性能问题。ClickHouse 的异步插入非常适合这项任务，让我们能够在不到 1 小时内完成整个 3,800 多万篇文档数据集的写入。

使用 BKTree 数据结构识别并纠正拼写错误

我们采用标准的拼写纠错方案，为每个数据集分别构建 Burkhard-Keller Tree（BKTree），从而以 O(log N) 的时间复杂度，高效比较搜索查询中的单词与数据集字典中的单词。深入讲解这种数据结构超出了本文的范围，不过你可以阅读我们的 Rust 实现或相关 wiki，了解更多信息。

我们使用了第三种 bktree-worker 来构建 BKTree。它会读取已经在 ClickHouse 中完成字典构建的数据集，并利用其中的单词及词频构建树。

BKTree 构建完成后，worker 会将它存入 Redis。这样，当某个数据集首次收到查询时，API server 就能按需将其高效加载到内存中。

对于较大的数据集，这一步颇具挑战：对应的树可能有数百 MB，读写时会导致 Redis 超时。我们开发了一套序列化方法，先将树扁平化，再使用 gzip 压缩，以减小它在 Redis 中占用的空间，同时降低从 Redis 拉取和推送数据时的延迟。

编写执行拼写纠错的业务逻辑

在 API server 端，我们优化了 typo_operator，将拼写正确查询的纠错耗时降低到约 300μs，将拼写错误查询的耗时降低到约每个单词 10ms。

从 Redis 拉取像 Hacker News BKTree 这样庞大的数据结构需要 300μs 以上，因此不可能在每次搜索时都执行。为此，我们使用 lazy_static! 在 server 端开发了一层缓存，用来存储已经拉取过的 BKTree！

lazy_static! {
    static ref BKTREE_CACHE: BKTreeCache = BKTreeCache::new();
}

当某个数据集首次执行启用了 typo-tolerance 的搜索时，我们会经历一次约 200～400ms 的冷启动，将该数据集的 BKTree 从 Redis 拉取到 server 内存中。之后的搜索便会使用这棵 BKTree 检查拼写错误，整个过程只需 100～300μs。

识别英语单词

由于我们的 BKTree 完全根据特定数据集的字典构建，因此可能没有收录所有有效的英语单词。为了避免错误纠正树中不存在的合法单词，我们会先执行一步英语单词识别：

我们在内存中维护了一个包含约 40 万个英语单词的 HashSet，并使用 lazy_static! 存储。

static ref ENGLISH_WORDS: HashSet<String> = {
        include_str!("../words.txt")
            .lines()
            .map(|s| s.to_lowercase())
            .collect()
    };

接下来，我们会检查这个单词是否只是一个添加了前缀或后缀的英语单词：

我们分别为常见前缀和后缀构建 Trie。

static ref PREFIX_TRIE: Trie = {
        let prefixes = vec![
            "anti", "auto", "de", "dis", "down", "extra", "hyper", "il", "im", "in", "ir", "inter",
            "mega", "mid", "mis", "non", "over", "out", "post", "pre", "pro", "re", "semi", "sub",
            "super", "tele", "trans", "ultra", "un", "under", "up",
        ];
        Trie::new(&prefixes)
    };
    static ref SUFFIX_TRIE: Trie = {
        let suffixes = vec![
            "able", "al", "ance", "ation", "ative", "ed", "en", "ence", "ent", "er", "es", "est",
            "ful", "ian", "ible", "ic", "ing", "ion", "ious", "ise", "ish", "ism", "ist", "ity",
            "ive", "ize", "less", "ly", "ment", "ness", "or", "ous", "s", "sion", "tion", "ty",
            "y",
        ];
        Trie::new(&suffixes)
    };

对于查询中的每个单词，我们都会搜索这些 Trie，找出最长的匹配前缀和后缀。

然后，我们从单词中移除这些词缀，得到词根。

移除词缀后，我们会执行最后一次字典检查：

在英语单词语料库中查找处理后的词根。

fn is_likely_english_word(word: &str) -> bool {
    if ENGLISH_WORDS.contains(&word.to_lowercase()) {
        return true;
    }

    // Check for prefix
    if let Some(prefix_len) = PREFIX_TRIE.longest_prefix(word) {
        if ENGLISH_WORDS.contains(&word[prefix_len..].to_lowercase()) {
            return true;
        }
    }

    // Check for suffix
    if let Some(suffix_len) = SUFFIX_TRIE.longest_suffix(word) {
        if ENGLISH_WORDS.contains(&word[..word.len() - suffix_len].to_lowercase()) {
            return true;
        }
    }

    // Check for compound words
    if word.contains('-') {
        let parts: Vec<&str> = word.split('-').collect();
        if parts
            .iter()
            .all(|part| ENGLISH_WORDS.contains(&part.to_lowercase()))
        {
            return true;
        }
    }

    false
}

对于未能通过英语单词检查的单词，我们会启动一次 BKTree 搜索：

• 查询 BKTree，找出最接近的匹配单词。

• 为每个不在字典中的单词生成一组候选纠正结果。

let mut best_correction = None;
            let mut best_score = 0;

            for ((correction, freq), distance) in tree.find(word.to_string(), max_distance) {
                if distance == 0 {
                    best_correction = None;
                    break;
                }
                if !is_best_correction(word, correction) {
                    continue;
                }

                let score = (max_distance - distance) * 1000 + *freq as isize;

                if score > best_score || best_correction.is_none() {
                    best_correction = Some(correction);
                    best_score = score;
                }
            }

            if let Some(correction) = best_correction {
                corrections.insert(word, correction.to_string());
            }

我们使用一套评分算法，从候选纠正结果中选出最佳结果：

我们的算法会优先考虑前缀匹配，同时将每个候选单词在数据集中的出现频率纳入计算。

fn is_best_correction(word: &str, correction: &str) -> bool {
    // Length-based filter
    let len_diff = (word.len() as i32 - correction.len() as i32).abs();
    if len_diff > 2 {
        return false;
    }

    // Prefix matching (adjust the length as needed)
    let prefix_len = std::cmp::min(1, std::cmp::min(word.len(), correction.len()));
    if word[..prefix_len] != correction[..prefix_len] {
        return false;
    }

    // Character set comparison
    let word_chars: HashSet<char> = word.chars().collect();
    let correction_chars: HashSet<char> = correction.chars().collect();
    let common_chars = word_chars.intersection(&correction_chars).count();
    let similarity_ratio =
        common_chars as f32 / word_chars.len().max(correction_chars.len()) as f32;

    similarity_ratio >= 0.8
}

Levitating 在 Hacker News 上评论说，按 points 排序搜索 FreeBSD 时，返回了不相关的结果。我们的 tokenizer 会按 camel case 切分，因此 FreeBSD 被处理成了 Free BSD FreeBSD。只包含单词 Free 的帖子获得的 points 比任何包含 FreeBSD 的帖子都多，因此排名也更高。能够在数据集层面有效检查完整单词后，我们便可以自动要求查询中的非英语单词必须出现。这样，对 FreeBSD 的查询就会变成 "FreeBSD"。

我们计划利用同一套系统实现查询拆分和拼接，因为这些功能都有相同的需求：快速在字典中查找单词。

Trieve 将始终致力于提供开箱即用的最佳相关性！你可以在我们的 Hacker News 搜索引擎中尝试它、注册免费的 cloud 账户，或查看我们的 self-hosting 指南。

如果需要采取进一步行动，你可以考虑屏蔽此人和／或举报滥用行为。