将 csv 格式数据从 ChildStdout 流式传输到文件的最快方法

Question

我正在寻找加快将数据写入文件的方法，我尝试了几种方法，但都没有取得多大成功。 本质上这个过程是

• 生成一个通过 BQ 向 BigQuery 发送请求的子进程 CLI API。
• 使用std::io::copy 将数据流式传输到MutexGuard<Box<dyn Write>>。

这是我目前拥有的代码的小复制。这可能不会自行编译，但我认为它足以详细解释我在做什么。

use std::process::{Command, ChildStdout, Child, Stdio};
use std::time::Duration;
use std::io::{BufReader, BufWriter, Read, Write};
use std::sync::MutexGuard;
use std::box::Box;

const CAP: usize = 1024 * 8;

fn main() {
    let cmd = "bq --format csv query";
    let cmd = cmd.split_whitespace();
    let query = "path/to/query.sql"
    let (mut child, stdout) = spawn(cmd, query).await?

    let save_path = "path/to/save/file.csv"
    let file = File::open(save_path).unwrap();

    let mut rdr = BufReader::with_capacity(CAP, stdout);
    let wtr = Mutex::new(Box::new(BufWriter::with_capacity(CAP, file)));
    

    write(&mut child, &mut rdr, &wtr).unwrap();

}

fn create_cmd(args: Vec<&str>) -> Result<Command> {
    let mut v_itr = args.into_iter();
    let mut cmd = match v_itr.next() { 
        Some(s) => Command::new(which(s).unwrap()),
        None => return Err("Must supply at least one command argument.".into())
    };
    
    while let Some(s) = v_itr.next() {
        // This is essentially Vec's "push", but for adding arguments to a shell command.
        cmd.arg(s);
    };

    Ok(cmd)
}

fn spawn(cmd: Vec<u8>, query: &str) -> Result<ChildStdout> {
    let cmd = create_cmd(cmd)

    let query = std::fs::read_to_string(query).unwrap();
    let query = query.trim();
    cmd.arg(query);

    let mut child = match cmd
            .stdout(Stdio::piped())
            .spawn()
    {
        Ok(c) => c,
        Err(e) => return Err(e.into()),
    };

    // Sleep while waiting for the child process to fully start,
    // otherwise we may get an error when getting a handle to stdout.
    let sleep_time = Duration::new(2, 0);
    sleep(sleep_time);

    Ok(child, stdout)
}

fn write<R: Read, W: Write>(child: &mut Child, rdr: &mut R, wtr: Mutex<Box<W>>) -> Result<()> {
    let mut wtr_lock = wtr.lock().unwrap();
    let mut wtr_ref = &mut *wtr_lock;

    while let Ok(None) = child.try_wait() {
        std::io::copy(rdr, wtr_ref).unwrap();
    }

    Ok(())
}

到目前为止我已经尝试过：

• 使用 std BufWriter 和 BufReader 的 tokio 等效项
• 不使用任何缓冲输入
• 使用 LineWriter 进行输出
• 使用缓冲写入而不使用缓冲读取
• 从 std 实现与 stack_buffer_copy 的异步等效项，但使用 Tokio（这非常慢 - 我基本上逐字复制 std 中的内容，然后为包装 ReadBuf 的结构实现 BufMut）

我正在编写的数据大小实际上可以达到几 GiB，但我并不担心低于 20 MiB 的数据的速度。目前，我的速度略高于 1 MiB/s，这比我想要的要慢得多。我不介意解决方案是否复杂，我只是想看看有什么可能（我这样做主要是为了学习）。

附带说明一下，BQ 的响应从开始到结束通常只需要大约 5~10 秒（BQ 命令实际上将Waiting on bqjob_**** ... (10s) Current status: DONE 之类的内容打印到 stderr），所以我知道这不是瓶颈案子。我正在使用的计算机是具有 16GiB RAM 的 Mac M1 Air。我正在使用以下命令进行编译：RUSTFLAGS="--emit=asm" cargo build --release。

Answer 1

如果您不对数据进行任何处理，我会尽量避免管道和复制数据，而让子进程直接写入该文件。

所以不要这样：

cmd.stdout(Stdio::piped())

你可以通过你的file:

cmd.stdout(file)

这应该可行，因为Stdio 实现了From<File>。有一个 here 的示例（尽管使用标准输入，而不是标准输出）。

要测量你的基线速度，你可以从终端运行命令：

$ bq ... > path/to/save/file.csv

那么你就会知道在速度方面会发生什么。

附：如果您想知道为什么您的程序运行缓慢，请尝试使用flamegraph。