src/vector/fnvec/links.rs
//! The two tables behind the per-function vectors, in the workspace's
//! warehouse catalog (the one that holds `embeddings`):
//!
//! | table | row |
//! |---|---|
//! | `embedding_node_events` | `repo · seq · sha · op (+/−) · moniker · content_hash · chunker · ts` |
//! | `embedding_node_runs` | one ledger row per run: scope, model space, backend, counts, status |
//!
//! Append-only. The state of a repo at ordinal N is a fold: per moniker the
//! last event with `seq ≤ N`, open when it is `+` ([`open`]).
use std::collections::BTreeMap;
use std::sync::Arc;
use anyhow::Result;
use arrow::array::{Array, Int64Array, 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::{At, CHUNKER, Run};
use crate::vector::store::col;
use crate::warehouse::iceberg::{IcebergWarehouse, TABLE_EMBEDDING_NODE_EVENTS, TABLE_EMBEDDING_NODE_RUNS, append_batch};
pub const OP_PLUS: &str = "+";
pub const OP_MINUS: &str = "-";
pub const STATUS_OK: &str = "ok";
pub const STATUS_SKIPPED: &str = "skipped";
pub const STATUS_ERROR: &str = "error";
/// One link event.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Event {
pub repo: String,
pub seq: i64,
pub sha: String,
pub plus: bool,
pub moniker: String,
/// `None` on a `-` event.
pub content_hash: Option<String>,
pub ts_micros: i64,
}
impl Event {
pub fn plus(at: &At, moniker: &str, content_hash: &str) -> Self {
Self { repo: at.repo.clone(), seq: at.seq, sha: at.sha.clone(), plus: true, moniker: moniker.to_string(), content_hash: Some(content_hash.to_string()), ts_micros: 0 }
}
pub fn minus(at: &At, moniker: &str) -> Self {
Self { repo: at.repo.clone(), seq: at.seq, sha: at.sha.clone(), plus: false, moniker: moniker.to_string(), content_hash: None, ts_micros: 0 }
}
}
/// One ledger row.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RunRow {
pub repo: String,
pub unit: String,
pub seq: i64,
pub sha: String,
pub model_profile: String,
pub model_id: String,
pub backend: String,
pub functions: i64,
pub distinct_bodies: i64,
pub new_vectors: i64,
pub plus: i64,
pub minus: i64,
pub without_chunk: i64,
pub ms: i64,
pub status: String,
pub detail: Option<String>,
pub ts_micros: i64,
}
impl RunRow {
pub fn of(at: &At, model_profile: &str, model_id: &str, run: &Run, status: &str, detail: Option<&str>) -> Self {
Self {
repo: at.repo.clone(),
unit: at.unit.clone(),
seq: at.seq,
sha: at.sha.clone(),
model_profile: model_profile.to_string(),
model_id: model_id.to_string(),
backend: run.backend.clone(),
functions: run.functions as i64,
distinct_bodies: run.distinct_bodies as i64,
new_vectors: run.new_vectors as i64,
plus: run.plus as i64,
minus: run.minus as i64,
without_chunk: run.without_chunk as i64,
ms: run.ms as i64,
status: status.to_string(),
detail: detail.map(str::to_string),
ts_micros: 0,
}
}
}
/// The open links of `repo` at ordinal `upto` (moniker → content hash): 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(events: &[Event], repo: &str, upto: i64) -> BTreeMap<String, String> {
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.content_hash, e.plus) {
(Some(h), true) => {
state.insert(e.moniker.clone(), h.clone());
}
_ => {
state.remove(&e.moniker);
}
}
}
state
}
fn strings<'a>(it: impl Iterator<Item = &'a str>) -> Arc<dyn Array> {
Arc::new(StringArray::from_iter_values(it))
}
fn now_micros() -> i64 {
chrono::Utc::now().timestamp_micros()
}
fn stamps(n: usize, ts: i64) -> Arc<dyn Array> {
Arc::new(TimestampMicrosecondArray::from(vec![ts; n]).with_timezone("+00:00"))
}
/// The table, or `None` when this warehouse was laid out before it existed
/// and is open read-only (a reader then has nothing to read: no run ever wrote).
async fn table(wh: &IcebergWarehouse, name: &str) -> Result<Option<iceberg::table::Table>> {
let ident = wh.table_ident(name);
if !wh.catalog().table_exists(&ident).await? {
return Ok(None);
}
Ok(Some(wh.catalog().load_table(&ident).await?))
}
/// Append link events (one commit). Nothing to write is no commit.
pub fn append_events(wh: &IcebergWarehouse, events: &[Event]) -> Result<()> {
if events.is_empty() {
return Ok(());
}
wh.block_on(async {
let _one_writer = crate::vector::store::VECTOR_WRITE.lock().await;
let table = wh.catalog().load_table(&wh.table_ident(TABLE_EMBEDDING_NODE_EVENTS)).await?;
let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?);
let cols: Vec<Arc<dyn Array>> = vec![
strings(events.iter().map(|e| e.repo.as_str())),
Arc::new(Int64Array::from_iter_values(events.iter().map(|e| e.seq))),
strings(events.iter().map(|e| e.sha.as_str())),
strings(events.iter().map(|e| if e.plus { OP_PLUS } else { OP_MINUS })),
strings(events.iter().map(|e| e.moniker.as_str())),
Arc::new(StringArray::from(events.iter().map(|e| e.content_hash.as_deref()).collect::<Vec<_>>())),
strings(events.iter().map(|_| CHUNKER)),
stamps(events.len(), now_micros()),
];
append_batch(wh.catalog(), table, RecordBatch::try_new(schema, cols)?).await
})
}
/// `repo = a OR repo = b …` (`None` = no repo named: every row).
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 link events of `repos` written under the current cut rule ([`CHUNKER`]).
pub fn read_events_of(wh: &IcebergWarehouse, repos: &[&str]) -> Result<Vec<Event>> {
wh.block_on(async {
let Some(table) = table(wh, TABLE_EMBEDDING_NODE_EVENTS).await? else { return Ok(Vec::new()) };
let mut scan = table.scan().select(["repo", "seq", "sha", "op", "moniker", "content_hash", "chunker", "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) = (col::<StringArray>(b, "repo")?, col::<Int64Array>(b, "seq")?, col::<StringArray>(b, "sha")?, col::<StringArray>(b, "op")?);
let (moniker, hash, chunker, ts) = (col::<StringArray>(b, "moniker")?, col::<StringArray>(b, "content_hash")?, col::<StringArray>(b, "chunker")?, col::<TimestampMicrosecondArray>(b, "ts_micros")?);
for i in 0..b.num_rows() {
if chunker.value(i) != CHUNKER || (!repos.is_empty() && !repos.contains(&repo.value(i))) {
continue;
}
out.push(Event {
repo: repo.value(i).to_string(),
seq: seq.value(i),
sha: sha.value(i).to_string(),
plus: op.value(i) == OP_PLUS,
moniker: moniker.value(i).to_string(),
content_hash: (!hash.is_null(i)).then(|| hash.value(i).to_string()),
ts_micros: ts.value(i),
});
}
}
Ok(out)
})
}
/// Append one ledger row.
pub fn append_run(wh: &IcebergWarehouse, r: &RunRow) -> Result<()> {
wh.block_on(async {
let _one_writer = crate::vector::store::VECTOR_WRITE.lock().await;
let table = wh.catalog().load_table(&wh.table_ident(TABLE_EMBEDDING_NODE_RUNS)).await?;
let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?);
let one = |s: &str| -> Arc<dyn Array> { Arc::new(StringArray::from(vec![s.to_string()])) };
let num = |n: i64| -> Arc<dyn Array> { Arc::new(Int64Array::from(vec![n])) };
let cols: Vec<Arc<dyn Array>> = vec![
one(&uuid::Uuid::new_v4().to_string()),
one(&r.repo),
one(&r.unit),
num(r.seq),
one(&r.sha),
one(CHUNKER),
one(&r.model_profile),
one(&r.model_id),
one(&r.backend),
num(r.functions),
num(r.distinct_bodies),
num(r.new_vectors),
num(r.plus),
num(r.minus),
num(r.without_chunk),
num(r.ms),
one(&r.status),
Arc::new(StringArray::from(vec![r.detail.clone()])),
stamps(1, now_micros()),
];
append_batch(wh.catalog(), table, RecordBatch::try_new(schema, cols)?).await
})
}
/// Every ledger row written under the current cut rule, oldest first.
pub fn read_runs(wh: &IcebergWarehouse) -> Result<Vec<RunRow>> {
wh.block_on(async {
let Some(table) = table(wh, TABLE_EMBEDDING_NODE_RUNS).await? else { return Ok(Vec::new()) };
let batches: Vec<RecordBatch> = table
.scan()
.select([
"repo", "unit", "seq", "sha", "chunker", "model_profile", "model_id", "backend", "functions", "distinct_bodies", "new_vectors", "plus", "minus",
"without_chunk", "ms", "status", "detail", "ts_micros",
])
.build()?.to_arrow().await?.try_collect().await?;
let mut out = Vec::new();
for b in &batches {
let s = |name: &str| col::<StringArray>(b, name);
let n = |name: &str| col::<Int64Array>(b, name);
let (repo, unit, sha, chunker, profile, model, backend, status, detail) =
(s("repo")?, s("unit")?, s("sha")?, s("chunker")?, s("model_profile")?, s("model_id")?, s("backend")?, s("status")?, s("detail")?);
let (seq, functions, distinct, new, plus, minus, without, ms) =
(n("seq")?, n("functions")?, n("distinct_bodies")?, n("new_vectors")?, n("plus")?, n("minus")?, n("without_chunk")?, n("ms")?);
let ts = col::<TimestampMicrosecondArray>(b, "ts_micros")?;
for i in 0..b.num_rows() {
if chunker.value(i) != CHUNKER {
continue;
}
out.push(RunRow {
repo: repo.value(i).to_string(),
unit: unit.value(i).to_string(),
seq: seq.value(i),
sha: sha.value(i).to_string(),
model_profile: profile.value(i).to_string(),
model_id: model.value(i).to_string(),
backend: backend.value(i).to_string(),
functions: functions.value(i),
distinct_bodies: distinct.value(i),
new_vectors: new.value(i),
plus: plus.value(i),
minus: minus.value(i),
without_chunk: without.value(i),
ms: ms.value(i),
status: status.value(i).to_string(),
detail: (!detail.is_null(i)).then(|| detail.value(i).to_string()),
ts_micros: ts.value(i),
});
}
}
out.sort_by_key(|r| r.ts_micros);
Ok(out)
})
}
#[cfg(test)]
mod tests {
use super::*;
fn ev(seq: i64, ts: i64, m: &str, hash: Option<&str>) -> Event {
Event { repo: "r".into(), seq, sha: format!("s{seq}"), plus: hash.is_some(), moniker: m.into(), content_hash: hash.map(str::to_string), ts_micros: ts }
}
/// The state at an ordinal is the fold up to it: a later change does not
/// reach back, a `-` closes the node, another repo's rows are not this one's.
#[test]
fn the_open_links_at_an_ordinal_are_the_fold_up_to_it() {
let mut events = vec![ev(1, 10, "a", Some("h1")), ev(1, 10, "b", Some("h2")), ev(2, 20, "a", Some("h3")), ev(3, 30, "b", None)];
events.push(Event { repo: "other".into(), ..ev(1, 5, "a", Some("zz")) });
let at = |n| open(&events, "r", n).into_iter().collect::<Vec<_>>();
assert_eq!(at(0), []);
assert_eq!(at(1), [("a".to_string(), "h1".to_string()), ("b".to_string(), "h2".to_string())]);
assert_eq!(at(2), [("a".to_string(), "h3".to_string()), ("b".to_string(), "h2".to_string())]);
assert_eq!(at(3), [("a".to_string(), "h3".to_string())]);
// two runs at ONE ordinal apply in the order they were written
let rerun = vec![ev(4, 50, "a", Some("new")), ev(4, 40, "a", Some("old"))];
assert_eq!(open(&rerun, "r", 4).get("a").map(String::as_str), Some("new"));
}
}