diff --git a/kr2r/src/bin/annotate.rs b/kr2r/src/bin/annotate.rs index 08c5ec1..e7d4439 100644 --- a/kr2r/src/bin/annotate.rs +++ b/kr2r/src/bin/annotate.rs @@ -32,6 +32,10 @@ pub struct Args { #[clap(long, default_value_t = BATCH_SIZE)] pub batch_size: usize, + + /// The number of threads to use. + #[clap(short = 'p', long = "num-threads", value_parser, default_value_t = num_cpus::get())] + pub num_threads: usize, } fn read_chunk_header(reader: &mut R) -> io::Result<(usize, usize)> { @@ -90,6 +94,7 @@ fn process_batch( chunk_dir: PathBuf, batch_size: usize, page_index: usize, + num_threads: usize, ) -> std::io::Result<()> where R: Read + Send, @@ -105,7 +110,7 @@ where buffer_read_parallel( reader, - num_cpus::get(), + num_threads, batch_size, |dataset: Vec>| { let mut results: HashMap> = HashMap::new(); @@ -187,6 +192,7 @@ fn process_chunk_file>( args.chunk_dir.clone(), args.batch_size, page_index, + args.num_threads, )?; Ok(()) diff --git a/kr2r/src/bin/kun.rs b/kr2r/src/bin/kun.rs index fa9fa60..a8635af 100644 --- a/kr2r/src/bin/kun.rs +++ b/kr2r/src/bin/kun.rs @@ -76,6 +76,7 @@ impl From for annotate::Args { database: item.database, chunk_dir: item.chunk_dir, batch_size: item.batch_size, + num_threads: item.num_threads, } } } @@ -92,6 +93,7 @@ impl From for resolve::Args { report_kmer_data: item.report_kmer_data, report_zero_counts: item.report_zero_counts, full_output: item.full_output, + num_threads: item.num_threads, } } } diff --git a/kr2r/src/bin/resolve.rs b/kr2r/src/bin/resolve.rs index 43cfff7..2daac9d 100644 --- a/kr2r/src/bin/resolve.rs +++ b/kr2r/src/bin/resolve.rs @@ -98,6 +98,10 @@ pub struct Args { #[clap(long, default_value_t = BATCH_SIZE)] pub batch_size: usize, + + /// The number of threads to use. + #[clap(short = 'p', long = "num-threads", value_parser, default_value_t = num_cpus::get())] + pub num_threads: usize, } fn read_rows_from_file>(file_path: P) -> io::Result>> {