crates/knut-popviz/src/permutate.rs
//! **Permutate** — change a fraction of the population in its Iceberg tables, the way
//! a real register changes between two loads, written the way Spark writes it.
//!
//! The tables are the Iceberg target's ([`crate::Target::Iceberg`]): one per label,
//! keyed on `_gar_id`, and one `rel_<REL>` table per `Person -> <label>` relationship
//! (`src`/`dst` = `_gar_id`s), every row carrying `round` (0 from a generate). A round
//! `N` picks `fraction` of the persons — deterministically, from `N` and the person's
//! id, so the same round on the same population changes the same people — and:
//!
//! * **moves** each chosen person that has an address to another existing address:
//! their `LIVES_AT` row gets the new `dst`, the person row the new address's
//! `lon`/`lat`/`kommun`;
//! * **flips** `ekonomiskt_bistand` for about a third of them;
//! * **closes a loan** for about a quarter of those holding one: the `HOLDS` row and the
//! `Loan` row are deleted;
//! * **buys a vehicle** for about a fifth: a new `Vehicle` row and a `DRIVES` row;
//! * stamps `round = N` on every row it writes for a changed entity.
//!
//! WRITTEN AS SPARK COPY-ON-WRITE: a table with changed or deleted rows is rewritten
//! whole in ONE `overwrite` commit (skade's `overwrite` over its `live_data_files`, through
//! whichever catalog holds the table: [`crate::lake`]) — the untouched rows carried over,
//! which skade's changelog cancels — and new rows are an `append`. So what reaches bifrost is exactly what a Spark
//! `UPDATE`/`DELETE`/`INSERT` would have produced: the full old row and the new one.
use std::collections::{HashMap, HashSet};
use std::path::Path;
use std::sync::Arc;
use std::time::Instant;
use arrow_array::cast::AsArray;
use arrow_array::types::{Float64Type, Int64Type};
use arrow_array::{ArrayRef, BooleanArray, Float64Array, Int64Array, RecordBatch, StringArray};
/// What one round changed.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct PermutateReport {
pub round: i64,
pub persons_changed: usize,
pub addresses_moved: usize,
pub bistand_flipped: usize,
pub loans_deleted: usize,
pub vehicles_added: usize,
pub seconds: f64,
}
impl PermutateReport {
pub fn json(&self) -> serde_json::Value {
serde_json::json!({
"event": "permutated", "round": self.round, "persons_changed": self.persons_changed,
"addresses_moved": self.addresses_moved, "bistand_flipped": self.bistand_flipped,
"loans_deleted": self.loans_deleted, "vehicles_added": self.vehicles_added,
"edges_moved": self.addresses_moved, "edges_added": self.vehicles_added,
"edges_deleted": self.loans_deleted, "seconds": self.seconds,
})
}
}
/// splitmix64 — the per-(round, person) draw. Deterministic, no RNG state.
fn mix(round: i64, id: i64) -> u64 {
let mut z = (round as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (id as u64).wrapping_add(0xD1B5_4A32_D192_ED03);
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
z ^ (z >> 31)
}
fn err(e: impl std::fmt::Display) -> String {
e.to_string()
}
fn col<'a>(b: &'a RecordBatch, name: &str) -> Result<&'a ArrayRef, String> {
b.column_by_name(name).ok_or_else(|| format!("no column `{name}`"))
}
fn i64s(b: &RecordBatch, name: &str) -> Result<Vec<i64>, String> {
Ok(col(b, name)?.as_primitive::<Int64Type>().values().to_vec())
}
/// Run one round on a local skade warehouse. See the module docs.
pub async fn permutate_iceberg(dir: &Path, fraction: f64, round: i64) -> Result<PermutateReport, String> {
permutate_lake(&crate::lake::Lake::Dir(dir.to_path_buf()), fraction, round).await
}
/// Run one round on the tables in `lake` - a local warehouse or an Iceberg REST catalog
/// (MinIO AIStor). `round` 0 is **the next round**: one above the highest `round` any
/// person carries, so a caller with no state (a button's unit) never repeats one. See
/// the module docs.
pub async fn permutate_lake(lake: &crate::lake::Lake, fraction: f64, round: i64) -> Result<PermutateReport, String> {
let t0 = Instant::now();
if !(0.0..=1.0).contains(&fraction) {
return Err("--permutate is a fraction in [0, 1]".into());
}
if round < 0 {
return Err("--round is a round number > 0, or 0 for the next one".into());
}
let wh = lake.open().await?;
let read = |name: &'static str| wh.read(name);
// ---- what is there ----------------------------------------------------------------
let persons = read("Person").await?;
let round = if round > 0 {
round
} else {
let mut top = 0i64;
for b in &persons {
top = top.max(i64s(b, "round")?.into_iter().max().unwrap_or(0));
}
top + 1
};
let addresses = read("Address").await?;
let lives = read("rel_LIVES_AT").await?;
let holds = read("rel_HOLDS").await?;
let loans = read("Loan").await?;
let mut geo: HashMap<i64, (f64, f64, String)> = HashMap::new();
let mut address_ids: Vec<i64> = Vec::new();
for b in &addresses {
let (id, lon, lat, k) = (i64s(b, "_gar_id")?, col(b, "lon")?.as_primitive::<Float64Type>(), col(b, "lat")?.as_primitive::<Float64Type>(), col(b, "kommun")?.as_string::<i32>());
for i in 0..b.num_rows() {
geo.insert(id[i], (lon.value(i), lat.value(i), k.value(i).to_string()));
address_ids.push(id[i]);
}
}
address_ids.sort_unstable();
let mut lives_at: HashMap<i64, i64> = HashMap::new();
for b in &lives {
let (s, d) = (i64s(b, "src")?, i64s(b, "dst")?);
for i in 0..b.num_rows() {
lives_at.insert(s[i], d[i]);
}
}
let mut first_loan: HashMap<i64, i64> = HashMap::new();
for b in &holds {
let (s, d) = (i64s(b, "src")?, i64s(b, "dst")?);
for i in 0..b.num_rows() {
first_loan.entry(s[i]).or_insert(d[i]);
}
}
// The next free id: above every _gar_id of every label (a vehicle added in an earlier
// round included), so a new vertex can never collide with another label's id.
let mut next_id = 0i64;
for (label, _) in crate::LABELS.iter() {
for b in wh.read_columns(label, &["_gar_id"]).await? {
if let Some(m) = i64s(&b, "_gar_id")?.into_iter().max() {
next_id = next_id.max(m + 1);
}
}
}
// ---- who changes, and how --------------------------------------------------------
let threshold = (fraction * u64::MAX as f64) as u64;
let mut report = PermutateReport { round, ..Default::default() };
let mut moved_to: HashMap<i64, i64> = HashMap::new();
let mut flipped: HashSet<i64> = HashSet::new();
let mut closed: HashMap<i64, i64> = HashMap::new(); // person -> loan
let mut bought: Vec<(i64, i64)> = Vec::new(); // (person, vehicle)
let mut chosen: HashSet<i64> = HashSet::new();
for b in &persons {
for p in i64s(b, "_gar_id")? {
let h = mix(round, p);
if h > threshold || fraction == 0.0 {
continue;
}
chosen.insert(p);
let r = h >> 16;
if let Some(&old) = lives_at.get(&p) {
if address_ids.len() > 1 {
let mut new = address_ids[(r as usize) % address_ids.len()];
if new == old {
new = address_ids[((r as usize) + 1) % address_ids.len()];
}
moved_to.insert(p, new);
}
}
if r % 3 == 0 {
flipped.insert(p);
}
if r % 4 == 1 {
if let Some(&l) = first_loan.get(&p) {
closed.insert(p, l);
}
}
if r % 5 == 2 {
bought.push((p, next_id));
next_id += 1;
}
}
}
report.persons_changed = chosen.len();
report.addresses_moved = moved_to.len();
report.bistand_flipped = flipped.len();
report.loans_deleted = closed.len();
report.vehicles_added = bought.len();
let closed_loans: HashSet<i64> = closed.values().copied().collect();
// ---- the rewritten tables --------------------------------------------------------
let mut new_persons = Vec::with_capacity(persons.len());
for b in &persons {
let id = i64s(b, "_gar_id")?;
let lon0 = col(b, "lon")?.as_primitive::<Float64Type>();
let lat0 = col(b, "lat")?.as_primitive::<Float64Type>();
let k0 = col(b, "kommun")?.as_string::<i32>();
let eb0 = col(b, "ekonomiskt_bistand")?.as_boolean();
let round0 = i64s(b, "round")?;
let (mut lon, mut lat, mut kommun, mut eb, mut rnd) = (vec![], vec![], vec![], vec![], vec![]);
for i in 0..b.num_rows() {
let p = id[i];
match moved_to.get(&p).and_then(|a| geo.get(a)) {
Some((x, y, k)) => {
lon.push(*x);
lat.push(*y);
kommun.push(k.clone());
}
None => {
lon.push(lon0.value(i));
lat.push(lat0.value(i));
kommun.push(k0.value(i).to_string());
}
}
eb.push(eb0.value(i) ^ flipped.contains(&p));
rnd.push(if chosen.contains(&p) { round } else { round0[i] });
}
new_persons.push(replace(b, vec![
("lon", Arc::new(Float64Array::from(lon)) as ArrayRef),
("lat", Arc::new(Float64Array::from(lat))),
("kommun", Arc::new(StringArray::from(kommun))),
("ekonomiskt_bistand", Arc::new(BooleanArray::from(eb))),
("round", Arc::new(Int64Array::from(rnd))),
])?);
}
let mut new_lives = Vec::with_capacity(lives.len());
for b in &lives {
let (s, d0, r0) = (i64s(b, "src")?, i64s(b, "dst")?, i64s(b, "round")?);
let d: Vec<i64> = s.iter().zip(&d0).map(|(p, d)| *moved_to.get(p).unwrap_or(d)).collect();
let r: Vec<i64> = s.iter().zip(&r0).map(|(p, r)| if moved_to.contains_key(p) { round } else { *r }).collect();
new_lives.push(replace(b, vec![("dst", Arc::new(Int64Array::from(d)) as ArrayRef), ("round", Arc::new(Int64Array::from(r)))])?);
}
let keep_rows = |b: &RecordBatch, keep: Vec<bool>| -> Result<RecordBatch, String> {
arrow_select::filter::filter_record_batch(b, &BooleanArray::from(keep)).map_err(err)
};
let mut new_holds = Vec::new();
for b in &holds {
let (s, d) = (i64s(b, "src")?, i64s(b, "dst")?);
let keep = s.iter().zip(&d).map(|(p, l)| closed.get(p) != Some(l)).collect();
new_holds.push(keep_rows(b, keep)?);
}
let mut new_loans = Vec::new();
for b in &loans {
let keep = i64s(b, "_gar_id")?.iter().map(|l| !closed_loans.contains(l)).collect();
new_loans.push(keep_rows(b, keep)?);
}
// ---- commit: an overwrite per changed table, an append per new row set ----------
let overwrite = |name: &'static str, batches: Vec<RecordBatch>| wh.overwrite_all(name, batches);
if !chosen.is_empty() {
overwrite("Person", new_persons).await?;
}
if !moved_to.is_empty() {
overwrite("rel_LIVES_AT", new_lives).await?;
}
if !closed.is_empty() {
overwrite("rel_HOLDS", new_holds).await?;
overwrite("Loan", new_loans).await?;
}
if !bought.is_empty() {
let n = bought.len();
let ids: Vec<i64> = bought.iter().map(|(_, v)| *v).collect();
let vehicles = RecordBatch::try_from_iter(vec![
("_gar_id", Arc::new(Int64Array::from(ids.clone())) as ArrayRef),
("idx", Arc::new(Int64Array::from(ids.clone()))),
("name", Arc::new(StringArray::from(bought.iter().map(|(p, _)| format!("v{round}_{p}")).collect::<Vec<_>>()))),
("round", Arc::new(Int64Array::from(vec![round; n]))),
])
.map_err(err)?;
let drives = RecordBatch::try_from_iter(vec![
("src", Arc::new(Int64Array::from(bought.iter().map(|(p, _)| *p).collect::<Vec<_>>())) as ArrayRef),
("dst", Arc::new(Int64Array::from(ids))),
("round", Arc::new(Int64Array::from(vec![round; n]))),
])
.map_err(err)?;
wh.append("Vehicle", vehicles).await?;
wh.append("rel_DRIVES", drives).await?;
}
report.seconds = t0.elapsed().as_secs_f64();
Ok(report)
}
/// `b` with the named columns replaced (same names, same order).
fn replace(b: &RecordBatch, with: Vec<(&str, ArrayRef)>) -> Result<RecordBatch, String> {
let schema = b.schema();
let mut cols = b.columns().to_vec();
for (name, arr) in with {
let i = schema.index_of(name).map_err(err)?;
cols[i] = arr;
}
RecordBatch::try_new(schema, cols).map_err(err)
}