feat: add annotate CLI command for applying genome metadata via CSV

Introduces the `annotate` subcommand to apply genome metadata from an external CSV file to a pre-built k-mer index. The command supports configurable field separators, ID columns, and null markers, while providing a `--dump` option to export current index metadata as sorted CSV. Supporting changes include minor internal refactoring in `obikindex` to use object-level directory accessors and updates to the `csv` dependency.

src/Cargo.lock

@@ -1654,6 +1654,7 @@ name = "obikmer2"
 version = "1.2.2"
 dependencies = [
  "clap",
+ "csv",
  "obifastwrite",
  "obikalgorithm",
  "obikindex",

src/obikindex/examples/compare_sparse.rs

@@ -25,25 +25,25 @@ fn main() -> anyhow::Result<()> {
     let mut first_mismatch = None;
 
     for part in 0..n_parts {
-        let index_dir_sparse = sparse.index_dir(part);
-        let index_dir_dense = dense.index_dir(part);
-
-        if !index_dir_sparse.exists() || !index_dir_dense.exists() {
+        let Ok(part_sparse) = sparse.partition(part) else {
             continue;
-        }
+        };
+        let Ok(part_dense) = dense.partition(part) else {
+            continue;
+        };
 
         for layer in 0..n_layers {
-            let layer_dir_sparse = sparse.layer_dir(part, layer);
-            let layer_dir_dense = dense.layer_dir(part, layer);
-
-            if !layer_dir_sparse.exists() || !layer_dir_dense.exists() {
+            let Ok(layer_sparse) = part_sparse.layer(layer) else {
                 continue;
-            }
+            };
+            let Ok(layer_dense) = part_dense.layer(layer) else {
+                continue;
+            };
 
-            let presence_sparse = layer_dir_sparse.join("presence");
-            let presence_dense = layer_dir_dense.join("presence");
+            let presence_sparse = layer_sparse.presence_dir();
+            let layer_dir_dense = layer_dense.dir().to_path_buf();
 
-            if !presence_sparse.exists() || !presence_dense.exists() {
+            if !presence_sparse.exists() || !layer_dense.presence_dir().exists() {
                 continue;
             }
 

src/obikindex/src/layer/typed_layer.rs

@@ -150,7 +150,7 @@ impl HasStorageKind for PersistentBitMatrix {
 // ── Structures ────────────────────────────────────────────────────────────────
 
 pub struct TypedLayer<D: LayerData = ()> {
-    mphf: MphfLayer,
+    kmerstore: MphfLayer,
     data: D,
 }
 
@@ -165,11 +165,14 @@ impl<D: LayerData> TypedLayer<D> {
     pub fn open(path: &Path) -> OKIResult<Self> {
         let mphf = MphfLayer::open(path)?;
         let data = D::open(path)?;
-        Ok(Self { mphf, data })
+        Ok(Self {
+            kmerstore: mphf,
+            data,
+        })
     }
 
     pub fn query(&self, kmer: CanonicalKmer) -> Option<Hit<D::Item>> {
-        self.mphf.find(kmer).map(|slot| Hit {
+        self.kmerstore.find(kmer).map(|slot| Hit {
             slot,
             data: self.data.read(slot),
         })
@@ -181,37 +184,37 @@ impl<D: LayerData> TypedLayer<D> {
     /// sweep), so a plain `query` reading — and discarding — a full row
     /// per lookup would be wasted work.
     pub fn find_slot(&self, kmer: CanonicalKmer) -> Option<usize> {
-        self.mphf.find(kmer)
+        self.kmerstore.find(kmer)
     }
 
     pub fn n(&self) -> usize {
-        self.mphf.n()
+        self.kmerstore.n()
     }
 
     /// Raw MPHF lookup: kmer → slot, no membership check.
     pub fn hash(&self, kmer: CanonicalKmer) -> usize {
-        self.mphf.hash(kmer)
+        self.kmerstore.hash(kmer)
     }
 
     /// Batch raw MPHF lookup: kmers → slots, no membership check.
     pub fn hash_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
-        self.mphf.hash_batch(kmers)
+        self.kmerstore.hash_batch(kmers)
     }
 
     /// Iterate over all canonical kmers in the layer, in deterministic order.
     pub fn iter_kmers(&self) -> crate::layer::mphf_layer::KmerIter {
-        self.mphf.iter_kmers()
+        self.kmerstore.iter_kmers()
     }
 
     /// Iterate over all canonical kmers, each paired with its zero-based
     /// sequence index in `unitigs.bin`.
     pub fn enumerate_kmers(&self) -> std::iter::Enumerate<crate::layer::mphf_layer::KmerIter> {
-        self.mphf.enumerate_kmers()
+        self.kmerstore.enumerate_kmers()
     }
 
     /// Iterate over the layer's canonical kmers in batches of `n`.
     pub fn iter_kmers_batch(&self, n: usize) -> crate::layer::mphf_layer::KmerBatchIter {
-        self.mphf.iter_kmers_batch(n)
+        self.kmerstore.iter_kmers_batch(n)
     }
 
     /// Iterate over batches, each paired with the zero-based index of the
@@ -220,13 +223,13 @@ impl<D: LayerData> TypedLayer<D> {
         &self,
         n: usize,
     ) -> impl Iterator<Item = (usize, Vec<CanonicalKmer>)> + Send + 'static {
-        self.mphf.enumerate_kmers_batch(n)
+        self.kmerstore.enumerate_kmers_batch(n)
     }
 
     /// Iterate over the layer's unitigs (whole sequences, not decomposed
     /// into k-mers).
     pub fn iter_unitigs(&self) -> crate::layer::mphf_layer::UnitigIter {
-        self.mphf.iter_unitigs()
+        self.kmerstore.iter_unitigs()
     }
 
     pub fn unitig_writer(out_dir: &Path) -> OKIResult<UnitigFileWriter> {
@@ -251,7 +254,7 @@ impl<D: LayerData> TypedLayer<D> {
     /// already in memory (`MphfLayer`'s own `LayerEvidence`), no disk
     /// access. Available regardless of `D`, unlike `content`/`storage_kind`.
     pub fn evidence_kind(&self) -> crate::layer::mphf_layer::EvidenceKind {
-        self.mphf.evidence_kind()
+        self.kmerstore.evidence_kind()
     }
 }
 

src/obikmer2/Cargo.toml

@@ -19,6 +19,7 @@ obikmerge     = { path = "../obikmerge" }
 obifastwrite  = { path = "../obifastwrite" }
 obiskbuilder  = { path = "../obiskbuilder" }
 clap          = { version = "4", features = ["derive"] }
+csv           = "1"
 tracing       = "0.1.44"
 tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] }
 

src/obikmer2/src/cmd/annotate/mod.rs

@@ -0,0 +1,184 @@
+use std::collections::HashSet;
+use std::io::{self, BufWriter, Write};
+use std::path::PathBuf;
+
+use clap::Args;
+use obikindex::KmerIndex;
+use tracing::info;
+
+#[derive(Args)]
+pub struct AnnotateArgs {
+    /// Index directory to annotate (modified in-place)
+    pub index: PathBuf,
+
+    /// CSV file with genome metadata (must contain an id column)
+    #[arg(long)]
+    pub csv: Option<PathBuf>,
+
+    /// CSV field separator
+    #[arg(long, default_value = ",")]
+    pub sep: char,
+
+    /// Name of the column that contains genome labels
+    #[arg(long, default_value = "id")]
+    pub id_col: String,
+
+    /// Value that means "delete / absent" (removes existing key if present)
+    #[arg(long, default_value = "NA")]
+    pub na_value: String,
+
+    /// Do not overwrite existing metadata keys
+    #[arg(long)]
+    pub no_overwrite: bool,
+
+    /// Dump all genome metadata as CSV (stdout) instead of reading a CSV
+    #[arg(long)]
+    pub dump: bool,
+}
+
+pub fn run(args: AnnotateArgs) {
+    if args.dump {
+        run_dump(&args);
+    } else {
+        run_annotate(&args);
+    }
+}
+
+fn run_dump(args: &AnnotateArgs) {
+    let idx = open_index(&args.index);
+
+    let genomes = idx.meta().genomes().unwrap_or_else(|e| {
+        eprintln!("error reading index metadata: {e}");
+        std::process::exit(1);
+    });
+    let genomes = &genomes;
+
+    // Collect all keys in stable order (sorted for determinism)
+    let mut key_set: HashSet<String> = HashSet::new();
+    for g in genomes {
+        for k in g.meta.keys() {
+            key_set.insert(k.clone());
+        }
+    }
+    let mut keys: Vec<String> = key_set.into_iter().collect();
+    keys.sort();
+
+    let stdout = io::stdout();
+    let mut out = BufWriter::new(stdout.lock());
+
+    // Header
+    write!(out, "id").unwrap();
+    for k in &keys {
+        write!(out, "{}{k}", args.sep).unwrap();
+    }
+    writeln!(out).unwrap();
+
+    // Rows
+    for g in genomes {
+        write!(out, "{}", g.label).unwrap();
+        for k in &keys {
+            let v = g.meta.get(k).map(|s| s.as_str()).unwrap_or("NA");
+            write!(out, "{}{v}", args.sep).unwrap();
+        }
+        writeln!(out).unwrap();
+    }
+}
+
+fn run_annotate(args: &AnnotateArgs) {
+    let csv_path = match &args.csv {
+        Some(p) => p.clone(),
+        None => {
+            eprintln!("error: --csv is required unless --dump is used");
+            std::process::exit(1);
+        }
+    };
+
+    let idx = open_index(&args.index);
+
+    let mut genomes = idx.meta().genomes().unwrap_or_else(|e| {
+        eprintln!("error reading index metadata: {e}");
+        std::process::exit(1);
+    });
+
+    // Build a label → genome index position map
+    let label_to_pos: std::collections::HashMap<String, usize> = genomes
+        .iter()
+        .enumerate()
+        .map(|(i, g)| (g.label.clone(), i))
+        .collect();
+
+    let sep = args.sep as u8;
+    let mut rdr = csv::ReaderBuilder::new()
+        .delimiter(sep)
+        .from_path(&csv_path)
+        .unwrap_or_else(|e| {
+            eprintln!("error opening {}: {e}", csv_path.display());
+            std::process::exit(1);
+        });
+
+    let headers = rdr
+        .headers()
+        .unwrap_or_else(|e| {
+            eprintln!("error reading CSV headers: {e}");
+            std::process::exit(1);
+        })
+        .clone();
+
+    let id_col_idx = headers.iter().position(|h| h == args.id_col).unwrap_or_else(|| {
+        eprintln!("error: id column '{}' not found in CSV", args.id_col);
+        std::process::exit(1);
+    });
+
+    let meta_cols: Vec<(usize, String)> = headers
+        .iter()
+        .enumerate()
+        .filter(|(i, _)| *i != id_col_idx)
+        .map(|(i, h)| (i, h.to_string()))
+        .collect();
+
+    let mut updated = 0usize;
+    let mut skipped = 0usize;
+
+    for result in rdr.records() {
+        let record = result.unwrap_or_else(|e| {
+            eprintln!("error reading CSV record: {e}");
+            std::process::exit(1);
+        });
+
+        let label = record.get(id_col_idx).unwrap_or("").to_string();
+        let pos = match label_to_pos.get(&label) {
+            Some(&p) => p,
+            None => {
+                skipped += 1;
+                continue;
+            }
+        };
+
+        let genome = &mut genomes[pos];
+        for (col_idx, key) in &meta_cols {
+            let val = record.get(*col_idx).unwrap_or("");
+            if val == args.na_value {
+                genome.meta.remove(key);
+            } else if args.no_overwrite && genome.meta.contains_key(key) {
+                // skip
+            } else {
+                genome.meta.insert(key.clone(), val.to_string());
+            }
+        }
+        updated += 1;
+    }
+
+    idx.meta().set_genomes(genomes).unwrap_or_else(|e| {
+        eprintln!("error writing index metadata: {e}");
+        std::process::exit(1);
+    });
+
+    info!("annotated {updated} genome(s), skipped {skipped} CSV row(s) with unknown label");
+}
+
+fn open_index(path: &PathBuf) -> KmerIndex {
+    KmerIndex::open(path).unwrap_or_else(|e| {
+        eprintln!("error opening index: {e}");
+        std::process::exit(1);
+    })
+}

src/obikmer2/src/cmd/mod.rs

@@ -1,3 +1,4 @@
+pub mod annotate;
 pub mod estimate;
 pub mod index;
 pub mod merge;

src/obikmer2/src/main.rs

@@ -21,6 +21,8 @@ enum Commands {
     Merge(cmd::merge::MergeArgs),
     /// Estimate approximate-evidence false-positive rates for given parameters
     Estimate(cmd::estimate::EstimateArgs),
+    /// Read/write genome metadata (CSV) on an already-built index
+    Annotate(cmd::annotate::AnnotateArgs),
 }
 
 fn main() {
@@ -37,5 +39,6 @@ fn main() {
         Commands::Superkmer(args) => cmd::superkmer::run(args),
         Commands::Merge(args) => cmd::merge::run(args),
         Commands::Estimate(args) => cmd::estimate::run(args),
+        Commands::Annotate(args) => cmd::annotate::run(args),
     }
 }