refactor: centralize index setup logic and add progress bar ETA

Extracts directory cleanup, partition initialization, and finalization into dedicated helper methods within KmerIndex. This centralizes force-flag handling and reduces boilerplate across merge, rebuild, and select workflows. Additionally updates the CLI progress bar template to display an ETA indicator following elapsed time.
This commit is contained in:
Eric Coissac
2026-08-16 13:15:59 +02:00
parent eee71430a4
commit d54ae272a4
5 changed files with 73 additions and 75 deletions
+60 -1
View File
@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
use std::fs;
use std::path::{Path, PathBuf};
use obikpartitionner::{KmerPartition, KmerSpectrum};
use obikpartitionner::{KmerPartition, KmerSpectrum, PARTITIONS_SUBDIR};
use obilayeredmap;
use obisys::{Reporter, Stage, progress_bar};
use rayon::prelude::*;
@@ -68,6 +68,65 @@ impl KmerIndex {
IndexMeta::exists(path.as_ref())
}
/// Make `output` ready to receive a newly created index.
///
/// If an index already exists there, remove it when `force` is set,
/// otherwise fail. A bare directory with no `index.meta` (e.g. one left
/// behind by a `DirLock`) is not considered a pre-existing index.
pub fn clear_output_for_create<P: AsRef<Path>>(output: P, force: bool) -> OKIResult<()> {
let output = output.as_ref();
if Self::exists(output) {
if force {
fs::remove_dir_all(output).map_err(OKIError::Io)?;
} else {
return Err(OKIError::Io(std::io::Error::new(
std::io::ErrorKind::AlreadyExists,
format!("{}: output directory already exists", output.display()),
)));
}
}
Ok(())
}
/// Lay out a fresh index skeleton at `output`: the root directory,
/// `index.meta` (from `meta`), and an opened, empty partition set.
///
/// For construction paths that build partitions from scratch (`select`,
/// `rebuild`). `merge` bootstraps by copying a source index instead, so
/// it does not use this.
pub(crate) fn create_skeleton<P: AsRef<Path>>(
output: P,
meta: &IndexMeta,
) -> OKIResult<KmerPartition> {
let output = output.as_ref();
fs::create_dir_all(output).map_err(OKIError::Io)?;
meta.write(output).map_err(OKIError::Io)?;
fs::create_dir_all(output.join(PARTITIONS_SUBDIR)).map_err(OKIError::Io)?;
Ok(KmerPartition::open_with_config(
output,
meta.config.kmer_size,
meta.config.minimizer_size,
meta.config.n_bits,
)?)
}
/// Mark `output` as fully indexed, pack its column matrices, and reopen it.
///
/// Shared tail of the `select`/`rebuild` construction paths, once their
/// partitions have been written.
pub(crate) fn finalize_indexed<P: AsRef<Path>>(
output: P,
rep: &mut Reporter,
) -> OKIResult<Self> {
let output = output.as_ref();
fs::File::create(output.join(SENTINEL_INDEXED)).map_err(OKIError::Io)?;
let idx = KmerIndex::open(output)?;
let t_pack = Stage::start("pack");
idx.pack_matrices()?;
rep.push(t_pack.stop());
Ok(idx)
}
/// Current construction state, as reported by sentinel files on disk.
pub fn state(&self) -> IndexState {
IndexState::detect(&self.root_path).unwrap_or(IndexState::Empty)
+1 -10
View File
@@ -130,16 +130,7 @@ impl KmerIndex {
let (source_labels, all_genomes) = compute_labels(sources, rename_duplicates)?;
// ── Prepare output directory ──────────────────────────────────────────
if output.exists() {
if force {
fs::remove_dir_all(output)?;
} else {
return Err(OKIError::Io(io::Error::new(
io::ErrorKind::AlreadyExists,
format!("{}: output directory already exists", output.display()),
)));
}
}
KmerIndex::clear_output_for_create(output, force)?;
// ── Bootstrap: copy first source to output ────────────────────────────
info!(
+5 -32
View File
@@ -1,15 +1,13 @@
use std::fs;
use std::io;
use std::path::Path;
use obikpartitionner::{KmerFilter, KmerPartition, MergeMode};
use obikpartitionner::{KmerFilter, MergeMode};
use obisys::{Reporter, Stage, progress_bar};
use tracing::info;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use crate::meta::IndexMeta;
use crate::state::{IndexState, SENTINEL_INDEXED};
use crate::state::IndexState;
impl KmerIndex {
/// Rebuild `src` into a new compact single-layer index at `output`.
@@ -40,36 +38,18 @@ impl KmerIndex {
));
}
if output.exists() {
if force {
fs::remove_dir_all(output)?;
} else {
return Err(OKIError::Io(io::Error::new(
io::ErrorKind::AlreadyExists,
format!("{}: output directory already exists", output.display()),
)));
}
}
KmerIndex::clear_output_for_create(output, force)?;
// ── Create output directory + metadata ────────────────────────────────
fs::create_dir_all(output)?;
let mut meta = IndexMeta::new(src.meta.config.clone());
meta.config.with_counts = mode == MergeMode::Count;
meta.genomes = src.meta.genomes.clone();
meta.write(output)?;
let n_genomes = src.meta.genomes.len();
let n_partitions = src.partition.n_partitions();
// ── Create an empty destination KmerPartition ─────────────────────────
// Create the partitions/ subdirectory so KmerPartition::open_with_config works.
fs::create_dir_all(output.join(obikpartitionner::PARTITIONS_SUBDIR))?;
let dst_partition = KmerPartition::open_with_config(
output,
meta.config.kmer_size,
meta.config.minimizer_size,
meta.config.n_bits,
)?;
let dst_partition = KmerIndex::create_skeleton(output, &meta)?;
info!(
"rebuild: {} partition(s), {} genome(s), mode={:?}",
@@ -94,13 +74,6 @@ impl KmerIndex {
rep.push(t.stop());
// Write SENTINEL_INDEXED — output is ready to use.
fs::File::create(output.join(SENTINEL_INDEXED))?;
let idx = KmerIndex::open(output)?;
let t_pack = Stage::start("pack");
idx.pack_matrices()?;
rep.push(t_pack.stop());
Ok(idx)
KmerIndex::finalize_indexed(output, rep)
}
}
+5 -30
View File
@@ -1,15 +1,13 @@
use std::fs;
use std::io;
use std::path::Path;
use obikpartitionner::{KmerPartition, OutputCol, PARTITIONS_SUBDIR};
use obikpartitionner::{KmerPartition, OutputCol};
use obisys::{Reporter, Stage, progress_bar};
use tracing::info;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use crate::meta::{GenomeInfo, IndexMeta};
use crate::state::{IndexState, SENTINEL_INDEXED};
use crate::state::IndexState;
impl KmerIndex {
/// Create a new index at `output` by projecting/aggregating the genome columns
@@ -33,35 +31,18 @@ impl KmerIndex {
return Err(OKIError::NotIndexed(src.root_path.clone()));
}
if output.exists() {
if force {
fs::remove_dir_all(output)?;
} else {
return Err(OKIError::Io(io::Error::new(
io::ErrorKind::AlreadyExists,
format!("{}: output directory already exists", output.display()),
)));
}
}
KmerIndex::clear_output_for_create(output, force)?;
fs::create_dir_all(output)?;
let mut meta = IndexMeta::new(src.meta.config.clone());
meta.config.with_counts = !output_presence;
meta.genomes = specs.iter()
.map(|s| GenomeInfo::new(s.label.clone()))
.collect();
meta.write(output)?;
let n_src_genomes = src.meta.genomes.len();
let n_partitions = src.partition.n_partitions();
fs::create_dir_all(output.join(PARTITIONS_SUBDIR))?;
let dst_partition = KmerPartition::open_with_config(
output,
meta.config.kmer_size,
meta.config.minimizer_size,
meta.config.n_bits,
)?;
let dst_partition = KmerIndex::create_skeleton(output, &meta)?;
info!(
"select: {} partition(s), {} source genome(s) → {} output column(s)",
@@ -83,13 +64,7 @@ impl KmerIndex {
pb.finish_and_clear();
rep.push(t.stop());
fs::File::create(output.join(SENTINEL_INDEXED))?;
let idx = KmerIndex::open(output)?;
let t_pack = Stage::start("pack");
idx.pack_matrices()?;
rep.push(t_pack.stop());
Ok(idx)
KmerIndex::finalize_indexed(output, rep)
}
/// Rewrite the genome columns of this index in-place according to `specs`.
+2 -2
View File
@@ -91,12 +91,12 @@ pub fn spinner(label: &str) -> TracedBar {
}
/// Progress bar with the standard project look:
/// `⠋ label — [████░░░░] pos/len unit elapsed`.
/// `⠋ label — [████░░░░] pos/len unit elapsed (eta: remaining)`.
pub fn progress_bar(label: &str, n: u64, unit: &str) -> TracedBar {
let pb = ProgressBar::new(n);
pb.set_style(
ProgressStyle::with_template(&format!(
"{{spinner}} {label} — {{bar:40.cyan/blue}} {{pos}}/{{len}} {unit} {{elapsed}}"
"{{spinner}} {label} — {{bar:40.cyan/blue}} {{pos}}/{{len}} {unit} {{elapsed}} (eta: {{eta}})"
))
.unwrap()
.tick_strings(BRAILLE),