Performance comparison

[andypohl]

I made a crate with three binaries that all count the number of times CAT exists in a FASTA file. I used noodles_bgzf and noodles_fasta. The goal was just a learning exercise to understand the runtime and memory footprint in a few strategies reading a compressed/indexed FASTA.

https://github.com/andypohl/cat-count

I did three styles:

1. Using a rayon parallel iterator to count on different FASTA records in parallel. Link
2. Same as (1), but using the lower-level sequence reader in order to use minimal memory. Link
3. Using tokio and the async readers. Link
4. Using crossbeam and some channels, contributed by @zaeleus ... changed slightly with the NonZero bit because I think that's nightly stuff. Link.
5. Basic sequential implementation, without a multithreaded BGZF decoder. Link.
6. Basic sequential implementation, with a multithreaded BGZF decoder. Link

Some basic results, not measured super scientifically just my rough observations:

Style	Approx wall time (s)	Approx peak RAM (MB)
(1) rayon_basic	2.8	1200
(2) rayon_stream	12.0	6
(3) tokio	3.5	400
(4) crossbeam	5.5	1200
(5) vanilla	11.0	800
(6) vanilla_w_multithread_reader	5.8	800

Does anyone have any suggestions? I'm still learning Rust, so I'm sure there's a lot of tricks I'm not aware of yet.

It seems like the Async reader for FASTA or BGZF is missing some capabilities that the synchronous code has. Namely that there's a low-level seqeuence_reader iterator in the non-async part not in the async part, and I couldn't figure out a way to seek to a particular sequence in a BGZF-compressed fasta with the async readers.

Again any feedback is welcome!

[zaeleus]

Interesting comparisons! The rayon_basic implementation actually looks about that same I would have done it. Here's a version using MPMC channels:

main.rs

// cargo add crossbeam-channel@0.5.13
// cargo add noodles@0.77.0 --features bgzf,fasta
// cargo add noodles-bgzf@0.31.0 --features libdeflate

use std::{
    env,
    fs::File,
    io::{self, SeekFrom},
    num::NonZero,
    path::Path,
    thread,
};

use noodles::{
    bgzf::{self, gzi, io::Seek},
    fasta::{self, fai},
};

const PATTERN: &[u8; 3] = b"CAT";

fn main() -> io::Result<()> {
    let src = env::args().nth(1).expect("missing src");

    let fai_src = format!("{src}.fai");
    let fasta_index = fai::read(fai_src)?;

    let gzi_src = format!("{src}.gzi");
    let gz_index = gzi::read(gzi_src)?;

    let worker_count = thread::available_parallelism().unwrap_or(NonZero::<usize>::MIN);

    let n = count(worker_count, src, &fasta_index, &gz_index)?;
    println!("{n}");

    Ok(())
}

fn count<P>(
    worker_count: NonZero<usize>,
    src: P,
    fasta_index: &fai::Index,
    gz_index: &gzi::Index,
) -> io::Result<usize>
where
    P: AsRef<Path>,
{
    let src = src.as_ref();

    let (count_tx, count_rx) = crossbeam_channel::bounded(worker_count.get());
    let (result_tx, result_rx) = crossbeam_channel::bounded(worker_count.get());

    thread::scope(move |scope| {
        scope.spawn(move || {
            for record in fasta_index {
                count_tx.send(record).unwrap();
            }
        });

        for _ in 0..worker_count.get() {
            let count_rx = count_rx.clone();
            let result_tx = result_tx.clone();

            scope.spawn(move || {
                while let Ok(record) = count_rx.recv() {
                    let mut decoder = File::open(src).map(bgzf::Reader::new)?;
                    decoder.seek_with_index(gz_index, SeekFrom::Start(record.offset()))?;

                    let mut reader = fasta::io::Reader::new(decoder);

                    let len = usize::try_from(record.length())
                        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
                    let mut buf = Vec::with_capacity(len);
                    reader.read_sequence(&mut buf)?;

                    buf.make_ascii_uppercase();

                    let n = buf
                        .windows(PATTERN.len())
                        .filter(|buf| buf == PATTERN)
                        .count();

                    result_tx.send(n).unwrap();
                }

                Ok::<_, io::Error>(())
            });
        }

        drop(result_tx);

        let mut sum = 0;

        while let Ok(n) = result_rx.recv() {
            sum += n;
        }

        Ok(sum)
    })
}

This has the same performance characteristics as rayon_basic.

Many open decoders at different file offsets will outperform a single sequential multithreaded reader, and notably, the latter is what's happening in your tokio_basic implementation. The decoder is multithreaded; however, it reads sequences and counts the substrings sequentially.

It seems like the Async reader for FASTA or BGZF is missing some capabilities that the synchronous code has.

I'll add a method to seek with an index, tracked in #273. Thanks for reporting!

[andypohl]

Thanks @zaeleus ! By the way I should probably mention my file and machine even though I'm only trying to compare roughly and between the various strategies. I'm using a Mac Mini with an 8-core M1 processor from fall 2020 16 GB RAM. Pretty average specs. Then I'm running on the chm13v2 genome which I downloaded gzipped but I unzipped/re-zipped with bgzip and indexed with samtools faidx.

I added three more strategies for counting CATs to my repository and added my results above. I hadn't tried with crossbeam (thanks!), but I did attempt a tokio multitask implementation spawning tokio tasks and using analogous tokio channels for inter-task communication. But I gave that up.

Anyway I'm not sure why the crossbeam one doesn't much my rayon_basic. At one point, I did get a ~2x speed boost changing from a classic for loop + array index + increment count variable, to the iterator with filter()/count(), but that's what you did, so I dunno. I changed it a little so I could get it to run on my normal non-nightly Rust, but I kept the main stuff intact.