crates/mimir-ingest/src/bin/mimir-lake.rs
//! `mimir-lake` — the lake's maintenance from the shell, for a lake no server
//! holds (a copy, a stopped server): what the server's `lake` lane does, with
//! numbers.
//!
//! ```text
//! mimir-lake <lake-dir> register a copied lake (metadata rewritten to this path,
//! no catalog.redb) gets its own catalog
//! mimir-lake <lake-dir> ring-split copy the ring to its own tables, verify, switch
//! (the compaction then partitions by repo)
//! mimir-lake <lake-dir> plan <table> <repo> files a scan of one repo opens vs the whole table
//! mimir-lake <lake-dir> stat tables: files, bytes, rows, codecs, snapshots
//! mimir-lake <lake-dir> fold [repo] time a full fold (the open set), rows open
//! mimir-lake <lake-dir> compact [--level N] [--grace-secs N] [--table T]
//! mimir-lake <lake-dir> append-bench <table> <parquet-file> [n] [--plain]
//! time n appends of that file's rows (the
//! consumer's write cost per release), then
//! roll them back by compacting nothing —
//! the rows STAY: run it on a copy only
//! ```
use std::path::{Path, PathBuf};
use std::time::Instant;
use anyhow::{Context, Result, bail};
use mimir_ingest::{CompactOptions, EventLog};
struct Args {
lake: PathBuf,
cmd: String,
rest: Vec<String>,
}
fn parse() -> Result<Args> {
let mut it = std::env::args().skip(1);
let lake = it.next().context("usage: mimir-lake <lake-dir> stat | fold [repo] | compact [--level N] [--grace-secs N] [--table T] | append-bench <table> <file> [n]")?;
let cmd = it.next().context("missing command")?;
Ok(Args { lake: PathBuf::from(lake), cmd, rest: it.collect() })
}
fn flag(rest: &[String], name: &str) -> Option<String> {
rest.iter().position(|a| a == name).and_then(|i| rest.get(i + 1).cloned())
}
fn main() -> Result<()> {
let a = parse()?;
if a.cmd == "register" {
return register(&a.lake);
}
let log = EventLog::open_lake_blocking(&a.lake).with_context(|| format!("open {}", a.lake.display()))?;
match a.cmd.as_str() {
"stat" => log.block_on(stat(&log, &a.lake)),
"fold" => log.block_on(fold(&log, a.rest.first().map(String::as_str))),
"compact" => log.block_on(compact(&log, &a.rest)),
"append-bench" => log.block_on(append_bench(&log, &a.rest)),
"ring-split" => log.block_on(async {
let r = log.split_ring(&mimir_ingest::SplitOptions::default()).await?;
println!("{r}");
Ok(())
}),
"plan" => log.block_on(plan(&log, &a.rest)),
other => bail!("unknown command `{other}`"),
}
}
/// `register`: a lake directory copied from elsewhere (its `metadata/*.json`
/// rewritten to this path, no `catalog.redb`) gets a catalog of its own: the
/// newest metadata file of every table registered under `mimir_v2`. Refuses
/// a metadata file whose `location` is not under this directory.
fn register(lake: &Path) -> Result<()> {
use skade::iceberg::Catalog as _;
if lake.join("catalog.redb").exists() {
bail!("{}: catalog.redb exists — register only fills a fresh catalog", lake.display());
}
let root = std::fs::canonicalize(lake)?;
let rt = tokio::runtime::Builder::new_multi_thread().worker_threads(2).enable_all().build()?;
rt.block_on(async {
let wh = skade::open(&root).await?;
let cat = wh.catalog();
let ns = skade::iceberg::NamespaceIdent::new("mimir_v2".into());
if !cat.namespace_exists(&ns).await? {
cat.create_namespace(&ns, Default::default()).await?;
}
for name in mimir_ingest::log::TABLE_NAMES {
let dir = root.join("warehouse/mimir_v2").join(name).join("metadata");
let mut metas: Vec<PathBuf> = std::fs::read_dir(&dir)
.with_context(|| format!("read {}", dir.display()))?
.flatten()
.map(|e| e.path())
.filter(|p| p.to_string_lossy().ends_with(".metadata.json"))
.collect();
metas.sort();
let Some(newest) = metas.pop() else { bail!("{name}: no metadata file under {}", dir.display()) };
let text = std::fs::read_to_string(&newest)?;
let json: serde_json::Value = serde_json::from_str(&text)?;
let loc = json.get("location").and_then(|v| v.as_str()).unwrap_or_default();
let want = format!("file://{}/warehouse/mimir_v2/{name}", root.display());
if loc != want {
bail!("{name}: {} says location {loc}, this lake is {want} — rewrite the metadata first", newest.display());
}
let ident = wh.table_ident(&format!("mimir_v2.{name}"))?;
cat.register_table(&ident, format!("file://{}", newest.display())).await?;
println!("registered {name} ← {}", newest.file_name().unwrap().to_string_lossy());
}
Ok::<(), anyhow::Error>(())
})
}
async fn stat(log: &EventLog, lake: &Path) -> Result<()> {
use skade::parquet::basic::Compression;
use skade::parquet::file::reader::{FileReader, SerializedFileReader};
let mut total_bytes = 0u64;
for name in mimir_ingest::log::TABLE_NAMES {
let t = log.warehouse().table(&format!("mimir_v2.{name}")).await?;
let md = t.inner().metadata();
let snaps = md.snapshots().count();
let cur = md.current_snapshot().map(|s| s.summary().additional_properties.clone()).unwrap_or_default();
let dir = lake.join("warehouse/mimir_v2").join(name).join("data");
let (mut files, mut bytes, mut zstd, mut plain) = (0usize, 0u64, 0usize, 0usize);
if let Ok(rd) = std::fs::read_dir(&dir) {
for e in rd.flatten() {
let p = e.path();
if p.extension().is_none_or(|x| x != "parquet") {
continue;
}
files += 1;
bytes += e.metadata().map(|m| m.len()).unwrap_or(0);
if let Ok(r) = SerializedFileReader::new(std::fs::File::open(&p)?) {
let m = r.metadata();
if m.num_row_groups() > 0 {
match m.row_group(0).column(0).compression() {
Compression::ZSTD(_) => zstd += 1,
_ => plain += 1,
}
}
}
}
}
total_bytes += bytes;
println!(
"{name:20} files {files:4} ({zstd} zstd, {plain} other) {:>8.1} MB rows {:>10} snapshots {snaps}",
bytes as f64 / 1e6,
cur.get("total-records").cloned().unwrap_or_default()
);
}
println!("total data {:.1} MB", total_bytes as f64 / 1e6);
Ok(())
}
/// The data files a filtered scan plans to open: `repo = R` vs no filter.
async fn plan(log: &EventLog, rest: &[String]) -> Result<()> {
let table = rest.first().context("plan <table> <repo>")?;
let repo = rest.get(1).context("plan <table> <repo>")?;
let t = log.warehouse().table(&format!("mimir_v2.{table}")).await?;
let all = skade::plan_stats(t.inner(), None).await?;
let one = skade::plan_stats(t.inner(), Some(&skade::ScanFilter::eq("repo", repo.as_str()))).await?;
println!("{table}: whole table {} file(s) / {} rows; repo {repo}: {} file(s) / {} rows ({} skipped)", all.data_files, all.rows_planned, one.data_files, one.rows_planned, all.data_files - one.data_files);
Ok(())
}
async fn fold(log: &EventLog, repo: Option<&str>) -> Result<()> {
let t0 = Instant::now();
let set = log.fold(repo, None).await?;
println!("fold{}: {} rows open, {:.1} MB held, {} ms", repo.map(|r| format!(" of {r}")).unwrap_or_default(), set.rows(), set.heap_bytes() as f64 / 1e6, t0.elapsed().as_millis());
Ok(())
}
async fn compact(log: &EventLog, rest: &[String]) -> Result<()> {
let mut opts = CompactOptions::default();
if let Some(l) = flag(rest, "--level") {
opts.level = l.parse().context("--level")?;
}
if let Some(g) = flag(rest, "--grace-secs") {
opts.grace = std::time::Duration::from_secs(g.parse().context("--grace-secs")?);
}
let t0 = Instant::now();
let reports = match flag(rest, "--table") {
Some(t) => vec![log.compact_table(&t, &opts).await?],
None => log.compact(&opts).await?,
};
for r in &reports {
println!("{r}");
}
println!(
"compaction: {} file(s) ({:.1} MB) → {} file(s) ({:.1} MB), swept {} file(s) ({:.1} MB), {} ms",
reports.iter().map(|r| r.files_in).sum::<usize>(),
reports.iter().map(|r| r.bytes_in).sum::<u64>() as f64 / 1e6,
reports.iter().map(|r| r.files_out).sum::<usize>(),
reports.iter().map(|r| r.bytes_out).sum::<u64>() as f64 / 1e6,
reports.iter().map(|r| r.swept_files).sum::<usize>(),
reports.iter().map(|r| r.swept_bytes).sum::<u64>() as f64 / 1e6,
t0.elapsed().as_millis()
);
Ok(())
}
/// Append the rows of a Parquet file to `table` `n` times (the consumer's
/// write path: encode + one fast-append commit each), timing each.
async fn append_bench(log: &EventLog, rest: &[String]) -> Result<()> {
use skade::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
let table = rest.first().context("append-bench <table> <file> [n]")?;
let file = rest.get(1).context("append-bench <table> <file> [n]")?;
let n: usize = rest.get(2).map(|s| s.parse()).transpose()?.unwrap_or(1);
let reader = ParquetRecordBatchReaderBuilder::try_new(std::fs::File::open(file)?)?.with_batch_size(64 * 1024).build()?;
let batches: Vec<_> = reader.collect::<std::result::Result<_, _>>()?;
// a release's batch holds one repo: the consumer sends the ring's rows
// to the ring tables, so a file of the lake's old shape is split here
let batches: Vec<skade::arrow_array::RecordBatch> = batches
.iter()
.map(|b| {
let repo = b.column_by_name("repo").context("repo")?;
let repo = skade::arrow_cast::cast(repo, &skade::arrow_schema::DataType::Utf8)?;
let keep = arrow_ord::cmp::neq(&repo, &skade::arrow_array::StringArray::new_scalar("@universe"))?;
Ok(arrow_select::filter::filter_record_batch(b, &keep)?)
})
.collect::<Result<_>>()?;
let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
let mut t = log.warehouse().table(&format!("mimir_v2.{table}")).await?;
// `--plain`: the lake's old write shape (uncompressed, 1 M-row groups)
let plain = rest.iter().any(|a| a == "--plain");
t.set_write_props(if plain { skade::WriteProps::default() } else { log.write_props().clone() });
println!("append-bench {table}: {rows} rows per append, {}", if plain { "UNCOMPRESSED (old)" } else { "zstd (new)" });
for i in 0..n {
let t0 = Instant::now();
t.append(&batches).await?;
println!("append {i}: {rows} rows in {} ms", t0.elapsed().as_millis());
}
Ok(())
}