Rust实现n-mer计数函数性能随n增大与Python持平问题排查

1. 字符串分配开销累积 #

当n较小时，子串长度短，内存分配开销占比低，Rust的原生循环优势被放大；但n增大后，若Rust代码中每次取子串都调用to_string()创建新String，频繁的内存分配与数据拷贝会成为性能瓶颈。而Python的字符串切片是零拷贝实现，仅生成指向原字符串的引用，无额外内存开销，这部分差距会随n增大逐渐抵消Rust的循环优势。

2. 哈希表性能瓶颈 #

当n增大，子串的哈希计算、哈希表插入/查询的开销占比显著提升：

• Rust默认HashMap使用SipHash哈希函数，设计偏向安全性，对长字符串的哈希速度不如Python内置的字符串优化哈希函数。
• 若未预分配哈希表容量，动态扩容的额外开销在n大时会被进一步放大。

3. 循环开销占比降低 #

n较小时，循环次数多（例如字符串长度为1e6，n=5时循环近1e6次），Rust的无边界检查、原生编译优势凸显；n增大后，循环次数骤减（n=200时循环约99.8k次），循环开销的占比下降，哈希与分配开销成为性能主导因素。

1. 使用零拷贝子串切片 #

将Rust中哈希表的键类型从String改为&str，直接引用原字符串的子串切片，彻底消除内存分配开销：

use pyo3::prelude::*;
use std::collections::HashMap;

#[pyfunction]
fn count_substrings(py: Python, s: &str, n: usize) -> PyResult<HashMap<&str, usize>> {
    if n == 0 || n > s.len() {
        return Ok(HashMap::new());
    }
    // 预分配哈希表容量，避免动态扩容
    let mut counts = HashMap::with_capacity(s.len() - n + 1);
    // 使用str::windows获取零拷贝子串迭代器，内部已优化边界检查
    for window in s.windows(n) {
        *counts.entry(window).or_insert(0) += 1;
    }
    Ok(counts)
}

2. 替换为更快的哈希函数 #

使用ahash crate的AHasher，它针对字符串等常见类型做了哈希优化，速度远快于默认SipHash：
首先在Cargo.toml添加依赖：

[dependencies]
pyo3 = "0.20.0"
ahash = "0.8.0"

修改代码使用FastHashMap：

use ahash::AHasher;
use pyo3::prelude::*;
use std::collections::HashMap;
use std::hash::BuildHasherDefault;

type FastHashMap<K, V> = HashMap<K, V, BuildHasherDefault<AHasher>>;

#[pyfunction]
fn count_substrings(py: Python, s: &str, n: usize) -> PyResult<FastHashMap<&str, usize>> {
    if n == 0 || n > s.len() {
        return Ok(FastHashMap::new());
    }
    let mut counts = FastHashMap::with_capacity(s.len() - n + 1);
    for window in s.windows(n) {
        *counts.entry(window).or_insert(0) += 1;
    }
    Ok(counts)
}

3. 减少Python-Rust类型转换开销 #

若最终需要返回Python字典，直接在Rust中构建PyDict，避免后续跨语言类型转换的额外开销：

#[pyfunction]
fn count_substrings(py: Python, s: &str, n: usize) -> PyResult<Py<PyDict>> {
    if n == 0 || n > s.len() {
        return Ok(PyDict::new(py).into());
    }
    let dict = PyDict::new(py);
    let mut counts = FastHashMap::with_capacity(s.len() - n + 1);
    for window in s.windows(n) {
        *counts.entry(window).or_insert(0) += 1;
    }
    for (k, v) in counts {
        dict.set_item(k, v)?;
    }
    Ok(dict.into())
}

优化后，即使n=200，Rust版本的性能会重新拉开与Python的差距——零拷贝子串消除了内存分配开销，更快的哈希函数降低了哈希计算耗时，预分配容量减少了哈希表扩容的额外开销，整体性能优势会随n增大保持稳定。

内容的提问来源于stack exchange，提问作者Alberto Marin Sanguino