src/obikmer/src/cmd/query.rs

@@ -2,6 +2,8 @@ use std::collections::{HashMap, VecDeque};
 use std::io::{self, BufWriter, Write};
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
+use std::time::Instant;
 
 use clap::Args;
 use obikindex::KmerIndex;
@@ -11,8 +13,8 @@ use obilayeredmap::IndexMode;
 use obiread::chunk::read_sequence_chunks_sized;
 use obiread::record::{SeqRecord, parse_chunk};
 use obiskbuilder::SuperKmerIter;
-use obisys::available_memory_bytes;
+use obisys::{Reporter, Stage, available_memory_bytes, spinner};
-use tracing::info;
+use tracing::{debug, info};
 
 // ── Pipeline data ─────────────────────────────────────────────────────────────
 
@@ -74,6 +76,20 @@ pub struct QueryArgs {
     /// I/O chunk size in MiB (default: auto-sized from available RAM and thread count)
     #[arg(long)]
     pub chunk_size: Option<usize>,
+
+    /// Maximum number of input files open simultaneously.
+    /// Defaults to threads/4 (minimum 1). Keep below the number of workers
+    /// to ensure CPU workers are always available for the transform stage.
+    #[arg(long)]
+    pub max_open_files: Option<usize>,
+}
+
+impl QueryArgs {
+    pub fn effective_max_open(&self) -> usize {
+        self.max_open_files
+            .unwrap_or_else(|| (self.threads / 4).max(1))
+            .max(1)
+    }
 }
 
 // ── SKDesc — one occurrence of a superkmer in the batch ───────────────────────
@@ -234,6 +250,9 @@ fn process_chunk(
     force_presence: bool,
     presence_threshold: u32,
 ) -> Vec<u8> {
+    let chunk_start = Instant::now();
+    let chunk_bytes = rope.len();
+
     let records = parse_chunk(&rope, k);
     if records.is_empty() {
         return Vec::new();
@@ -384,6 +403,15 @@ fn process_chunk(
         &cov,
         &mut buf,
     );
+
+    debug!(
+        chunk_bytes,
+        n_seqs,
+        n_smers = batch.n_kmers.iter().map(|&n| n as u64).sum::<u64>(),
+        wall_ms = chunk_start.elapsed().as_millis() as u64,
+        "process_chunk"
+    );
+
     buf
 }
 
@@ -457,26 +485,71 @@ pub fn run(args: QueryArgs) {
         }
     });
 
+    // Instrumentation: total bytes processed (for the EMA throughput readout)
+    // and the number of chunks currently being processed by a worker — both
+    // read from the spinner loop below, updated from inside the pipe closure.
+    let total_bytes = Arc::new(AtomicU64::new(0));
+    let chunks_active = Arc::new(AtomicU32::new(0));
+
     let pipe = obipipeline::make_pipe! {
         QueryData : Rope => Vec<u8>,
         | {
             let idx = Arc::clone(&idx);
+            let total_bytes = Arc::clone(&total_bytes);
+            let chunks_active = Arc::clone(&chunks_active);
             move |rope: Rope| {
-                process_chunk(
+                chunks_active.fetch_add(1, Ordering::Relaxed);
+                let bytes = rope.len() as u64;
+                let out = process_chunk(
                     &idx, rope, k, n_genomes, n_partitions, with_counts,
                     effective_z, detail, count_missing, force_presence, presence_threshold,
-                )
+                );
+                total_bytes.fetch_add(bytes, Ordering::Relaxed);
+                chunks_active.fetch_sub(1, Ordering::Relaxed);
+                out
             }
         } : Chunk => Output,
     };
 
+    let t = Stage::start("query");
+    let pb = spinner("query");
+
+    let mut ema_rate: f64 = 0.0;
+    let mut last_t = Instant::now();
+    let mut last_bytes: u64 = 0;
+    const ALPHA: f64 = 0.15;
+
     let mut out = BufWriter::new(io::stdout());
     for block in pipe.apply(all_chunks, n_workers, 2) {
         if !block.is_empty() {
             out.write_all(&block).expect("write error");
         }
+
+        let now = Instant::now();
+        let dt = now.duration_since(last_t).as_secs_f64();
+        if dt > 0.1 {
+            let total = total_bytes.load(Ordering::Relaxed);
+            let instant = (total - last_bytes) as f64 / dt;
+            ema_rate = ALPHA * instant + (1.0 - ALPHA) * ema_rate;
+            last_t = now;
+            last_bytes = total;
+            let bp = total as f64;
+            let (count_str, rate_str) = if bp >= 1e9 {
+                (format!("{:.2} GB", bp / 1e9), format!("{:.0} MB/s", ema_rate / 1e6))
+            } else {
+                (format!("{:.0} MB", bp / 1e6), format!("{:.0} MB/s", ema_rate / 1e6))
+            };
+            let active = chunks_active.load(Ordering::Relaxed);
+            pb.set_message(format!("{count_str}  {rate_str}  [chunks in flight: {active}]"));
+        }
     }
     out.flush().expect("flush error");
+
+    pb.finish_and_clear();
+
+    let mut rep = Reporter::new();
+    rep.push(t.stop());
+    rep.print();
 }
 
 // ── Output ────────────────────────────────────────────────────────────────────