src/clone_signals/store.rs
//! `fn_shape_events` — the per-function shapes, a change log beside the
//! function nodes (the same fold as `embedding_node_events`): per
//! (repo, moniker) the last event with `seq ≤ N` is the state at ordinal N.
//!
//! | column | |
//! |---|---|
//! | `repo · seq · sha` | where the state is |
//! | `op` | `+` (new or changed shape) / `-` (no longer a function with a shape) |
//! | `moniker` | the node id |
//! | `shape` | the normalisation rule ([`super::norm::SHAPE`]) |
//! | `exact` | SHA-256 of the normalised text: the exact key, and the key of its normalised vector |
//! | `tokens` | normalised tokens (the reporting floor reads it) |
//! | `minhash` | the MinHash signature, [`super::minhash::PERMS`] little-endian u32 |
use std::collections::{BTreeMap, HashMap};
use std::sync::Arc;
use anyhow::Result;
use arrow::array::{Array, Int64Array, LargeBinaryArray, RecordBatch, StringArray, TimestampMicrosecondArray};
use futures::TryStreamExt;
use iceberg::Catalog;
use iceberg::arrow::schema_to_arrow_schema;
use iceberg::expr::{Predicate, Reference};
use iceberg::spec::Datum;
use super::norm::SHAPE;
use crate::vector::store::col;
use crate::warehouse::iceberg::{IcebergWarehouse, TABLE_FN_ROLE_VECTORS, TABLE_FN_SHAPE_EVENTS, append_batch};
/// One function's shape.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Shape {
pub exact: String,
pub tokens: i64,
pub minhash: Vec<u32>,
/// A body written for one CPU feature / width (`#[target_feature]`,
/// `std::arch`, `_mm*` / NEON intrinsics): its look-alikes are deliberate.
pub simd: bool,
}
/// One shape event (`shape: None` = a `-`).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Event {
pub repo: String,
pub seq: i64,
pub sha: String,
pub moniker: String,
pub shape: Option<Shape>,
pub ts_micros: i64,
}
/// The open shapes of `repo` at ordinal `upto`: per moniker the last event
/// with `seq ≤ upto`, kept when it is a `+`. Events of one ordinal apply in
/// the order they were written.
pub fn open<'a>(events: &'a [Event], repo: &str, upto: i64) -> BTreeMap<&'a str, &'a Shape> {
let mut mine: Vec<&Event> = events.iter().filter(|e| e.repo == repo && e.seq <= upto).collect();
mine.sort_by_key(|e| (e.seq, e.ts_micros));
let mut state = BTreeMap::new();
for e in mine {
match &e.shape {
Some(s) => {
state.insert(e.moniker.as_str(), s);
}
None => {
state.remove(e.moniker.as_str());
}
}
}
state
}
async fn table(wh: &IcebergWarehouse) -> Result<Option<iceberg::table::Table>> {
let ident = wh.table_ident(TABLE_FN_SHAPE_EVENTS);
if !wh.catalog().table_exists(&ident).await? {
return Ok(None);
}
Ok(Some(wh.catalog().load_table(&ident).await?))
}
/// Append shape events (one commit). Nothing to write is no commit.
pub fn append(wh: &IcebergWarehouse, events: &[Event]) -> Result<()> {
if events.is_empty() {
return Ok(());
}
wh.block_on(async {
let table = wh.catalog().load_table(&wh.table_ident(TABLE_FN_SHAPE_EVENTS)).await?;
let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?);
let strings = |it: Vec<&str>| -> Arc<dyn Array> { Arc::new(StringArray::from(it)) };
let ts = chrono::Utc::now().timestamp_micros();
let cols: Vec<Arc<dyn Array>> = vec![
strings(events.iter().map(|e| e.repo.as_str()).collect()),
Arc::new(Int64Array::from_iter_values(events.iter().map(|e| e.seq))),
strings(events.iter().map(|e| e.sha.as_str()).collect()),
strings(events.iter().map(|e| if e.shape.is_some() { "+" } else { "-" }).collect()),
strings(events.iter().map(|e| e.moniker.as_str()).collect()),
strings(events.iter().map(|_| SHAPE).collect()),
Arc::new(StringArray::from(events.iter().map(|e| e.shape.as_ref().map(|s| s.exact.as_str())).collect::<Vec<_>>())),
Arc::new(Int64Array::from_iter_values(events.iter().map(|e| e.shape.as_ref().map_or(0, |s| s.tokens)))),
Arc::new(LargeBinaryArray::from(events.iter().map(|e| e.shape.as_ref().map(|s| super::minhash::to_bytes(&s.minhash))).collect::<Vec<_>>().iter().map(|b| b.as_deref()).collect::<Vec<_>>())),
Arc::new(StringArray::from(events.iter().map(|e| e.shape.as_ref().filter(|s| s.simd).map(|_| "simd")).collect::<Vec<_>>())),
Arc::new(TimestampMicrosecondArray::from(vec![ts; events.len()]).with_timezone("+00:00")),
];
append_batch(wh.catalog(), table, RecordBatch::try_new(schema, cols)?).await
})
}
fn repo_filter(repos: &[&str]) -> Option<Predicate> {
repos.iter().map(|r| Reference::new("repo").equal_to(Datum::string(*r))).reduce(|a, b| a.or(b))
}
/// The shape events of `repos` (every repo when empty) under the current rule.
/// A warehouse laid out before the table existed reads as no event.
pub fn read(wh: &IcebergWarehouse, repos: &[&str]) -> Result<Vec<Event>> {
wh.block_on(async {
let Some(table) = table(wh).await? else { return Ok(Vec::new()) };
let mut scan = table.scan().select(["repo", "seq", "sha", "op", "moniker", "shape", "exact", "tokens", "minhash", "flags", "ts_micros"]);
if let Some(p) = repo_filter(repos) {
scan = scan.with_filter(p);
}
let batches: Vec<RecordBatch> = scan.build()?.to_arrow().await?.try_collect().await?;
let mut out = Vec::new();
for b in &batches {
let (repo, seq, sha, op, moniker) = (col::<StringArray>(b, "repo")?, col::<Int64Array>(b, "seq")?, col::<StringArray>(b, "sha")?, col::<StringArray>(b, "op")?, col::<StringArray>(b, "moniker")?);
let (shape, exact, tokens, minhash, ts) =
(col::<StringArray>(b, "shape")?, col::<StringArray>(b, "exact")?, col::<Int64Array>(b, "tokens")?, col::<LargeBinaryArray>(b, "minhash")?, col::<TimestampMicrosecondArray>(b, "ts_micros")?);
let flags = col::<StringArray>(b, "flags")?;
for i in 0..b.num_rows() {
if shape.value(i) != SHAPE || (!repos.is_empty() && !repos.contains(&repo.value(i))) {
continue;
}
let s = (op.value(i) == "+" && !exact.is_null(i) && !minhash.is_null(i))
.then(|| super::minhash::from_bytes(minhash.value(i)).map(|m| Shape { exact: exact.value(i).to_string(), tokens: tokens.value(i), minhash: m, simd: !flags.is_null(i) && flags.value(i).contains("simd") }))
.flatten();
out.push(Event { repo: repo.value(i).to_string(), seq: seq.value(i), sha: sha.value(i).to_string(), moniker: moniker.value(i).to_string(), shape: s, ts_micros: ts.value(i) });
}
}
Ok(out)
})
}
// ── fn_role_vectors ─────────────────────────────────────────────────────────
/// The role-vector space name (the `space` column; a new FastRP rule is a new space).
pub const ROLE_SPACE: &str = "fastrp-v1";
/// One stored role vector.
#[derive(Debug, Clone, PartialEq)]
pub struct RoleRow {
pub repo: String,
pub stable_id: String,
pub seq: i64,
pub feat_hash: u64,
pub vector: Vec<f32>,
pub ts_micros: i64,
}
/// Append role vectors (one commit). Nothing to write is no commit.
pub fn append_roles(wh: &IcebergWarehouse, rows: &[RoleRow]) -> Result<()> {
append_roles_at(wh, rows, chrono::Utc::now().timestamp_micros())
}
fn append_roles_at(wh: &IcebergWarehouse, rows: &[RoleRow], ts: i64) -> Result<()> {
if rows.is_empty() {
return Ok(());
}
wh.block_on(async {
let table = wh.catalog().load_table(&wh.table_ident(TABLE_FN_ROLE_VECTORS)).await?;
let batch = roles_batch(&table, rows.iter(), Some(ts))?;
append_batch(wh.catalog(), table, batch).await
})
}
/// The role table's batch of `rows`: stamped `ts`, or each row's own.
fn roles_batch<'a>(table: &iceberg::table::Table, rows: impl Iterator<Item = &'a RoleRow> + Clone, ts: Option<i64>) -> Result<RecordBatch> {
let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?);
let bytes: Vec<Vec<u8>> = rows.clone().map(|r| r.vector.iter().flat_map(|x| x.to_le_bytes()).collect()).collect();
let hashes: Vec<String> = rows.clone().map(|r| format!("{:016x}", r.feat_hash)).collect();
let cols: Vec<Arc<dyn Array>> = vec![
Arc::new(StringArray::from_iter_values(rows.clone().map(|r| r.repo.as_str()))),
Arc::new(StringArray::from_iter_values(rows.clone().map(|r| r.stable_id.as_str()))),
Arc::new(Int64Array::from_iter_values(rows.clone().map(|r| r.seq))),
Arc::new(StringArray::from_iter_values(hashes.iter().map(String::as_str))),
Arc::new(StringArray::from_iter_values(rows.clone().map(|_| ROLE_SPACE))),
Arc::new(LargeBinaryArray::from_iter_values(bytes.iter().map(Vec::as_slice))),
Arc::new(TimestampMicrosecondArray::from_iter_values(rows.map(|r| ts.unwrap_or(r.ts_micros))).with_timezone("+00:00")),
];
Ok(RecordBatch::try_new(schema, cols)?)
}
/// **Fold the append-only role table to its newest row per stable id**
/// when it holds more than twice that (the live table: 3.6 GB in 903 files,
/// 71 s to read): ONE overwrite of the files live at `snap` by the fold —
/// skipped when the table moved past `snap` (an append since the read).
/// Every read after it folds the same rows from a fraction of the bytes.
fn compact_roles(wh: &IcebergWarehouse, fold: &RoleFold, snap: Option<i64>) -> Result<bool> {
wh.block_on(async {
let table = wh.catalog().load_table(&wh.table_ident(TABLE_FN_ROLE_VECTORS)).await?;
if table.metadata().current_snapshot().map(|s| s.snapshot_id()) != snap {
return Ok(false);
}
let remove = skade::live_data_files(&table).await?;
let mut rows: Vec<&RoleRow> = fold.values().collect();
rows.sort_by(|a, b| (a.repo.as_str(), a.stable_id.as_str()).cmp(&(b.repo.as_str(), b.stable_id.as_str())));
let batch = roles_batch(&table, rows.iter().copied(), None)?;
skade::overwrite(wh.catalog().as_ref(), &table, &remove, std::slice::from_ref(&batch)).await?;
Ok(true)
})
}
/// The newest role vector per stable id (this space only).
pub fn read_roles(wh: &IcebergWarehouse) -> Result<HashMap<String, RoleRow>> {
Ok(read_roles_counted(wh)?.0)
}
/// [`read_roles`] and the rows it read, streamed (a batch is folded and
/// dropped as it decodes, never the whole table at once).
fn read_roles_counted(wh: &IcebergWarehouse) -> Result<(HashMap<String, RoleRow>, u64)> {
wh.block_on(async {
let ident = wh.table_ident(TABLE_FN_ROLE_VECTORS);
if !wh.catalog().table_exists(&ident).await? {
return Ok((HashMap::new(), 0));
}
let table = wh.catalog().load_table(&ident).await?;
let mut stream = table.scan().build()?.to_arrow().await?;
let mut out: HashMap<String, RoleRow> = HashMap::new();
let mut read = 0u64;
while let Some(b) = stream.try_next().await? {
let b = &b;
read += b.num_rows() as u64;
let (repo, id, seq, hash, space) = (col::<StringArray>(b, "repo")?, col::<StringArray>(b, "stable_id")?, col::<Int64Array>(b, "seq")?, col::<StringArray>(b, "feat_hash")?, col::<StringArray>(b, "space")?);
let (vector, ts) = (col::<LargeBinaryArray>(b, "vector")?, col::<TimestampMicrosecondArray>(b, "ts_micros")?);
for i in 0..b.num_rows() {
if space.value(i) != ROLE_SPACE {
continue;
}
let row = RoleRow {
repo: repo.value(i).to_string(),
stable_id: id.value(i).to_string(),
seq: seq.value(i),
feat_hash: u64::from_str_radix(hash.value(i), 16).unwrap_or(0),
vector: vector.value(i).chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect(),
ts_micros: ts.value(i),
};
let newer = out.get(&row.stable_id).is_none_or(|o| (o.ts_micros, o.seq) <= (row.ts_micros, row.seq));
if newer {
out.insert(row.stable_id.clone(), row);
}
}
}
Ok((out, read))
})
}
// ── the kept read: the role table, held by the process ──────────────────────
//
// `read_roles` scans the WHOLE append-only table (every vector ever written)
// to fold the newest per stable id — on the live workspace seconds per ask,
// and growing. The process keeps the fold beside the table snapshot it read
// at; the next read checks the snapshot (table metadata only): unchanged →
// the kept fold IS the read; moved by this process's own append → the
// appended rows were folded in by the same rule; moved by anyone else (or a
// compaction) → read again. ZERO-LOSS: the kept fold equals a fresh read.
type RoleFold = HashMap<String, RoleRow>;
struct KeptRoles {
snapshot: Option<i64>,
rows: Arc<RoleFold>,
}
fn kept_roles() -> &'static std::sync::Mutex<HashMap<std::path::PathBuf, KeptRoles>> {
static KEPT: std::sync::OnceLock<std::sync::Mutex<HashMap<std::path::PathBuf, KeptRoles>>> = std::sync::OnceLock::new();
KEPT.get_or_init(|| std::sync::Mutex::new(HashMap::new()))
}
/// The role table's current snapshot (`None`: no table, or no snapshot yet).
fn role_snapshot(wh: &IcebergWarehouse) -> Result<Option<i64>> {
wh.block_on(async {
let ident = wh.table_ident(TABLE_FN_ROLE_VECTORS);
if !wh.catalog().table_exists(&ident).await? {
return Ok(None);
}
let table = wh.catalog().load_table(&ident).await?;
Ok(table.metadata().current_snapshot().map(|s| s.snapshot_id()))
})
}
/// How the last [`read_roles_kept`] answered: `true` = the kept fold.
static LAST_ROLE_READ_KEPT: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
pub fn last_role_read_was_kept() -> bool {
LAST_ROLE_READ_KEPT.load(std::sync::atomic::Ordering::Relaxed)
}
/// [`read_roles`], kept by the process beside the snapshot it was read at.
pub fn read_roles_kept(wh: &IcebergWarehouse) -> Result<Arc<RoleFold>> {
let snap = role_snapshot(wh)?;
let root = wh.root().to_path_buf();
if let Ok(c) = kept_roles().lock()
&& let Some(k) = c.get(&root)
&& snap.is_some()
&& k.snapshot == snap
{
LAST_ROLE_READ_KEPT.store(true, std::sync::atomic::Ordering::Relaxed);
return Ok(k.rows.clone());
}
LAST_ROLE_READ_KEPT.store(false, std::sync::atomic::Ordering::Relaxed);
let (fold, read) = read_roles_counted(wh)?;
// the newest row per id, folded into the table itself when history is
// most of it (best effort: a failed or skipped fold changes no answer)
let snap = if read > 2 * fold.len() as u64 + 50_000 {
match compact_roles(wh, &fold, snap) {
Ok(true) => {
eprintln!("clone_signals: fn_role_vectors folded to its newest row per id: {read} -> {} row(s)", fold.len());
role_snapshot(wh)?
}
Ok(false) => snap,
Err(e) => {
eprintln!("clone_signals: fn_role_vectors fold skipped: {e:#}");
snap
}
}
} else {
snap
};
let rows = Arc::new(fold);
// the snapshot read BEFORE the scan: a commit between the two is re-read next time
if let Ok(mut c) = kept_roles().lock() {
c.insert(root, KeptRoles { snapshot: snap, rows: rows.clone() });
}
Ok(rows)
}
/// [`append_roles`], folding the appended rows into the kept read when the
/// table was at the kept snapshot right before the append (else the kept read
/// is dropped: the next read scans).
pub fn append_roles_kept(wh: &IcebergWarehouse, rows: &[RoleRow]) -> Result<()> {
if rows.is_empty() {
return Ok(());
}
let root = wh.root().to_path_buf();
let before = role_snapshot(wh);
let ts = chrono::Utc::now().timestamp_micros();
let wrote = append_roles_at(wh, rows, ts);
let after = role_snapshot(wh);
let Ok(mut c) = kept_roles().lock() else { return wrote };
let k = c.remove(&root);
if let (Ok(()), Ok(Some(before)), Ok(Some(after)), Some(mut k)) = (&wrote, &before, &after, k)
&& k.snapshot == Some(*before)
&& after != before
{
let fold = Arc::make_mut(&mut k.rows);
for r in rows {
let row = RoleRow { ts_micros: ts, ..r.clone() };
if fold.get(&row.stable_id).is_none_or(|o| (o.ts_micros, o.seq) <= (row.ts_micros, row.seq)) {
fold.insert(row.stable_id.clone(), row);
}
}
k.snapshot = Some(*after);
c.insert(root, k);
}
wrote
}
/// Forget the kept role reads (tests).
pub fn forget_kept_roles() {
if let Ok(mut c) = kept_roles().lock() {
c.clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
/// ZERO-LOSS of the kept role read: after the process's own appends the
/// kept fold equals a fresh scan; a write by anyone else (here: the plain
/// `append_roles`, which does not fold into the kept read) moves the
/// snapshot and the next kept read scans again.
#[test]
fn the_kept_role_read_equals_a_fresh_read_and_rereads_after_a_foreign_write() {
let role = |id: &str, x: f32, seq: i64| RoleRow { repo: "r".into(), stable_id: id.into(), seq, feat_hash: 7, vector: vec![x; 4], ts_micros: 0 };
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
append_roles(&wh, &[role("a", 0.5, 1), role("b", 0.25, 1)]).unwrap();
let first = read_roles_kept(&wh).unwrap();
assert!(!last_role_read_was_kept(), "the first read scans");
assert_eq!(*first, read_roles(&wh).unwrap());
let again = read_roles_kept(&wh).unwrap();
assert!(last_role_read_was_kept(), "an unchanged table: the kept fold");
assert_eq!(*again, *first);
// the process's own append is folded in by the same rule
append_roles_kept(&wh, &[role("a", 0.75, 2), role("c", 1.0, 2)]).unwrap();
let after = read_roles_kept(&wh).unwrap();
assert!(last_role_read_was_kept(), "our own append: still kept");
let fresh = read_roles(&wh).unwrap();
assert_eq!(after.len(), fresh.len());
for (k, v) in fresh.iter() {
assert_eq!((after[k].vector.clone(), after[k].seq, after[k].ts_micros), (v.vector.clone(), v.seq, v.ts_micros), "{k}");
}
// a foreign write: re-read
append_roles(&wh, &[role("b", 0.125, 3)]).unwrap();
let foreign = read_roles_kept(&wh).unwrap();
assert!(!last_role_read_was_kept(), "a foreign write moved the snapshot: scanned");
assert_eq!(foreign["b"].vector, vec![0.125; 4]);
forget_kept_roles();
}
/// RED before the fold: the role table only grew (every write appends
/// every changed vector). Folded, it holds the newest row per id — the
/// same answer from a fraction of the rows — and a table that moved
/// since the read is left alone.
#[test]
fn the_role_table_folds_to_its_newest_row_per_id_with_the_same_answer() {
let role = |id: &str, x: f32, seq: i64| RoleRow { repo: "r".into(), stable_id: id.into(), seq, feat_hash: 7, vector: vec![x; 4], ts_micros: 0 };
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
for seq in 1..=4 {
append_roles(&wh, &[role("a", seq as f32, seq), role("b", 0.5, seq), role("c", 0.25, seq)]).unwrap();
}
let (before, read) = read_roles_counted(&wh).unwrap();
assert_eq!((before.len(), read), (3, 12));
let snap = role_snapshot(&wh).unwrap();
// a table that moved since the read: left alone
assert!(!compact_roles(&wh, &before, Some(-1)).unwrap());
assert!(compact_roles(&wh, &before, snap).unwrap());
let (after, read) = read_roles_counted(&wh).unwrap();
assert_eq!(read, 3, "one row per id");
assert_eq!(after, before, "the same newest rows, stamps and seqs included");
assert_eq!(after["a"].vector, vec![4.0; 4]);
}
fn ev(seq: i64, ts: i64, m: &str, exact: Option<&str>) -> Event {
Event { repo: "r".into(), seq, sha: format!("s{seq}"), moniker: m.into(), shape: exact.map(|e| Shape { exact: e.into(), tokens: 50, minhash: vec![1; 64], simd: e == "z" }), ts_micros: ts }
}
#[test]
fn the_open_shapes_at_an_ordinal_are_the_fold_up_to_it_and_survive_a_round_trip() {
let events = vec![ev(1, 1, "a", Some("x")), ev(1, 1, "b", Some("y")), ev(2, 2, "a", Some("z")), ev(3, 3, "b", None)];
let at = |n| open(&events, "r", n).into_iter().map(|(m, s)| (m.to_string(), s.exact.clone())).collect::<Vec<_>>();
assert_eq!(at(1), [("a".into(), "x".into()), ("b".into(), "y".into())]);
assert_eq!(at(3), [("a".to_string(), "z".to_string())]);
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
append(&wh, &events).unwrap();
let back = read(&wh, &["r"]).unwrap();
assert_eq!(back.len(), 4);
let opened = open(&back, "r", 3);
assert_eq!(opened.len(), 1);
assert_eq!((opened["a"].minhash.clone(), opened["a"].simd), (vec![1; 64], true));
assert!(read(&wh, &["other"]).unwrap().is_empty());
// role vectors: the newest row per stable id wins
let role = |seq: i64, v: f32| RoleRow { repo: "r".into(), stable_id: "id".into(), seq, feat_hash: 0xabc, vector: vec![v; 4], ts_micros: 0 };
append_roles(&wh, &[role(1, 0.5)]).unwrap();
append_roles(&wh, &[role(2, 0.25)]).unwrap();
let roles = read_roles(&wh).unwrap();
assert_eq!((roles["id"].seq, roles["id"].feat_hash, roles["id"].vector.clone()), (2, 0xabc, vec![0.25; 4]));
}
}