如何使用rust-polars读取压缩TSV格式的*.gtf.gz文件？

前置依赖配置 #

首先在Cargo.toml中添加所需依赖：

[dependencies]
polars = { version = "0.37", features = ["csv", "categorical", "lazy"] }
flate2 = { version = "1.0", features = ["gzip"] }

1. 读取gz压缩文件 #

两种实现方式可选：

• 简单场景直接传文件路径：Polars会自动识别.gz后缀并解码，不需要手动处理压缩流
• 自定义流场景使用CsvReader::new()构造方法，传入解码后的流即可，示例代码：

use std::fs::File;
use std::io::BufReader;
use flate2::read::GzDecoder;

let file = File::open("test.gtf.gz")?;
let decoder = GzDecoder::new(BufReader::new(file));
// 后续将decoder传给CsvReader::new()即可

2. 正确解析Categorical类型 #

• 确保开启了Polars的categorical feature
• 新版本Polars的DataType::Categorical需要传入两个参数：映射规则、排序规则，日常使用直接填默认值即可，示例：

Field::new("contigName", DataType::Categorical(None, Default::default()))

• 若多线程读取时出现分类值异常，可先将列读取为Utf8类型，读完后再做类型转换：

df = df.with_column(col("contigName").cast(DataType::Categorical(None, Default::default())))?;

3. 处理缺失列/额外列 #

在CsvReader配置中添加两个参数即可：

• with_allow_extra_columns(true)：自动忽略超出Schema定义的列
• with_allow_missing_columns(true)：Schema要求但文件中缺失的列，自动填充对应类型的null值

4. 区分#开头表头和##开头注释 #

Polars自带的with_comment_char仅支持识别单个字符开头的注释，无法区分单#和双#，需要手动预处理流：

• 逐行读取解码后的内容，过滤所有##开头的注释行
• 保留#开头的行，可去掉开头的#作为表头使用
• 处理后的内容封装为Cursor传给CsvReader，同时设置has_header(true)

完整可运行代码示例 #

use std::fs::File;
use std::io::{BufRead, BufReader, Cursor};
use flate2::read::GzDecoder;
use polars::prelude::*;
use polars::csv::CsvEncoding;

fn main() -> PolarsResult<()> {
    // 定义Schema
    let schema = Arc::new(Schema::new(vec![
        Field::new("contigName", DataType::Categorical(None, Default::default())),
        Field::new("source", DataType::Utf8),
        Field::new("feature", DataType::Categorical(None, Default::default())),
        Field::new("start", DataType::Int64),
        Field::new("end", DataType::Int64),
        Field::new("score", DataType::Float32),
        Field::new("strand", DataType::Categorical(None, Default::default())),
        Field::new("frame", DataType::Categorical(None, Default::default())),
        Field::new("attribute", DataType::Utf8),
    ]));

    // 解码压缩文件
    let file = File::open("your_file.gtf.gz")?;
    let decoder = GzDecoder::new(BufReader::new(file));
    let mut buf_reader = BufReader::new(decoder);

    // 预处理注释与表头
    let mut content = String::new();
    let mut line = String::new();
    while buf_reader.read_line(&mut line)? > 0 {
        // 过滤双#开头的注释
        if line.starts_with("##") {
            line.clear();
            continue;
        }
        // 去掉单#表头的前缀
        if line.starts_with('#') {
            line = line.strip_prefix('#').unwrap_or(&line).to_string();
        }
        content.push_str(&line);
        line.clear();
    }

    // 读取预处理后的内容
    let cursor = Cursor::new(content);
    let df = CsvReader::new(cursor)
        .with_delimiter(b'\t')
        .with_schema(&schema)
        .with_allow_extra_columns(true)
        .with_allow_missing_columns(true)
        .with_encoding(CsvEncoding::LossyUtf8)
        .has_header(true)
        .finish()?;

    println!("{}", df.head(Some(10)));
    Ok(())
}

若处理GB级以上超大GTF文件，可自定义实现流式过滤的Read结构体，避免将全量文件读入内存。

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