使用rust-htslib快速提取VCF基因型的性能优化方案咨询

核心性能瓶颈分析 #

你的代码性能低下主要源于三个关键问题：

• 每次处理记录时重复创建Buffer，产生不必要的内存分配开销
• 将基因型值转换为字符串再解析（gts.get(sample_index).to_string()），这是最大的性能杀手——字符串序列化/反序列化的开销极大
• 使用HashMap存储inform，但样本索引是连续整数，HashMap的哈希查找远慢于直接数组访问

优化后的实现方案 #

以下是针对性优化后的代码，能大幅提升处理速度：

use rust_htslib::bcf::{Reader, record::{Genotypes, Buffer}};
use std::collections::HashMap;

const MAX_MARKERS: usize = 10000; // 替换为实际标记数量

fn main() {
    let vcf_path = "your_vcf_file.vcf";
    let mut bcf = Reader::from_path(vcf_path).expect("Error opening file.");
    let sample_count = usize::try_from(bcf.header().sample_count()).unwrap();
    
    let mut anml_lookup: HashMap<i32, usize> = HashMap::with_capacity(sample_count);
    // 预分配基因型存储，避免动态扩容
    let mut genotypes: Vec<Vec<i8>> = vec![Vec::with_capacity(MAX_MARKERS); sample_count];
    // 用Vec替代HashMap，利用连续索引的直接访问优势
    let mut inform: Vec<i32> = vec![0; sample_count];

    // 复用Buffer，避免每次循环创建新实例
    let mut buffer = Buffer::new();

    for record in bcf.records().map(|r| r.expect("Failed to read record")) {
        let gts = record.genotypes_shared_buffer(&mut buffer).expect("Can't get GTs");
        
        // 直接遍历原始基因型数据，跳过字符串转换步骤
        for (sample_idx, gt) in gts.iter().enumerate() {
            let (encoded_gt, info_val) = parse_gt_directly(gt);
            
            inform[sample_idx] += info_val;
            genotypes[sample_idx].push(encoded_gt);
        }
    }
}

// 直接解析GT原始值，完全避免字符串转换
fn parse_gt_directly(gt: &rust_htslib::bcf::record::Genotype) -> (i8, i32) {
    // 根据你的conv函数逻辑调整以下匹配逻辑
    match gt {
        rust_htslib::bcf::record::Genotype::Unphased(a, b) => {
            if *a == -1 || *b == -1 {
                // 缺失基因型的编码，按需调整
                (-1, 0)
            } else {
                let sum = (*a as i8) + (*b as i8);
                (sum, 1)
            }
        }
        rust_htslib::bcf::record::Genotype::Phased(a, b) => {
            if *a == -1 || *b == -1 {
                (-1, 0)
            } else {
                let sum = (*a as i8) + (*b as i8);
                (sum, 1)
            }
        }
        _ => (-1, 0) // 处理其他基因型类型，按需调整
    }
}

关键优化点详解 #

• 复用Buffer：将Buffer::new()移到循环外，每次调用genotypes_shared_buffer时传入同一个可变引用，消除重复的内存分配开销。
• 跳过字符串转换：直接遍历Genotype枚举的原始值，它已经提供了等位基因的整数表示（-1为缺失），无需转字符串再解析，这是性能提升的核心。
• 用Vec替代HashMap：样本索引是连续的0到sample_count-1，Vec的直接内存访问比HashMap的哈希查找快得多，且无哈希碰撞开销。
• 预分配容量：确保genotypes和inform都预分配了足够容量，避免运行时动态扩容带来的内存拷贝开销。

如果你的conv函数有更复杂的逻辑，只需修改parse_gt_directly函数，保持直接操作Genotype原始值即可。

内容的提问来源于stack exchange，提问作者AEON