Compare commits
18
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9e6002c4d4 | ||
|
|
3bc82833ac | ||
|
|
fd4b6836ae | ||
|
|
59e72fa566 | ||
|
|
2592f8fcb3 | ||
|
|
17a0259cd0 | ||
|
|
7c1222d240 | ||
|
|
0d037cfac1 | ||
|
|
4e1c400f5d | ||
|
|
4d5c49090c | ||
|
|
a293502265 | ||
|
|
d87f2ccac6 | ||
|
|
09fd446c8d | ||
|
|
24f1249a62 | ||
|
|
650f5e97b6 | ||
|
|
97a376482c | ||
|
|
76f6fba5de | ||
|
|
dbe00d018f |
Generated
+3014
-369
File diff suppressed because it is too large
Load Diff
+4
-1
@@ -12,6 +12,8 @@ members = [
|
||||
"crates/journald",
|
||||
"crates/gateway",
|
||||
"crates/ui",
|
||||
"crates/lance-bench",
|
||||
"crates/vectord-lance",
|
||||
]
|
||||
|
||||
[workspace.dependencies]
|
||||
@@ -45,4 +47,5 @@ csv = "1"
|
||||
lopdf = "0.35"
|
||||
encoding_rs = "0.8"
|
||||
instant-distance = "0.6"
|
||||
tokio-postgres = { version = "0.7", features = ["with-serde_json-1", "with-chrono-0_4"] }
|
||||
tokio-postgres = { version = "0.7", features = ["with-serde_json-1", "with-chrono-0_4", "with-uuid-1"] }
|
||||
mysql_async = { version = "0.34", default-features = false, features = ["minimal"] }
|
||||
|
||||
@@ -134,4 +134,47 @@ impl AiClient {
|
||||
}
|
||||
resp.json().await.map_err(|e| format!("rerank parse error: {e}"))
|
||||
}
|
||||
|
||||
/// Force Ollama to unload the named model from VRAM (keep_alive=0).
|
||||
/// Used for predictable profile swaps — without this, Ollama holds a
|
||||
/// model for its configured TTL (default 5min) and the previous
|
||||
/// profile's model can linger in VRAM next to the new one.
|
||||
pub async fn unload_model(&self, model: &str) -> Result<serde_json::Value, String> {
|
||||
let resp = self.client
|
||||
.post(format!("{}/admin/unload", self.base_url))
|
||||
.json(&serde_json::json!({ "model": model }))
|
||||
.send().await
|
||||
.map_err(|e| format!("unload request failed: {e}"))?;
|
||||
if !resp.status().is_success() {
|
||||
let text = resp.text().await.unwrap_or_default();
|
||||
return Err(format!("unload error: {text}"));
|
||||
}
|
||||
resp.json().await.map_err(|e| format!("unload parse error: {e}"))
|
||||
}
|
||||
|
||||
/// Ask Ollama to load the named model into VRAM proactively. Makes
|
||||
/// the first real request after profile activation fast (no cold-load
|
||||
/// latency).
|
||||
pub async fn preload_model(&self, model: &str) -> Result<serde_json::Value, String> {
|
||||
let resp = self.client
|
||||
.post(format!("{}/admin/preload", self.base_url))
|
||||
.json(&serde_json::json!({ "model": model }))
|
||||
.send().await
|
||||
.map_err(|e| format!("preload request failed: {e}"))?;
|
||||
if !resp.status().is_success() {
|
||||
let text = resp.text().await.unwrap_or_default();
|
||||
return Err(format!("preload error: {text}"));
|
||||
}
|
||||
resp.json().await.map_err(|e| format!("preload parse error: {e}"))
|
||||
}
|
||||
|
||||
/// GPU + loaded-model snapshot from the sidecar. Combines nvidia-smi
|
||||
/// output (if available) with Ollama's /api/ps.
|
||||
pub async fn vram_snapshot(&self) -> Result<serde_json::Value, String> {
|
||||
let resp = self.client
|
||||
.get(format!("{}/admin/vram", self.base_url))
|
||||
.send().await
|
||||
.map_err(|e| format!("vram request failed: {e}"))?;
|
||||
resp.json().await.map_err(|e| format!("vram parse error: {e}"))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,9 +16,17 @@ pub fn router(client: AiClient) -> Router {
|
||||
.route("/embed", post(embed))
|
||||
.route("/generate", post(generate))
|
||||
.route("/rerank", post(rerank))
|
||||
.route("/vram", get(vram))
|
||||
.with_state(client)
|
||||
}
|
||||
|
||||
async fn vram(State(client): State<AiClient>) -> impl IntoResponse {
|
||||
match client.vram_snapshot().await {
|
||||
Ok(snap) => Ok(Json(snap)),
|
||||
Err(e) => Err((StatusCode::BAD_GATEWAY, e)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn health(State(client): State<AiClient>) -> impl IntoResponse {
|
||||
match client.health().await {
|
||||
Ok(info) => Ok(Json(info)),
|
||||
|
||||
@@ -15,5 +15,6 @@ bytes = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
uuid = { workspace = true }
|
||||
object_store = { workspace = true }
|
||||
arrow = { workspace = true }
|
||||
proto = { path = "../proto" }
|
||||
tonic = { workspace = true }
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod registry;
|
||||
pub mod service;
|
||||
pub mod grpc;
|
||||
pub mod tombstones;
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
use shared::types::{
|
||||
DatasetId, DatasetManifest, ObjectRef, SchemaFingerprint,
|
||||
ColumnMeta, Lineage, FreshnessContract, Sensitivity,
|
||||
AiView, ColumnMeta, DatasetId, DatasetManifest, FreshnessContract, Lineage, ModelProfile,
|
||||
ObjectRef, RefreshPolicy, SchemaFingerprint, Sensitivity, Tombstone,
|
||||
};
|
||||
|
||||
use crate::tombstones::TombstoneStore;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
@@ -9,6 +11,23 @@ use tokio::sync::RwLock;
|
||||
use storaged::ops;
|
||||
use object_store::ObjectStore;
|
||||
|
||||
/// Make a view name safe for use as an object storage key.
|
||||
/// Allows letters, digits, `_`, `-`, `.`. Anything else becomes `_`.
|
||||
fn sanitize_view_name(name: &str) -> String {
|
||||
name.chars()
|
||||
.map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.' { c } else { '_' })
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, serde::Serialize)]
|
||||
pub struct MigrateBucketsReport {
|
||||
pub refs_examined: usize,
|
||||
pub refs_renamed: usize, // legacy "data"/"local" → "primary"
|
||||
pub refs_stamped: usize, // empty → "primary"
|
||||
pub refs_unchanged: usize, // already canonical
|
||||
pub manifests_persisted: usize,
|
||||
}
|
||||
|
||||
/// Partial metadata update — only set fields are applied.
|
||||
#[derive(Debug, Clone, Default, serde::Deserialize)]
|
||||
pub struct MetadataUpdate {
|
||||
@@ -20,14 +39,23 @@ pub struct MetadataUpdate {
|
||||
pub lineage: Option<Lineage>,
|
||||
pub freshness: Option<FreshnessContract>,
|
||||
pub row_count: Option<u64>,
|
||||
// Phase C embedding freshness
|
||||
pub embedding_refresh_policy: Option<RefreshPolicy>,
|
||||
}
|
||||
|
||||
const MANIFEST_PREFIX: &str = "_catalog/manifests";
|
||||
const VIEW_PREFIX: &str = "_catalog/views";
|
||||
const PROFILE_PREFIX: &str = "_catalog/profiles";
|
||||
|
||||
/// In-memory dataset registry backed by manifest persistence in object storage.
|
||||
/// Also tracks AiViews (Phase D) — safe projections over base datasets.
|
||||
/// And tombstones (Phase E) — soft-delete markers applied at query time.
|
||||
#[derive(Clone)]
|
||||
pub struct Registry {
|
||||
datasets: Arc<RwLock<HashMap<DatasetId, DatasetManifest>>>,
|
||||
views: Arc<RwLock<HashMap<String, AiView>>>,
|
||||
profiles: Arc<RwLock<HashMap<String, ModelProfile>>>,
|
||||
tombstones: TombstoneStore,
|
||||
store: Arc<dyn ObjectStore>,
|
||||
}
|
||||
|
||||
@@ -35,11 +63,47 @@ impl Registry {
|
||||
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||
Self {
|
||||
datasets: Arc::new(RwLock::new(HashMap::new())),
|
||||
views: Arc::new(RwLock::new(HashMap::new())),
|
||||
profiles: Arc::new(RwLock::new(HashMap::new())),
|
||||
tombstones: TombstoneStore::new(store.clone()),
|
||||
store,
|
||||
}
|
||||
}
|
||||
|
||||
/// Rebuild in-memory index from persisted manifests on startup.
|
||||
pub fn tombstones(&self) -> &TombstoneStore {
|
||||
&self.tombstones
|
||||
}
|
||||
|
||||
/// Add a tombstone for a dataset row. Validates that the dataset
|
||||
/// exists. Idempotent at the endpoint layer — callers dedup if needed.
|
||||
pub async fn add_tombstone(
|
||||
&self,
|
||||
dataset: &str,
|
||||
row_key_column: &str,
|
||||
row_key_value: &str,
|
||||
actor: &str,
|
||||
reason: &str,
|
||||
) -> Result<Tombstone, String> {
|
||||
if self.get_by_name(dataset).await.is_none() {
|
||||
return Err(format!("dataset not found: {dataset}"));
|
||||
}
|
||||
let ts = Tombstone {
|
||||
dataset: dataset.to_string(),
|
||||
row_key_column: row_key_column.to_string(),
|
||||
row_key_value: row_key_value.to_string(),
|
||||
deleted_at: chrono::Utc::now(),
|
||||
actor: actor.to_string(),
|
||||
reason: reason.to_string(),
|
||||
};
|
||||
self.tombstones.append(&ts).await?;
|
||||
Ok(ts)
|
||||
}
|
||||
|
||||
pub async fn list_tombstones(&self, dataset: &str) -> Result<Vec<Tombstone>, String> {
|
||||
self.tombstones.list(dataset).await
|
||||
}
|
||||
|
||||
/// Rebuild in-memory index from persisted manifests + views on startup.
|
||||
pub async fn rebuild(&self) -> Result<usize, String> {
|
||||
let keys = ops::list(&self.store, Some(MANIFEST_PREFIX)).await?;
|
||||
let mut datasets = self.datasets.write().await;
|
||||
@@ -52,9 +116,102 @@ impl Registry {
|
||||
}
|
||||
let count = datasets.len();
|
||||
tracing::info!("catalog rebuilt: {count} datasets loaded");
|
||||
|
||||
// Phase D: load AiView definitions alongside manifests.
|
||||
let view_keys = ops::list(&self.store, Some(VIEW_PREFIX)).await.unwrap_or_default();
|
||||
let mut views = self.views.write().await;
|
||||
views.clear();
|
||||
for key in &view_keys {
|
||||
if !key.ends_with(".json") { continue; }
|
||||
let data = match ops::get(&self.store, key).await {
|
||||
Ok(d) => d,
|
||||
Err(e) => {
|
||||
tracing::warn!("view '{key}': read failed: {e}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
match serde_json::from_slice::<AiView>(&data) {
|
||||
Ok(view) => { views.insert(view.name.clone(), view); }
|
||||
Err(e) => tracing::warn!("view '{key}': parse failed: {e}"),
|
||||
}
|
||||
}
|
||||
if !views.is_empty() {
|
||||
tracing::info!("catalog: {} views loaded", views.len());
|
||||
}
|
||||
|
||||
// Phase 17: load model profiles.
|
||||
let profile_keys = ops::list(&self.store, Some(PROFILE_PREFIX)).await.unwrap_or_default();
|
||||
let mut profiles = self.profiles.write().await;
|
||||
profiles.clear();
|
||||
for key in &profile_keys {
|
||||
if !key.ends_with(".json") { continue; }
|
||||
let data = match ops::get(&self.store, key).await {
|
||||
Ok(d) => d,
|
||||
Err(e) => { tracing::warn!("profile '{key}': read failed: {e}"); continue; }
|
||||
};
|
||||
match serde_json::from_slice::<ModelProfile>(&data) {
|
||||
Ok(p) => { profiles.insert(p.id.clone(), p); }
|
||||
Err(e) => tracing::warn!("profile '{key}': parse failed: {e}"),
|
||||
}
|
||||
}
|
||||
if !profiles.is_empty() {
|
||||
tracing::info!("catalog: {} model profiles loaded", profiles.len());
|
||||
}
|
||||
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
// --- Phase 17: Model profiles ---
|
||||
|
||||
/// Create or replace a model profile. Validates id slug + ensures
|
||||
/// every bound_dataset exists (as either a raw dataset or an AiView).
|
||||
pub async fn put_profile(&self, profile: ModelProfile) -> Result<ModelProfile, String> {
|
||||
if profile.id.is_empty() {
|
||||
return Err("profile id is empty".into());
|
||||
}
|
||||
if !profile.id.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_') {
|
||||
return Err(format!(
|
||||
"profile id '{}' must be alphanumeric + '-' or '_' only", profile.id,
|
||||
));
|
||||
}
|
||||
for binding in &profile.bound_datasets {
|
||||
let exists = self.get_by_name(binding).await.is_some()
|
||||
|| self.get_view(binding).await.is_some();
|
||||
if !exists {
|
||||
return Err(format!(
|
||||
"bound dataset '{}' not found as dataset or view", binding,
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
let key = format!("{PROFILE_PREFIX}/{}.json", profile.id);
|
||||
let json = serde_json::to_vec_pretty(&profile).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &key, json.into()).await?;
|
||||
|
||||
let mut profiles = self.profiles.write().await;
|
||||
profiles.insert(profile.id.clone(), profile.clone());
|
||||
tracing::info!(
|
||||
"profile registered: {} -> ollama={} bindings={:?}",
|
||||
profile.id, profile.ollama_name, profile.bound_datasets,
|
||||
);
|
||||
Ok(profile)
|
||||
}
|
||||
|
||||
pub async fn get_profile(&self, id: &str) -> Option<ModelProfile> {
|
||||
self.profiles.read().await.get(id).cloned()
|
||||
}
|
||||
|
||||
pub async fn list_profiles(&self) -> Vec<ModelProfile> {
|
||||
self.profiles.read().await.values().cloned().collect()
|
||||
}
|
||||
|
||||
pub async fn delete_profile(&self, id: &str) -> Result<(), String> {
|
||||
let key = format!("{PROFILE_PREFIX}/{id}.json");
|
||||
ops::delete(&self.store, &key).await?;
|
||||
self.profiles.write().await.remove(id);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Register a new dataset. Persists manifest to storage before updating memory.
|
||||
pub async fn register(
|
||||
&self,
|
||||
@@ -78,6 +235,9 @@ impl Registry {
|
||||
freshness: None,
|
||||
tags: vec![],
|
||||
row_count: None,
|
||||
last_embedded_at: None,
|
||||
embedding_stale_since: None,
|
||||
embedding_refresh_policy: None,
|
||||
};
|
||||
|
||||
// Write-ahead: persist before in-memory update
|
||||
@@ -111,6 +271,7 @@ impl Registry {
|
||||
if let Some(lineage) = updates.lineage { manifest.lineage = Some(lineage); }
|
||||
if let Some(freshness) = updates.freshness { manifest.freshness = Some(freshness); }
|
||||
if let Some(count) = updates.row_count { manifest.row_count = Some(count); }
|
||||
if let Some(policy) = updates.embedding_refresh_policy { manifest.embedding_refresh_policy = Some(policy); }
|
||||
manifest.updated_at = chrono::Utc::now();
|
||||
|
||||
// Persist
|
||||
@@ -140,6 +301,264 @@ impl Registry {
|
||||
datasets.values().cloned().collect()
|
||||
}
|
||||
|
||||
/// Re-read the parquet footer(s) for a dataset and repopulate `row_count`
|
||||
/// and `columns` from reality. Use this to repair manifests whose
|
||||
/// metadata was lost (e.g. migrated from a pre-Phase 10 catalog).
|
||||
///
|
||||
/// Does NOT touch owner/description/sensitivity/lineage/tags — only
|
||||
/// the structural facts that parquet can tell us authoritatively.
|
||||
/// The existing `schema_fingerprint` is updated if the recomputed one
|
||||
/// differs; a warning is logged so drift is visible.
|
||||
pub async fn resync_from_parquet(&self, name: &str) -> Result<DatasetManifest, String> {
|
||||
use shared::arrow_helpers::{fingerprint_schema, parquet_to_record_batches};
|
||||
|
||||
// Snapshot the target manifest so we don't hold the write lock during IO.
|
||||
let (id, objects, old_fp) = {
|
||||
let datasets = self.datasets.read().await;
|
||||
let m = datasets
|
||||
.values()
|
||||
.find(|d| d.name == name)
|
||||
.ok_or_else(|| format!("dataset not found: {name}"))?;
|
||||
(m.id.clone(), m.objects.clone(), m.schema_fingerprint.clone())
|
||||
};
|
||||
|
||||
if objects.is_empty() {
|
||||
return Err(format!("dataset '{name}' has no object references to resync from"));
|
||||
}
|
||||
|
||||
let mut total_rows: u64 = 0;
|
||||
let mut first_schema: Option<arrow::datatypes::SchemaRef> = None;
|
||||
for obj in &objects {
|
||||
let data = ops::get(&self.store, &obj.key).await
|
||||
.map_err(|e| format!("read {}: {e}", obj.key))?;
|
||||
let (schema, batches) = parquet_to_record_batches(&data)
|
||||
.map_err(|e| format!("parse {}: {e}", obj.key))?;
|
||||
let rows: u64 = batches.iter().map(|b| b.num_rows() as u64).sum();
|
||||
total_rows += rows;
|
||||
if first_schema.is_none() {
|
||||
first_schema = Some(schema);
|
||||
}
|
||||
}
|
||||
|
||||
let schema = first_schema.ok_or("no schema recovered")?;
|
||||
let new_fp = fingerprint_schema(&schema);
|
||||
if new_fp != old_fp {
|
||||
tracing::warn!(
|
||||
"dataset '{}' schema fingerprint drift: {} -> {} (updating to match parquet reality)",
|
||||
name, old_fp.0, new_fp.0,
|
||||
);
|
||||
}
|
||||
|
||||
let columns: Vec<ColumnMeta> = schema
|
||||
.fields()
|
||||
.iter()
|
||||
.map(|f| ColumnMeta {
|
||||
name: f.name().clone(),
|
||||
data_type: f.data_type().to_string(),
|
||||
sensitivity: None,
|
||||
description: String::new(),
|
||||
is_pii: false,
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Apply updates.
|
||||
let mut datasets = self.datasets.write().await;
|
||||
let manifest = datasets
|
||||
.get_mut(&id)
|
||||
.ok_or_else(|| format!("dataset disappeared during resync: {name}"))?;
|
||||
manifest.row_count = Some(total_rows);
|
||||
manifest.columns = columns;
|
||||
manifest.schema_fingerprint = new_fp;
|
||||
manifest.updated_at = chrono::Utc::now();
|
||||
|
||||
// Persist.
|
||||
let manifest_key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
|
||||
let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &manifest_key, json.into()).await?;
|
||||
|
||||
tracing::info!("resynced '{name}': row_count={total_rows}, {} columns", manifest.columns.len());
|
||||
Ok(manifest.clone())
|
||||
}
|
||||
|
||||
/// Resync every dataset that currently has a null row_count.
|
||||
/// Returns (successes, failures) where each entry is (name, detail).
|
||||
pub async fn resync_missing(&self) -> (Vec<(String, u64)>, Vec<(String, String)>) {
|
||||
let names: Vec<String> = {
|
||||
let datasets = self.datasets.read().await;
|
||||
datasets
|
||||
.values()
|
||||
.filter(|d| d.row_count.is_none() || d.columns.is_empty())
|
||||
.map(|d| d.name.clone())
|
||||
.collect()
|
||||
};
|
||||
|
||||
let mut ok = Vec::new();
|
||||
let mut err = Vec::new();
|
||||
for name in names {
|
||||
match self.resync_from_parquet(&name).await {
|
||||
Ok(m) => ok.push((name, m.row_count.unwrap_or(0))),
|
||||
Err(e) => err.push((name, e)),
|
||||
}
|
||||
}
|
||||
(ok, err)
|
||||
}
|
||||
|
||||
/// Mark a dataset's embeddings as stale (row-level data has been written
|
||||
/// since the last embedding refresh). Idempotent — setting stale when
|
||||
/// already stale is a no-op. Only marks stale if the dataset has been
|
||||
/// embedded before — a never-embedded dataset doesn't need a stale flag
|
||||
/// (it just needs an initial index build). Called from the ingest path.
|
||||
pub async fn mark_embeddings_stale(&self, name: &str) -> Result<(), String> {
|
||||
let mut datasets = self.datasets.write().await;
|
||||
let manifest = datasets
|
||||
.values_mut()
|
||||
.find(|d| d.name == name)
|
||||
.ok_or_else(|| format!("dataset not found: {name}"))?;
|
||||
|
||||
if manifest.last_embedded_at.is_none() {
|
||||
return Ok(()); // never embedded -> no stale semantics yet
|
||||
}
|
||||
if manifest.embedding_stale_since.is_none() {
|
||||
manifest.embedding_stale_since = Some(chrono::Utc::now());
|
||||
manifest.updated_at = chrono::Utc::now();
|
||||
|
||||
let key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
|
||||
let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &key, json.into()).await?;
|
||||
tracing::info!("marked embeddings stale for dataset '{name}'");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Clear the stale marker and set `last_embedded_at = now`.
|
||||
/// Called by the embedding refresh pipeline once it finishes.
|
||||
pub async fn clear_embeddings_stale(&self, name: &str) -> Result<(), String> {
|
||||
let mut datasets = self.datasets.write().await;
|
||||
let manifest = datasets
|
||||
.values_mut()
|
||||
.find(|d| d.name == name)
|
||||
.ok_or_else(|| format!("dataset not found: {name}"))?;
|
||||
|
||||
let now = chrono::Utc::now();
|
||||
manifest.embedding_stale_since = None;
|
||||
manifest.last_embedded_at = Some(now);
|
||||
manifest.updated_at = now;
|
||||
|
||||
let key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
|
||||
let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &key, json.into()).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Federation layer 2: stamp `bucket = "primary"` on every ObjectRef
|
||||
/// whose `bucket` field is empty or matches a legacy value (`"data"`,
|
||||
/// `"local"`). One-shot migration; re-running is a safe no-op once
|
||||
/// every ref is canonical.
|
||||
pub async fn migrate_buckets_to_primary(&self) -> Result<MigrateBucketsReport, String> {
|
||||
let mut datasets = self.datasets.write().await;
|
||||
let mut report = MigrateBucketsReport::default();
|
||||
let mut to_persist: Vec<(DatasetId, DatasetManifest)> = Vec::new();
|
||||
|
||||
for manifest in datasets.values_mut() {
|
||||
let mut changed = false;
|
||||
for obj in manifest.objects.iter_mut() {
|
||||
report.refs_examined += 1;
|
||||
if obj.bucket.is_empty() {
|
||||
obj.bucket = "primary".to_string();
|
||||
report.refs_stamped += 1;
|
||||
changed = true;
|
||||
} else if obj.bucket == "data" || obj.bucket == "local" {
|
||||
obj.bucket = "primary".to_string();
|
||||
report.refs_renamed += 1;
|
||||
changed = true;
|
||||
} else {
|
||||
report.refs_unchanged += 1;
|
||||
}
|
||||
}
|
||||
if changed {
|
||||
manifest.updated_at = chrono::Utc::now();
|
||||
to_persist.push((manifest.id.clone(), manifest.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
// Persist updated manifests after we've finished mutating the map.
|
||||
for (id, manifest) in to_persist {
|
||||
let key = format!("{MANIFEST_PREFIX}/{}.json", id);
|
||||
let json = serde_json::to_vec_pretty(&manifest).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &key, json.into()).await?;
|
||||
report.manifests_persisted += 1;
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
"bucket migration: examined {} refs, renamed {}, stamped {}, unchanged {}, persisted {} manifests",
|
||||
report.refs_examined,
|
||||
report.refs_renamed,
|
||||
report.refs_stamped,
|
||||
report.refs_unchanged,
|
||||
report.manifests_persisted,
|
||||
);
|
||||
Ok(report)
|
||||
}
|
||||
|
||||
// --- Phase D: AI-safe views ---
|
||||
|
||||
/// Create or replace a named view. Validates that the base dataset
|
||||
/// exists and that the column whitelist is non-empty. Persists to
|
||||
/// `_catalog/views/{name}.json`.
|
||||
pub async fn put_view(&self, mut view: AiView) -> Result<AiView, String> {
|
||||
if view.name.is_empty() {
|
||||
return Err("view name is empty".into());
|
||||
}
|
||||
if view.columns.is_empty() {
|
||||
return Err("view must whitelist at least one column".into());
|
||||
}
|
||||
// Base dataset must exist (read-side will fail anyway, fail-fast here
|
||||
// so operators don't end up with dangling views).
|
||||
if self.get_by_name(&view.base_dataset).await.is_none() {
|
||||
return Err(format!(
|
||||
"base dataset '{}' not found in catalog",
|
||||
view.base_dataset
|
||||
));
|
||||
}
|
||||
if view.created_at.timestamp() == 0 {
|
||||
view.created_at = chrono::Utc::now();
|
||||
}
|
||||
|
||||
let key = format!("{VIEW_PREFIX}/{}.json", sanitize_view_name(&view.name));
|
||||
let json = serde_json::to_vec_pretty(&view).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &key, json.into()).await?;
|
||||
|
||||
let mut views = self.views.write().await;
|
||||
views.insert(view.name.clone(), view.clone());
|
||||
tracing::info!("view registered: {} over '{}'", view.name, view.base_dataset);
|
||||
Ok(view)
|
||||
}
|
||||
|
||||
pub async fn get_view(&self, name: &str) -> Option<AiView> {
|
||||
self.views.read().await.get(name).cloned()
|
||||
}
|
||||
|
||||
pub async fn list_views(&self) -> Vec<AiView> {
|
||||
self.views.read().await.values().cloned().collect()
|
||||
}
|
||||
|
||||
pub async fn delete_view(&self, name: &str) -> Result<(), String> {
|
||||
let key = format!("{VIEW_PREFIX}/{}.json", sanitize_view_name(name));
|
||||
ops::delete(&self.store, &key).await?;
|
||||
self.views.write().await.remove(name);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// List datasets whose `embedding_stale_since` is set — they need a refresh.
|
||||
pub async fn stale_datasets(&self) -> Vec<DatasetManifest> {
|
||||
let datasets = self.datasets.read().await;
|
||||
datasets
|
||||
.values()
|
||||
.filter(|d| d.embedding_stale_since.is_some())
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Add objects to an existing dataset.
|
||||
pub async fn add_objects(
|
||||
&self,
|
||||
|
||||
@@ -19,6 +19,17 @@ pub fn router(registry: Registry) -> Router {
|
||||
.route("/datasets/{id}", get(get_dataset))
|
||||
.route("/datasets/by-name/{name}", get(get_dataset_by_name))
|
||||
.route("/datasets/by-name/{name}/metadata", post(update_metadata))
|
||||
.route("/datasets/by-name/{name}/resync", post(resync_dataset))
|
||||
.route("/resync-missing", post(resync_all_missing))
|
||||
.route("/migrate-buckets", post(migrate_buckets))
|
||||
// Phase D: AI-safe views
|
||||
.route("/views", post(create_view).get(list_views))
|
||||
.route("/views/{name}", get(get_view).delete(delete_view))
|
||||
// Phase E: soft-delete tombstones
|
||||
.route("/datasets/by-name/{name}/tombstone", post(tombstone_rows).get(list_tombstones))
|
||||
// Phase 17: model profiles
|
||||
.route("/profiles", post(create_profile).get(list_profiles))
|
||||
.route("/profiles/{id}", get(get_profile).delete(delete_profile))
|
||||
.with_state(registry)
|
||||
}
|
||||
|
||||
@@ -152,3 +163,254 @@ async fn update_metadata(
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Re-read parquet footers for a single dataset and repopulate row_count
|
||||
/// and columns from reality. Useful for repairing manifests whose metadata
|
||||
/// was lost or never backfilled.
|
||||
async fn resync_dataset(
|
||||
State(registry): State<Registry>,
|
||||
Path(name): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
match registry.resync_from_parquet(&name).await {
|
||||
Ok(manifest) => Ok(Json(DatasetResponse::from(&manifest))),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ResyncAllResponse {
|
||||
succeeded: Vec<ResyncOk>,
|
||||
failed: Vec<ResyncErr>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ResyncOk {
|
||||
name: String,
|
||||
row_count: u64,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ResyncErr {
|
||||
name: String,
|
||||
error: String,
|
||||
}
|
||||
|
||||
/// Resync every dataset that currently has null row_count or empty columns.
|
||||
async fn resync_all_missing(State(registry): State<Registry>) -> impl IntoResponse {
|
||||
let (ok, err) = registry.resync_missing().await;
|
||||
Json(ResyncAllResponse {
|
||||
succeeded: ok.into_iter().map(|(name, row_count)| ResyncOk { name, row_count }).collect(),
|
||||
failed: err.into_iter().map(|(name, error)| ResyncErr { name, error }).collect(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Federation layer 2 one-shot: normalize every ObjectRef.bucket field
|
||||
/// to the canonical "primary" value. Idempotent — re-running once
|
||||
/// everything is canonical is a safe no-op.
|
||||
async fn migrate_buckets(State(registry): State<Registry>) -> impl IntoResponse {
|
||||
match registry.migrate_buckets_to_primary().await {
|
||||
Ok(report) => Ok(Json(report)),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
// --- Phase D: AI-safe views ---
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct CreateViewRequest {
|
||||
name: String,
|
||||
base_dataset: String,
|
||||
columns: Vec<String>,
|
||||
#[serde(default)]
|
||||
row_filter: Option<String>,
|
||||
#[serde(default)]
|
||||
column_redactions: std::collections::HashMap<String, shared::types::Redaction>,
|
||||
#[serde(default)]
|
||||
description: String,
|
||||
#[serde(default)]
|
||||
created_by: String,
|
||||
}
|
||||
|
||||
async fn create_view(
|
||||
State(registry): State<Registry>,
|
||||
Json(req): Json<CreateViewRequest>,
|
||||
) -> impl IntoResponse {
|
||||
let view = shared::types::AiView {
|
||||
name: req.name,
|
||||
base_dataset: req.base_dataset,
|
||||
columns: req.columns,
|
||||
row_filter: req.row_filter,
|
||||
column_redactions: req.column_redactions,
|
||||
created_at: chrono::Utc::now(),
|
||||
created_by: req.created_by,
|
||||
description: req.description,
|
||||
};
|
||||
match registry.put_view(view).await {
|
||||
Ok(v) => Ok((StatusCode::CREATED, Json(v))),
|
||||
Err(e) => Err((StatusCode::BAD_REQUEST, e)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_views(State(registry): State<Registry>) -> impl IntoResponse {
|
||||
Json(registry.list_views().await)
|
||||
}
|
||||
|
||||
async fn get_view(
|
||||
State(registry): State<Registry>,
|
||||
Path(name): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
match registry.get_view(&name).await {
|
||||
Some(v) => Ok(Json(v)),
|
||||
None => Err((StatusCode::NOT_FOUND, format!("view not found: {name}"))),
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_view(
|
||||
State(registry): State<Registry>,
|
||||
Path(name): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
match registry.delete_view(&name).await {
|
||||
Ok(()) => Ok(StatusCode::NO_CONTENT),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
// --- Phase E: soft-delete tombstones ---
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct TombstoneRequest {
|
||||
row_key_column: String,
|
||||
row_key_values: Vec<String>,
|
||||
#[serde(default)]
|
||||
actor: String,
|
||||
#[serde(default)]
|
||||
reason: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct TombstoneResponse {
|
||||
dataset: String,
|
||||
row_key_column: String,
|
||||
rows_tombstoned: usize,
|
||||
failures: Vec<String>,
|
||||
}
|
||||
|
||||
async fn tombstone_rows(
|
||||
State(registry): State<Registry>,
|
||||
Path(name): Path<String>,
|
||||
Json(req): Json<TombstoneRequest>,
|
||||
) -> impl IntoResponse {
|
||||
if req.row_key_values.is_empty() {
|
||||
return Err((StatusCode::BAD_REQUEST, "row_key_values is empty".to_string()));
|
||||
}
|
||||
|
||||
let mut ok = 0;
|
||||
let mut failures = Vec::new();
|
||||
for value in &req.row_key_values {
|
||||
match registry
|
||||
.add_tombstone(&name, &req.row_key_column, value, &req.actor, &req.reason)
|
||||
.await
|
||||
{
|
||||
Ok(_) => ok += 1,
|
||||
Err(e) => failures.push(format!("{value}: {e}")),
|
||||
}
|
||||
}
|
||||
|
||||
let status = if ok > 0 && failures.is_empty() {
|
||||
StatusCode::CREATED
|
||||
} else if ok > 0 {
|
||||
StatusCode::MULTI_STATUS
|
||||
} else {
|
||||
StatusCode::BAD_REQUEST
|
||||
};
|
||||
|
||||
Ok((status, Json(TombstoneResponse {
|
||||
dataset: name,
|
||||
row_key_column: req.row_key_column,
|
||||
rows_tombstoned: ok,
|
||||
failures,
|
||||
})))
|
||||
}
|
||||
|
||||
async fn list_tombstones(
|
||||
State(registry): State<Registry>,
|
||||
Path(name): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
match registry.list_tombstones(&name).await {
|
||||
Ok(ts) => Ok(Json(ts)),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
// --- Phase 17: Model profiles ---
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct CreateProfileRequest {
|
||||
id: String,
|
||||
ollama_name: String,
|
||||
#[serde(default)]
|
||||
description: String,
|
||||
bound_datasets: Vec<String>,
|
||||
#[serde(default)]
|
||||
hnsw_config: shared::types::ProfileHnswConfig,
|
||||
#[serde(default = "default_embed_model_req")]
|
||||
embed_model: String,
|
||||
#[serde(default)]
|
||||
created_by: String,
|
||||
/// Federation: optional per-profile bucket (`profile:{id}` by convention).
|
||||
/// Omitting keeps artifacts in primary.
|
||||
#[serde(default)]
|
||||
bucket: Option<String>,
|
||||
/// ADR-019 hybrid: which vector backend to route this profile's
|
||||
/// indexes to. Defaults to Parquet+HNSW.
|
||||
#[serde(default)]
|
||||
vector_backend: shared::types::VectorBackend,
|
||||
}
|
||||
|
||||
fn default_embed_model_req() -> String { "nomic-embed-text".to_string() }
|
||||
|
||||
async fn create_profile(
|
||||
State(registry): State<Registry>,
|
||||
Json(req): Json<CreateProfileRequest>,
|
||||
) -> impl IntoResponse {
|
||||
let profile = shared::types::ModelProfile {
|
||||
id: req.id,
|
||||
ollama_name: req.ollama_name,
|
||||
description: req.description,
|
||||
bound_datasets: req.bound_datasets,
|
||||
hnsw_config: req.hnsw_config,
|
||||
embed_model: req.embed_model,
|
||||
created_at: chrono::Utc::now(),
|
||||
created_by: req.created_by,
|
||||
bucket: req.bucket,
|
||||
vector_backend: req.vector_backend,
|
||||
};
|
||||
match registry.put_profile(profile).await {
|
||||
Ok(p) => Ok((StatusCode::CREATED, Json(p))),
|
||||
Err(e) => Err((StatusCode::BAD_REQUEST, e)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_profiles(State(registry): State<Registry>) -> impl IntoResponse {
|
||||
Json(registry.list_profiles().await)
|
||||
}
|
||||
|
||||
async fn get_profile(
|
||||
State(registry): State<Registry>,
|
||||
Path(id): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
match registry.get_profile(&id).await {
|
||||
Some(p) => Ok(Json(p)),
|
||||
None => Err((StatusCode::NOT_FOUND, format!("profile not found: {id}"))),
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_profile(
|
||||
State(registry): State<Registry>,
|
||||
Path(id): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
match registry.delete_profile(&id).await {
|
||||
Ok(()) => Ok(StatusCode::NO_CONTENT),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,151 @@
|
||||
//! Soft-delete tombstone storage (Phase E).
|
||||
//!
|
||||
//! One append-log per dataset at `_catalog/tombstones/{dataset}/batch_*.jsonl`.
|
||||
//! Uses the shared `storaged::append_log::AppendLog` pattern so appends
|
||||
//! are write-once (never rewrites existing files) and can be compacted.
|
||||
//!
|
||||
//! The store exposes a per-dataset cache of active tombstones for the
|
||||
//! hot path (queryd filter construction) — that's why events are pulled
|
||||
//! into an in-memory map on every call rather than scanning object
|
||||
//! storage repeatedly.
|
||||
|
||||
use object_store::ObjectStore;
|
||||
use shared::types::Tombstone;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use storaged::append_log::{AppendLog, CompactStats};
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
const TOMBSTONE_PREFIX: &str = "_catalog/tombstones";
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct TombstoneStore {
|
||||
store: Arc<dyn ObjectStore>,
|
||||
logs: Arc<RwLock<HashMap<String, Arc<AppendLog>>>>,
|
||||
}
|
||||
|
||||
impl TombstoneStore {
|
||||
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||
Self {
|
||||
store,
|
||||
logs: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
fn prefix_for(dataset: &str) -> String {
|
||||
// Sanitize dataset name for filesystem safety.
|
||||
let safe: String = dataset
|
||||
.chars()
|
||||
.map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' { c } else { '_' })
|
||||
.collect();
|
||||
format!("{TOMBSTONE_PREFIX}/{}", safe)
|
||||
}
|
||||
|
||||
async fn log_for(&self, dataset: &str) -> Arc<AppendLog> {
|
||||
if let Some(log) = self.logs.read().await.get(dataset) {
|
||||
return log.clone();
|
||||
}
|
||||
let mut guard = self.logs.write().await;
|
||||
if let Some(log) = guard.get(dataset) {
|
||||
return log.clone();
|
||||
}
|
||||
// Threshold of 1 — every tombstone is high-value. Compliance/audit
|
||||
// doesn't tolerate "lost on restart"; we trade a small file count
|
||||
// for guaranteed durability. Compaction merges later if volume grows.
|
||||
let log = Arc::new(
|
||||
AppendLog::new(self.store.clone(), Self::prefix_for(dataset))
|
||||
.with_flush_threshold(1),
|
||||
);
|
||||
guard.insert(dataset.to_string(), log.clone());
|
||||
log
|
||||
}
|
||||
|
||||
/// Append one tombstone. Validates that the `row_key_column` matches
|
||||
/// the column already used for this dataset (all tombstones for a
|
||||
/// dataset share one key column so the query filter is well-defined).
|
||||
/// Forces a flush so the tombstone is durable before this call returns.
|
||||
pub async fn append(&self, ts: &Tombstone) -> Result<(), String> {
|
||||
let existing = self.list(&ts.dataset).await?;
|
||||
if let Some(prior) = existing.first() {
|
||||
if prior.row_key_column != ts.row_key_column {
|
||||
return Err(format!(
|
||||
"dataset '{}' already uses '{}' as tombstone key; cannot mix with '{}'",
|
||||
ts.dataset, prior.row_key_column, ts.row_key_column,
|
||||
));
|
||||
}
|
||||
}
|
||||
let line = serde_json::to_vec(ts).map_err(|e| e.to_string())?;
|
||||
let log = self.log_for(&ts.dataset).await;
|
||||
log.append(line).await?;
|
||||
// Belt-and-suspenders: explicit flush in case the threshold is
|
||||
// ever raised. Tombstones must be durable on return.
|
||||
log.flush().await
|
||||
}
|
||||
|
||||
/// All tombstones for a dataset (chronological).
|
||||
pub async fn list(&self, dataset: &str) -> Result<Vec<Tombstone>, String> {
|
||||
let log = self.log_for(dataset).await;
|
||||
let lines = log.read_all().await?;
|
||||
let mut out = Vec::with_capacity(lines.len());
|
||||
for line in lines {
|
||||
match serde_json::from_slice::<Tombstone>(&line) {
|
||||
Ok(t) => out.push(t),
|
||||
Err(e) => tracing::warn!("tombstones/{}: skip malformed entry: {e}", dataset),
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Per-dataset grouped view used by queryd — returns a map of
|
||||
/// `{dataset -> (row_key_column, set_of_values)}` for every dataset
|
||||
/// that has any tombstones.
|
||||
pub async fn all_grouped(
|
||||
&self,
|
||||
datasets: &[String],
|
||||
) -> Result<HashMap<String, (String, Vec<String>)>, String> {
|
||||
let mut grouped = HashMap::new();
|
||||
for dataset in datasets {
|
||||
let ts = match self.list(dataset).await {
|
||||
Ok(ts) => ts,
|
||||
Err(_) => continue,
|
||||
};
|
||||
if ts.is_empty() { continue; }
|
||||
let col = ts[0].row_key_column.clone();
|
||||
let values: Vec<String> = ts.iter().map(|t| t.row_key_value.clone()).collect();
|
||||
grouped.insert(dataset.clone(), (col, values));
|
||||
}
|
||||
Ok(grouped)
|
||||
}
|
||||
|
||||
pub async fn compact(&self, dataset: &str) -> Result<CompactStats, String> {
|
||||
let log = self.log_for(dataset).await;
|
||||
log.compact().await
|
||||
}
|
||||
|
||||
/// Remove every tombstone for a dataset. Called after a successful
|
||||
/// parquet compaction has physically deleted those rows — the
|
||||
/// tombstones have done their job and the journal can be cleared.
|
||||
///
|
||||
/// Implementation: drop the per-dataset AppendLog from the cache and
|
||||
/// delete all its batch files. Next write starts fresh.
|
||||
pub async fn clear(&self, dataset: &str) -> Result<usize, String> {
|
||||
let prefix = format!("{}/", Self::prefix_for(dataset));
|
||||
let keys = storaged::ops::list(&self.store, Some(&prefix)).await?;
|
||||
let matching: Vec<String> = keys
|
||||
.into_iter()
|
||||
.filter(|k| {
|
||||
let basename = k.rsplit('/').next().unwrap_or(k);
|
||||
basename.starts_with("batch_") && basename.ends_with(".jsonl")
|
||||
})
|
||||
.collect();
|
||||
let count = matching.len();
|
||||
for key in &matching {
|
||||
let _ = storaged::ops::delete(&self.store, key).await;
|
||||
}
|
||||
self.logs.write().await.remove(dataset);
|
||||
if count > 0 {
|
||||
tracing::info!("cleared {count} tombstone batch files for '{}'", dataset);
|
||||
}
|
||||
Ok(count)
|
||||
}
|
||||
}
|
||||
+103
-14
@@ -24,8 +24,30 @@ async fn main() {
|
||||
tracing::info!("config loaded: gateway={}:{}, storage={}",
|
||||
config.gateway.host, config.gateway.port, config.storage.root);
|
||||
|
||||
// Storage backend
|
||||
let store = storaged::backend::init_local(&config.storage.root);
|
||||
// Secrets provider
|
||||
let secrets = std::sync::Arc::new(shared::secrets::FileSecretsProvider::from_env());
|
||||
if let Err(e) = secrets.load().await {
|
||||
tracing::warn!("secrets load: {e} — bucket credentials may be missing");
|
||||
}
|
||||
|
||||
// Federation: bucket registry (primary + any [[storage.buckets]] entries + rescue)
|
||||
let bucket_registry = std::sync::Arc::new(
|
||||
storaged::registry::BucketRegistry::from_config(&config.storage, secrets.clone())
|
||||
.await
|
||||
.expect("bucket registry init failed")
|
||||
);
|
||||
let buckets = bucket_registry.list().await;
|
||||
tracing::info!(
|
||||
"bucket registry: {} buckets [{}], rescue={:?}",
|
||||
buckets.len(),
|
||||
buckets.iter().map(|b| format!("{}:{}", b.name, b.backend)).collect::<Vec<_>>().join(", "),
|
||||
bucket_registry.rescue_name(),
|
||||
);
|
||||
|
||||
// Back-compat: existing components that expect a single Arc<ObjectStore>
|
||||
// get the primary bucket's store. As we migrate call sites to be bucket-
|
||||
// aware, they'll call bucket_registry.get(name) directly.
|
||||
let store = bucket_registry.default_store();
|
||||
|
||||
// Catalog
|
||||
let registry = catalogd::registry::Registry::new(store.clone());
|
||||
@@ -35,7 +57,7 @@ async fn main() {
|
||||
|
||||
// Query engine with 16GB memory cache
|
||||
let cache = queryd::cache::MemCache::new(16 * 1024 * 1024 * 1024);
|
||||
let engine = queryd::context::QueryEngine::new(registry.clone(), store.clone(), cache);
|
||||
let engine = queryd::context::QueryEngine::new(registry.clone(), bucket_registry.clone(), cache);
|
||||
|
||||
// Event journal — append-only mutation log (flush every 100 events)
|
||||
let journal = journald::journal::Journal::new(store.clone(), 100);
|
||||
@@ -53,27 +75,94 @@ async fn main() {
|
||||
// AI sidecar client
|
||||
let ai_client = aibridge::client::AiClient::new(&config.sidecar.url);
|
||||
|
||||
// Vector service components — built before the router because both the
|
||||
// /vectors service AND ingestd need the agent handle to enqueue triggers.
|
||||
let index_reg = vectord::index_registry::IndexRegistry::new(store.clone());
|
||||
let _ = index_reg.rebuild().await;
|
||||
let hnsw = vectord::hnsw::HnswStore::new();
|
||||
let emb_cache = vectord::embedding_cache::EmbeddingCache::new(store.clone());
|
||||
// Phase B (federation layer 2): trial journals + promotion files now
|
||||
// live in each index's recorded bucket (IndexMeta.bucket), not always
|
||||
// primary. Journal / registry resolve per-call via the bucket registry.
|
||||
let tj = vectord::trial::TrialJournal::new(bucket_registry.clone(), index_reg.clone());
|
||||
let pr = vectord::promotion::PromotionRegistry::new(bucket_registry.clone(), index_reg.clone());
|
||||
|
||||
// Phase 16.2: spawn the autotune agent. When config.agent.enabled=false
|
||||
// this returns a handle that drops triggers silently — no surprise load.
|
||||
let agent_cfg = vectord::agent::AgentConfig {
|
||||
enabled: config.agent.enabled,
|
||||
cycle_interval_secs: config.agent.cycle_interval_secs,
|
||||
cooldown_between_trials_secs: config.agent.cooldown_between_trials_secs,
|
||||
min_recall: config.agent.min_recall,
|
||||
max_trials_per_hour: config.agent.max_trials_per_hour,
|
||||
};
|
||||
let agent_handle = vectord::agent::spawn(
|
||||
agent_cfg,
|
||||
vectord::agent::AgentDeps {
|
||||
store: store.clone(),
|
||||
ai_client: ai_client.clone(),
|
||||
catalog: registry.clone(),
|
||||
index_registry: index_reg.clone(),
|
||||
hnsw_store: hnsw.clone(),
|
||||
embedding_cache: emb_cache.clone(),
|
||||
trial_journal: tj.clone(),
|
||||
promotion_registry: pr.clone(),
|
||||
},
|
||||
);
|
||||
|
||||
// HTTP router
|
||||
let mut app = Router::new()
|
||||
.route("/health", get(health))
|
||||
.nest("/storage", storaged::service::router(store.clone()))
|
||||
.nest("/storage", storaged::service::router(store.clone())
|
||||
.merge(storaged::federation_service::router(bucket_registry.clone())))
|
||||
.nest("/catalog", catalogd::service::router(registry.clone()))
|
||||
.nest("/query", queryd::service::router(engine.clone()))
|
||||
.nest("/ai", aibridge::service::router(ai_client.clone()))
|
||||
;
|
||||
|
||||
// Phase E: scheduled ingest
|
||||
let sched_store = ingestd::schedule::ScheduleStore::new(store.clone());
|
||||
if let Ok(n) = sched_store.rebuild().await {
|
||||
if n > 0 { tracing::info!("rebuilt {n} persisted schedule(s)"); }
|
||||
}
|
||||
let scheduler = ingestd::schedule::Scheduler {
|
||||
store: sched_store.clone(),
|
||||
ingest: ingestd::schedule::SchedulerIngestDeps {
|
||||
store: store.clone(),
|
||||
registry: registry.clone(),
|
||||
buckets: bucket_registry.clone(),
|
||||
agent_handle: Some(agent_handle.clone()),
|
||||
index_registry: index_reg.clone(),
|
||||
},
|
||||
tick_secs: 10,
|
||||
};
|
||||
scheduler.spawn();
|
||||
|
||||
app = app
|
||||
.nest("/ingest", ingestd::service::router(ingestd::service::IngestState {
|
||||
store: store.clone(),
|
||||
registry: registry.clone(),
|
||||
buckets: bucket_registry.clone(),
|
||||
agent_handle: agent_handle.clone(),
|
||||
index_registry: index_reg.clone(),
|
||||
schedules: sched_store,
|
||||
}))
|
||||
.nest("/vectors", vectord::service::router({
|
||||
let index_reg = vectord::index_registry::IndexRegistry::new(store.clone());
|
||||
let _ = index_reg.rebuild().await;
|
||||
vectord::service::VectorState {
|
||||
store: store.clone(),
|
||||
ai_client: ai_client.clone(),
|
||||
job_tracker: vectord::jobs::JobTracker::new(),
|
||||
index_registry: index_reg,
|
||||
hnsw_store: vectord::hnsw::HnswStore::new(),
|
||||
}
|
||||
.nest("/vectors", vectord::service::router(vectord::service::VectorState {
|
||||
store: store.clone(),
|
||||
ai_client: ai_client.clone(),
|
||||
job_tracker: vectord::jobs::JobTracker::new(),
|
||||
index_registry: index_reg.clone(),
|
||||
hnsw_store: hnsw,
|
||||
embedding_cache: emb_cache,
|
||||
trial_journal: tj,
|
||||
catalog: registry.clone(),
|
||||
promotion_registry: pr,
|
||||
agent_handle,
|
||||
bucket_registry: bucket_registry.clone(),
|
||||
active_profile: std::sync::Arc::new(tokio::sync::RwLock::new(None)),
|
||||
lance: vectord::lance_backend::LanceRegistry::new(
|
||||
bucket_registry.clone(), index_reg.clone(),
|
||||
),
|
||||
}))
|
||||
.nest("/workspaces", queryd::workspace_service::router(workspace_mgr))
|
||||
.nest("/journal", journald::service::router(journal))
|
||||
|
||||
@@ -7,6 +7,7 @@ edition = "2024"
|
||||
shared = { path = "../shared" }
|
||||
storaged = { path = "../storaged" }
|
||||
catalogd = { path = "../catalogd" }
|
||||
vectord = { path = "../vectord" }
|
||||
tokio = { workspace = true }
|
||||
axum = { workspace = true, features = ["multipart"] }
|
||||
lopdf = { workspace = true }
|
||||
@@ -21,3 +22,5 @@ csv = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
object_store = { workspace = true }
|
||||
tokio-postgres = { workspace = true }
|
||||
mysql_async = { workspace = true }
|
||||
uuid = { workspace = true }
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
pub mod db_ingest;
|
||||
pub mod my_stream;
|
||||
pub mod pg_stream;
|
||||
pub mod schedule;
|
||||
pub mod detect;
|
||||
pub mod csv_ingest;
|
||||
pub mod json_ingest;
|
||||
|
||||
@@ -0,0 +1,410 @@
|
||||
/// Streaming MySQL ingest.
|
||||
///
|
||||
/// Mirrors `pg_stream` for MySQL sources. Same OFFSET-paginated strategy,
|
||||
/// same Parquet-streaming shape. Uses `mysql_async` (pure-rust) so we
|
||||
/// don't need a C client library at build time.
|
||||
///
|
||||
/// Type mapping follows ADR-010 (default to string on ambiguity):
|
||||
/// - Booleans (TINYINT(1)) and integer types map to Arrow Int32/Int64
|
||||
/// - Floating point and decimals → Float64
|
||||
/// - Everything else (text, varchar, json, date, timestamp) → Utf8
|
||||
///
|
||||
/// What's deliberately not supported (yet):
|
||||
/// - TLS connections — `minimal` feature is plain TCP only. Upgrade
|
||||
/// when a tenant actually needs it.
|
||||
/// - Keyset pagination — OFFSET scans are O(N²) at multi-million-row
|
||||
/// scale; same caveat as `pg_stream`.
|
||||
/// - BINARY/BLOB columns — currently rendered as base64 or empty string
|
||||
/// via Display-fallback.
|
||||
|
||||
use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray};
|
||||
use arrow::datatypes::{DataType, Field, Schema};
|
||||
use mysql_async::prelude::*;
|
||||
use mysql_async::{Conn, Opts, Row, Value};
|
||||
use parquet::arrow::ArrowWriter;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Request shape for MySQL streaming ingest.
|
||||
#[derive(Debug, Clone, serde::Deserialize)]
|
||||
pub struct MyStreamRequest {
|
||||
/// mysql://user:pass@host:port/db
|
||||
pub dsn: String,
|
||||
pub table: String,
|
||||
#[serde(default)]
|
||||
pub dataset_name: Option<String>,
|
||||
/// Rows per fetch. Default 10_000.
|
||||
#[serde(default)]
|
||||
pub batch_size: Option<usize>,
|
||||
/// Column to ORDER BY for stable pagination. If omitted, the first
|
||||
/// column returned by the schema probe is used.
|
||||
#[serde(default)]
|
||||
pub order_by: Option<String>,
|
||||
/// Hard cap on total rows (for sampling / previews).
|
||||
#[serde(default)]
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct MyStreamResult {
|
||||
pub table: String,
|
||||
pub rows: usize,
|
||||
pub batches: usize,
|
||||
pub columns: usize,
|
||||
pub schema: Vec<String>,
|
||||
pub parquet_bytes: u64,
|
||||
pub duration_secs: f32,
|
||||
}
|
||||
|
||||
/// Parsed DSN pieces. Kept local (rather than reusing pg's DbConfig) so
|
||||
/// the MySQL connector doesn't depend on the PG path.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MyConfig {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
pub user: String,
|
||||
pub password: String,
|
||||
pub database: String,
|
||||
}
|
||||
|
||||
impl MyConfig {
|
||||
/// Build a mysql_async `Opts` from the parsed config.
|
||||
pub fn to_opts(&self) -> Opts {
|
||||
let url = if self.password.is_empty() {
|
||||
format!(
|
||||
"mysql://{}@{}:{}/{}",
|
||||
percent(&self.user), percent(&self.host), self.port, percent(&self.database),
|
||||
)
|
||||
} else {
|
||||
format!(
|
||||
"mysql://{}:{}@{}:{}/{}",
|
||||
percent(&self.user), percent(&self.password),
|
||||
percent(&self.host), self.port, percent(&self.database),
|
||||
)
|
||||
};
|
||||
Opts::from_url(&url).expect("MyConfig.to_opts: rebuilt URL should parse")
|
||||
}
|
||||
}
|
||||
|
||||
/// Minimal URL-encoder for the few characters that commonly appear in
|
||||
/// MySQL passwords. mysql_async's URL parser expects valid URL-encoded
|
||||
/// components.
|
||||
fn percent(s: &str) -> String {
|
||||
let mut out = String::with_capacity(s.len());
|
||||
for c in s.chars() {
|
||||
match c {
|
||||
'@' | ':' | '/' | '?' | '#' | '[' | ']' | ' ' | '%' => {
|
||||
out.push_str(&format!("%{:02X}", c as u32));
|
||||
}
|
||||
_ => out.push(c),
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Parse a mysql:// DSN.
|
||||
/// Supports: mysql://[user[:password]@]host[:port][/db]
|
||||
pub fn parse_dsn(dsn: &str) -> Result<MyConfig, String> {
|
||||
let rest = dsn
|
||||
.strip_prefix("mysql://")
|
||||
.ok_or_else(|| "DSN must start with mysql://".to_string())?;
|
||||
|
||||
let (auth_host, database) = match rest.split_once('/') {
|
||||
Some((ah, db)) => (ah, db.split('?').next().unwrap_or(db).to_string()),
|
||||
None => (rest, String::new()),
|
||||
};
|
||||
|
||||
let (userpass, hostport) = match auth_host.rsplit_once('@') {
|
||||
Some((up, hp)) => (Some(up), hp),
|
||||
None => (None, auth_host),
|
||||
};
|
||||
|
||||
let (user, password) = match userpass {
|
||||
Some(up) => match up.split_once(':') {
|
||||
Some((u, p)) => (u.to_string(), p.to_string()),
|
||||
None => (up.to_string(), String::new()),
|
||||
},
|
||||
None => ("root".to_string(), String::new()),
|
||||
};
|
||||
|
||||
let (host, port) = match hostport.rsplit_once(':') {
|
||||
Some((h, p)) => (
|
||||
h.to_string(),
|
||||
p.parse::<u16>().map_err(|_| format!("invalid port in DSN: {p}"))?,
|
||||
),
|
||||
None => (hostport.to_string(), 3306),
|
||||
};
|
||||
|
||||
if host.is_empty() {
|
||||
return Err("DSN has no host".into());
|
||||
}
|
||||
if database.is_empty() {
|
||||
return Err("DSN has no database path (mysql://... /db)".into());
|
||||
}
|
||||
|
||||
Ok(MyConfig { host, port, user, password, database })
|
||||
}
|
||||
|
||||
/// Stream a MySQL table as Parquet bytes.
|
||||
pub async fn stream_table_to_parquet(
|
||||
req: &MyStreamRequest,
|
||||
) -> Result<(bytes::Bytes, MyStreamResult), String> {
|
||||
let t0 = std::time::Instant::now();
|
||||
let config = parse_dsn(&req.dsn)?;
|
||||
let batch_size = req.batch_size.unwrap_or(10_000).max(1);
|
||||
|
||||
let mut conn = Conn::new(config.to_opts()).await
|
||||
.map_err(|e| format!("mysql connect: {e}"))?;
|
||||
|
||||
// Probe columns via information_schema — gives real type names that
|
||||
// we can map to Arrow dtypes without needing to fetch a row first.
|
||||
let columns = probe_columns(&mut conn, &config.database, &req.table).await?;
|
||||
if columns.is_empty() {
|
||||
return Err(format!("table '{}' not found or has no columns", req.table));
|
||||
}
|
||||
|
||||
let arrow_fields: Vec<Field> = columns
|
||||
.iter()
|
||||
.map(|(name, ty)| Field::new(name, mysql_type_to_arrow(ty), true))
|
||||
.collect();
|
||||
let schema = Arc::new(Schema::new(arrow_fields));
|
||||
let schema_report: Vec<String> = columns
|
||||
.iter()
|
||||
.map(|(n, t)| format!("{}:{}", n, t))
|
||||
.collect();
|
||||
|
||||
let order_col = req.order_by.clone().unwrap_or_else(|| columns[0].0.clone());
|
||||
|
||||
let mut buf: Vec<u8> = Vec::with_capacity(1024 * 1024);
|
||||
let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None)
|
||||
.map_err(|e| format!("ArrowWriter init: {e}"))?;
|
||||
|
||||
let mut total_rows: usize = 0;
|
||||
let mut batch_count: usize = 0;
|
||||
let row_cap = req.limit.unwrap_or(usize::MAX);
|
||||
|
||||
loop {
|
||||
let remaining = row_cap.saturating_sub(total_rows);
|
||||
if remaining == 0 { break; }
|
||||
let fetch = remaining.min(batch_size);
|
||||
|
||||
// Backticks are MySQL's identifier quote. Forbid backticks in
|
||||
// table/column names to prevent injection — neither pg nor us
|
||||
// should be accepting those anyway.
|
||||
if req.table.contains('`') || order_col.contains('`') {
|
||||
return Err("table or order_by column contains backticks — refused".into());
|
||||
}
|
||||
let sql = format!(
|
||||
"SELECT * FROM `{}` ORDER BY `{}` LIMIT {} OFFSET {}",
|
||||
req.table, order_col, fetch, total_rows,
|
||||
);
|
||||
|
||||
let rows: Vec<Row> = conn.query(&sql).await
|
||||
.map_err(|e| format!("fetch batch at offset {total_rows}: {e}"))?;
|
||||
|
||||
if rows.is_empty() { break; }
|
||||
let n = rows.len();
|
||||
|
||||
let arrays: Vec<ArrayRef> = columns
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(idx, (_, ty))| rows_to_column(&rows, idx, ty))
|
||||
.collect::<Result<_, _>>()?;
|
||||
let batch = RecordBatch::try_new(schema.clone(), arrays)
|
||||
.map_err(|e| format!("RecordBatch: {e}"))?;
|
||||
writer.write(&batch).map_err(|e| format!("ArrowWriter::write: {e}"))?;
|
||||
|
||||
total_rows += n;
|
||||
batch_count += 1;
|
||||
tracing::info!(
|
||||
"mysql stream '{}': fetched batch {} ({} rows, total {})",
|
||||
req.table, batch_count, n, total_rows,
|
||||
);
|
||||
|
||||
if n < fetch { break; }
|
||||
}
|
||||
|
||||
writer.close().map_err(|e| format!("ArrowWriter::close: {e}"))?;
|
||||
conn.disconnect().await.ok();
|
||||
|
||||
let result = MyStreamResult {
|
||||
table: req.table.clone(),
|
||||
rows: total_rows,
|
||||
batches: batch_count,
|
||||
columns: columns.len(),
|
||||
schema: schema_report,
|
||||
parquet_bytes: buf.len() as u64,
|
||||
duration_secs: t0.elapsed().as_secs_f32(),
|
||||
};
|
||||
|
||||
Ok((bytes::Bytes::from(buf), result))
|
||||
}
|
||||
|
||||
async fn probe_columns(
|
||||
conn: &mut Conn,
|
||||
schema: &str,
|
||||
table: &str,
|
||||
) -> Result<Vec<(String, String)>, String> {
|
||||
let sql = format!(
|
||||
"SELECT COLUMN_NAME, DATA_TYPE FROM information_schema.columns \
|
||||
WHERE table_schema = '{}' AND table_name = '{}' \
|
||||
ORDER BY ORDINAL_POSITION",
|
||||
schema.replace('\'', "''"),
|
||||
table.replace('\'', "''"),
|
||||
);
|
||||
let rows: Vec<(String, String)> = conn.query(&sql).await
|
||||
.map_err(|e| format!("probe columns: {e}"))?;
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
/// MySQL data type string → Arrow DataType. Conservative: anything we
|
||||
/// don't recognize becomes Utf8 (ADR-010).
|
||||
fn mysql_type_to_arrow(ty: &str) -> DataType {
|
||||
match ty.to_lowercase().as_str() {
|
||||
"tinyint" | "smallint" | "mediumint" | "int" | "integer" => DataType::Int32,
|
||||
"bigint" => DataType::Int64,
|
||||
"float" | "double" | "decimal" | "numeric" | "real" => DataType::Float64,
|
||||
"bit" | "bool" | "boolean" => DataType::Boolean,
|
||||
_ => DataType::Utf8,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a single column slice of MySQL rows into an Arrow array.
|
||||
fn rows_to_column(
|
||||
rows: &[Row],
|
||||
idx: usize,
|
||||
ty: &str,
|
||||
) -> Result<ArrayRef, String> {
|
||||
let arrow_ty = mysql_type_to_arrow(ty);
|
||||
match arrow_ty {
|
||||
DataType::Boolean => {
|
||||
let v: Vec<Option<bool>> = rows.iter().map(|r| cell_as_bool(r, idx)).collect();
|
||||
Ok(Arc::new(BooleanArray::from(v)))
|
||||
}
|
||||
DataType::Int32 => {
|
||||
let v: Vec<Option<i32>> = rows.iter().map(|r| cell_as_i64(r, idx).map(|n| n as i32)).collect();
|
||||
Ok(Arc::new(Int32Array::from(v)))
|
||||
}
|
||||
DataType::Int64 => {
|
||||
let v: Vec<Option<i64>> = rows.iter().map(|r| cell_as_i64(r, idx)).collect();
|
||||
Ok(Arc::new(Int64Array::from(v)))
|
||||
}
|
||||
DataType::Float64 => {
|
||||
let v: Vec<Option<f64>> = rows.iter().map(|r| cell_as_f64(r, idx)).collect();
|
||||
Ok(Arc::new(Float64Array::from(v)))
|
||||
}
|
||||
_ => {
|
||||
let v: Vec<Option<String>> = rows.iter().map(|r| cell_as_string(r, idx)).collect();
|
||||
Ok(Arc::new(StringArray::from(v)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn cell(r: &Row, idx: usize) -> &Value { r.as_ref(idx).unwrap_or(&Value::NULL) }
|
||||
|
||||
fn cell_as_bool(r: &Row, idx: usize) -> Option<bool> {
|
||||
match cell(r, idx) {
|
||||
Value::NULL => None,
|
||||
Value::Int(n) => Some(*n != 0),
|
||||
Value::UInt(n) => Some(*n != 0),
|
||||
Value::Bytes(b) => std::str::from_utf8(b).ok().and_then(|s| {
|
||||
match s.to_ascii_lowercase().as_str() {
|
||||
"true" | "1" | "y" | "yes" => Some(true),
|
||||
"false" | "0" | "n" | "no" => Some(false),
|
||||
_ => None,
|
||||
}
|
||||
}),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn cell_as_i64(r: &Row, idx: usize) -> Option<i64> {
|
||||
match cell(r, idx) {
|
||||
Value::NULL => None,
|
||||
Value::Int(n) => Some(*n),
|
||||
Value::UInt(n) => i64::try_from(*n).ok(),
|
||||
Value::Float(f) => Some(*f as i64),
|
||||
Value::Double(f) => Some(*f as i64),
|
||||
Value::Bytes(b) => std::str::from_utf8(b).ok().and_then(|s| s.parse().ok()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn cell_as_f64(r: &Row, idx: usize) -> Option<f64> {
|
||||
match cell(r, idx) {
|
||||
Value::NULL => None,
|
||||
Value::Int(n) => Some(*n as f64),
|
||||
Value::UInt(n) => Some(*n as f64),
|
||||
Value::Float(f) => Some(*f as f64),
|
||||
Value::Double(f) => Some(*f),
|
||||
Value::Bytes(b) => std::str::from_utf8(b).ok().and_then(|s| s.parse().ok()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn cell_as_string(r: &Row, idx: usize) -> Option<String> {
|
||||
match cell(r, idx) {
|
||||
Value::NULL => None,
|
||||
Value::Bytes(b) => Some(String::from_utf8_lossy(b).into_owned()),
|
||||
Value::Int(n) => Some(n.to_string()),
|
||||
Value::UInt(n) => Some(n.to_string()),
|
||||
Value::Float(f) => Some(f.to_string()),
|
||||
Value::Double(f) => Some(f.to_string()),
|
||||
Value::Date(y, mo, d, h, mi, s, _us) => {
|
||||
Some(format!("{:04}-{:02}-{:02} {:02}:{:02}:{:02}", y, mo, d, h, mi, s))
|
||||
}
|
||||
Value::Time(neg, days, h, mi, s, _us) => {
|
||||
let sign = if *neg { "-" } else { "" };
|
||||
Some(format!("{}{}d {:02}:{:02}:{:02}", sign, days, h, mi, s))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parse_dsn_full() {
|
||||
let c = parse_dsn("mysql://daisy:secret@db.example.com:3307/my_db").unwrap();
|
||||
assert_eq!(c.host, "db.example.com");
|
||||
assert_eq!(c.port, 3307);
|
||||
assert_eq!(c.user, "daisy");
|
||||
assert_eq!(c.password, "secret");
|
||||
assert_eq!(c.database, "my_db");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_dsn_default_port() {
|
||||
let c = parse_dsn("mysql://root@localhost/shop").unwrap();
|
||||
assert_eq!(c.port, 3306);
|
||||
assert_eq!(c.user, "root");
|
||||
assert_eq!(c.password, "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_dsn_no_auth() {
|
||||
let c = parse_dsn("mysql://127.0.0.1:3306/analytics").unwrap();
|
||||
assert_eq!(c.user, "root");
|
||||
assert_eq!(c.host, "127.0.0.1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_dsn_rejects_non_mysql() {
|
||||
assert!(parse_dsn("postgresql://host/db").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_dsn_requires_db() {
|
||||
assert!(parse_dsn("mysql://localhost:3306").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn type_map_basics() {
|
||||
assert_eq!(mysql_type_to_arrow("int"), DataType::Int32);
|
||||
assert_eq!(mysql_type_to_arrow("BIGINT"), DataType::Int64);
|
||||
assert_eq!(mysql_type_to_arrow("decimal"), DataType::Float64);
|
||||
assert_eq!(mysql_type_to_arrow("varchar"), DataType::Utf8);
|
||||
assert_eq!(mysql_type_to_arrow("json"), DataType::Utf8);
|
||||
assert_eq!(mysql_type_to_arrow("bool"), DataType::Boolean);
|
||||
}
|
||||
}
|
||||
@@ -4,8 +4,19 @@ use arrow::record_batch::RecordBatch;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Extract text from a PDF file.
|
||||
/// Returns one row per page: (page_number, text_content).
|
||||
/// This handles text-based PDFs. Scanned/image PDFs need OCR (not implemented yet).
|
||||
///
|
||||
/// Two-tier approach:
|
||||
/// 1. **Text extraction (lopdf)** — fast, works on digital/native PDFs.
|
||||
/// Returns one row per page: (source_file, page_number, text_content).
|
||||
/// 2. **OCR fallback (Tesseract)** — if text extraction yields zero
|
||||
/// pages (common with scanned documents, image-only PDFs, or PDFs
|
||||
/// exported from scanners without an OCR layer), we extract embedded
|
||||
/// images from the PDF and run Tesseract on each. Falls back
|
||||
/// gracefully if Tesseract isn't installed (returns an error
|
||||
/// directing the operator to install it).
|
||||
///
|
||||
/// Schema is identical for both paths so downstream consumers (chunker,
|
||||
/// vectord, queryd) don't need to know how the text was produced.
|
||||
pub fn parse_pdf(content: &[u8], source_filename: &str) -> Result<(Arc<Schema>, Vec<RecordBatch>), String> {
|
||||
let doc = lopdf::Document::load_mem(content)
|
||||
.map_err(|e| format!("PDF load error: {e}"))?;
|
||||
@@ -26,13 +37,35 @@ pub fn parse_pdf(content: &[u8], source_filename: &str) -> Result<(Arc<Schema>,
|
||||
}
|
||||
}
|
||||
|
||||
// Tier 2: OCR fallback for scanned pages.
|
||||
if page_numbers.is_empty() {
|
||||
// PDF has no extractable text — likely scanned/image
|
||||
return Err("PDF contains no extractable text (may be scanned/image — OCR not yet supported)".into());
|
||||
tracing::info!(
|
||||
"no extractable text in '{}' — attempting OCR on {} pages",
|
||||
source_filename, pages.len(),
|
||||
);
|
||||
let ocr_results = ocr_pdf_images(&doc, source_filename)?;
|
||||
if ocr_results.is_empty() {
|
||||
return Err(
|
||||
"PDF contains no extractable text and OCR produced no output. \
|
||||
Ensure tesseract is installed (`apt install tesseract-ocr tesseract-ocr-eng`)."
|
||||
.into(),
|
||||
);
|
||||
}
|
||||
for (pg, text) in ocr_results {
|
||||
page_numbers.push(pg);
|
||||
page_texts.push(text);
|
||||
sources.push(source_filename.to_string());
|
||||
}
|
||||
tracing::info!(
|
||||
"OCR recovered {} pages from '{}'", page_numbers.len(), source_filename,
|
||||
);
|
||||
} else {
|
||||
tracing::info!(
|
||||
"extracted {} pages with text from PDF '{}'",
|
||||
page_numbers.len(), source_filename,
|
||||
);
|
||||
}
|
||||
|
||||
tracing::info!("extracted {} pages with text from PDF '{}'", page_numbers.len(), source_filename);
|
||||
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("source_file", DataType::Utf8, false),
|
||||
Field::new("page_number", DataType::Int32, false),
|
||||
@@ -50,3 +83,197 @@ pub fn parse_pdf(content: &[u8], source_filename: &str) -> Result<(Arc<Schema>,
|
||||
|
||||
Ok((schema, vec![batch]))
|
||||
}
|
||||
|
||||
/// Extract embedded images from the PDF, write each to a temp file,
|
||||
/// run Tesseract, collect the OCR'd text. Returns (page_number, text)
|
||||
/// pairs for pages that yielded non-empty output.
|
||||
///
|
||||
/// Image extraction strategy: walk every page's resources looking for
|
||||
/// XObject streams with /Subtype /Image. Decode the raw stream bytes
|
||||
/// (lopdf handles FlateDecode/DCTDecode) and write them as-is to a
|
||||
/// temp file with the appropriate extension. Tesseract handles JPEG,
|
||||
/// PNG, TIFF, and BMP natively.
|
||||
fn ocr_pdf_images(
|
||||
doc: &lopdf::Document,
|
||||
source_filename: &str,
|
||||
) -> Result<Vec<(i32, String)>, String> {
|
||||
if !tesseract_available() {
|
||||
return Err(
|
||||
"tesseract not found on PATH — install with `apt install tesseract-ocr tesseract-ocr-eng`"
|
||||
.into(),
|
||||
);
|
||||
}
|
||||
|
||||
let pages = doc.get_pages();
|
||||
let mut results: Vec<(i32, String)> = Vec::new();
|
||||
|
||||
for (&page_num, &page_id) in &pages {
|
||||
let images = extract_page_images(doc, page_id);
|
||||
if images.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// OCR each image on this page and concatenate the text.
|
||||
let mut page_text = String::new();
|
||||
for (idx, img_bytes) in images.iter().enumerate() {
|
||||
match ocr_image_bytes(img_bytes, page_num, idx) {
|
||||
Ok(text) if !text.trim().is_empty() => {
|
||||
if !page_text.is_empty() { page_text.push('\n'); }
|
||||
page_text.push_str(text.trim());
|
||||
}
|
||||
Ok(_) => {} // empty OCR — skip
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"OCR failed on '{}' page {} image {}: {e}",
|
||||
source_filename, page_num, idx,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
if !page_text.is_empty() {
|
||||
results.push((page_num as i32, page_text));
|
||||
}
|
||||
}
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Extract raw image bytes from a PDF page's XObject resources.
|
||||
/// Handles the common case: page /Resources → /XObject dict → streams
|
||||
/// with /Subtype /Image.
|
||||
fn extract_page_images(
|
||||
doc: &lopdf::Document,
|
||||
page_id: lopdf::ObjectId,
|
||||
) -> Vec<Vec<u8>> {
|
||||
let mut images = Vec::new();
|
||||
|
||||
let page_obj = match doc.get_object(page_id) {
|
||||
Ok(o) => o,
|
||||
Err(_) => return images,
|
||||
};
|
||||
let page_dict = match page_obj.as_dict() {
|
||||
Ok(d) => d,
|
||||
Err(_) => return images,
|
||||
};
|
||||
|
||||
// Resolve /Resources → /XObject
|
||||
let resources = page_dict
|
||||
.get(b"Resources")
|
||||
.ok()
|
||||
.and_then(|r| doc.dereference(r).ok())
|
||||
.map(|(_, o)| o.clone());
|
||||
let xobjects = resources
|
||||
.as_ref()
|
||||
.and_then(|r| r.as_dict().ok())
|
||||
.and_then(|d| d.get(b"XObject").ok())
|
||||
.and_then(|x| doc.dereference(x).ok())
|
||||
.map(|(_, o)| o.clone());
|
||||
|
||||
let xobj_dict = match xobjects.as_ref().and_then(|x| x.as_dict().ok()) {
|
||||
Some(d) => d,
|
||||
None => return images,
|
||||
};
|
||||
|
||||
for (_name, obj_ref) in xobj_dict.iter() {
|
||||
let obj = match doc.dereference(obj_ref) {
|
||||
Ok((_, o)) => o.clone(),
|
||||
Err(_) => continue,
|
||||
};
|
||||
let stream = match obj.as_stream() {
|
||||
Ok(s) => s,
|
||||
Err(_) => continue,
|
||||
};
|
||||
// Check /Subtype == /Image
|
||||
let is_image = stream.dict.get(b"Subtype")
|
||||
.ok()
|
||||
.and_then(|s| s.as_name().ok())
|
||||
.map(|n| n == b"Image")
|
||||
.unwrap_or(false);
|
||||
if !is_image {
|
||||
continue;
|
||||
}
|
||||
// Get decoded content — lopdf handles FlateDecode/DCTDecode.
|
||||
let content = match stream.decompressed_content() {
|
||||
Ok(c) => c,
|
||||
Err(_) => {
|
||||
// Try raw content as fallback (some DCT streams are stored raw).
|
||||
stream.content.clone()
|
||||
}
|
||||
};
|
||||
if !content.is_empty() {
|
||||
images.push(content);
|
||||
}
|
||||
}
|
||||
images
|
||||
}
|
||||
|
||||
/// Run Tesseract on raw image bytes. Writes to a temp file, shells out,
|
||||
/// reads the output. Returns the OCR'd text.
|
||||
fn ocr_image_bytes(
|
||||
img_bytes: &[u8],
|
||||
page_num: u32,
|
||||
img_idx: usize,
|
||||
) -> Result<String, String> {
|
||||
use std::io::Write;
|
||||
use std::process::Command;
|
||||
|
||||
let tmp_dir = std::env::temp_dir();
|
||||
// Detect image format from magic bytes to pick the right extension.
|
||||
let ext = if img_bytes.starts_with(&[0xFF, 0xD8]) {
|
||||
"jpg"
|
||||
} else if img_bytes.starts_with(&[0x89, 0x50, 0x4E, 0x47]) {
|
||||
"png"
|
||||
} else if img_bytes.starts_with(b"II") || img_bytes.starts_with(b"MM") {
|
||||
"tiff"
|
||||
} else {
|
||||
// Tesseract can often handle raw pixel data with the right options,
|
||||
// but writing as PNG is safer. For truly unknown formats, try PNG
|
||||
// wrapping; if Tesseract can't read it, it'll error cleanly.
|
||||
"png"
|
||||
};
|
||||
|
||||
let img_path = tmp_dir.join(format!("lh_ocr_p{page_num}_i{img_idx}.{ext}"));
|
||||
let out_base = tmp_dir.join(format!("lh_ocr_p{page_num}_i{img_idx}_out"));
|
||||
let out_path = tmp_dir.join(format!("lh_ocr_p{page_num}_i{img_idx}_out.txt"));
|
||||
|
||||
// Write image bytes to temp file.
|
||||
let mut f = std::fs::File::create(&img_path)
|
||||
.map_err(|e| format!("create temp image: {e}"))?;
|
||||
f.write_all(img_bytes)
|
||||
.map_err(|e| format!("write temp image: {e}"))?;
|
||||
drop(f);
|
||||
|
||||
// Run Tesseract.
|
||||
let status = Command::new("tesseract")
|
||||
.arg(&img_path)
|
||||
.arg(&out_base)
|
||||
.arg("--oem").arg("3") // LSTM + legacy combined
|
||||
.arg("--psm").arg("6") // assume uniform block of text
|
||||
.arg("-l").arg("eng")
|
||||
.stdout(std::process::Stdio::null())
|
||||
.stderr(std::process::Stdio::null())
|
||||
.status()
|
||||
.map_err(|e| format!("tesseract exec: {e}"))?;
|
||||
|
||||
// Cleanup input regardless of outcome.
|
||||
let _ = std::fs::remove_file(&img_path);
|
||||
|
||||
if !status.success() {
|
||||
let _ = std::fs::remove_file(&out_path);
|
||||
return Err(format!("tesseract exited with status {status}"));
|
||||
}
|
||||
|
||||
let text = std::fs::read_to_string(&out_path).unwrap_or_default();
|
||||
let _ = std::fs::remove_file(&out_path);
|
||||
Ok(text)
|
||||
}
|
||||
|
||||
/// Quick check: is `tesseract` on PATH?
|
||||
fn tesseract_available() -> bool {
|
||||
std::process::Command::new("tesseract")
|
||||
.arg("--version")
|
||||
.stdout(std::process::Stdio::null())
|
||||
.stderr(std::process::Stdio::null())
|
||||
.status()
|
||||
.map(|s| s.success())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,308 @@
|
||||
/// Streaming PostgreSQL ingest.
|
||||
///
|
||||
/// The original `db_ingest::import_postgres_table` loads every row into
|
||||
/// memory before emitting Parquet — fine for small tables, blows up on
|
||||
/// 1M+ rows. This module paginates via `ORDER BY <pk> LIMIT N OFFSET M`
|
||||
/// and streams batches into an `ArrowWriter`, closing once exhausted.
|
||||
///
|
||||
/// Pagination is OFFSET-based (not keyset) — simpler, works on any table
|
||||
/// without needing to know the PK type. OFFSET N scans N rows per fetch,
|
||||
/// which becomes quadratic for very large tables. Upgrade to keyset when
|
||||
/// a user actually ingests something that hurts.
|
||||
|
||||
use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray};
|
||||
use arrow::datatypes::{DataType, Field, Schema};
|
||||
use parquet::arrow::ArrowWriter;
|
||||
use std::sync::Arc;
|
||||
use tokio_postgres::{Client, NoTls, types::Type as PgType};
|
||||
|
||||
use crate::db_ingest::DbConfig;
|
||||
|
||||
/// Request shape for streaming ingest — takes a DSN string and optional
|
||||
/// tuning knobs.
|
||||
#[derive(Debug, Clone, serde::Deserialize)]
|
||||
pub struct PgStreamRequest {
|
||||
/// postgresql://user:pass@host:port/db — alternative to DbConfig.
|
||||
pub dsn: String,
|
||||
pub table: String,
|
||||
#[serde(default)]
|
||||
pub dataset_name: Option<String>,
|
||||
/// Rows per fetch. Default 10_000.
|
||||
#[serde(default)]
|
||||
pub batch_size: Option<usize>,
|
||||
/// Column to ORDER BY for stable pagination. If omitted, the first
|
||||
/// column returned by the schema probe is used.
|
||||
#[serde(default)]
|
||||
pub order_by: Option<String>,
|
||||
/// Hard cap on total rows (for sampling / previews).
|
||||
#[serde(default)]
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct PgStreamResult {
|
||||
pub table: String,
|
||||
pub rows: usize,
|
||||
pub batches: usize,
|
||||
pub columns: usize,
|
||||
pub schema: Vec<String>,
|
||||
pub parquet_bytes: u64,
|
||||
pub duration_secs: f32,
|
||||
}
|
||||
|
||||
/// Parse a postgresql:// DSN into a DbConfig.
|
||||
/// Supports: postgresql://[user[:password]@]host[:port][/db]
|
||||
/// Does NOT support: query parameters (sslmode etc) — add if needed.
|
||||
pub fn parse_dsn(dsn: &str) -> Result<DbConfig, String> {
|
||||
let rest = dsn
|
||||
.strip_prefix("postgresql://")
|
||||
.or_else(|| dsn.strip_prefix("postgres://"))
|
||||
.ok_or_else(|| "DSN must start with postgresql:// or postgres://".to_string())?;
|
||||
|
||||
// Split off path (database) first.
|
||||
let (auth_host, database) = match rest.split_once('/') {
|
||||
Some((ah, db)) => (ah, db.split('?').next().unwrap_or(db).to_string()),
|
||||
None => (rest, "postgres".to_string()),
|
||||
};
|
||||
|
||||
// Split user[:password] from host[:port].
|
||||
let (userpass, hostport) = match auth_host.rsplit_once('@') {
|
||||
Some((up, hp)) => (Some(up), hp),
|
||||
None => (None, auth_host),
|
||||
};
|
||||
|
||||
let (user, password) = match userpass {
|
||||
Some(up) => match up.split_once(':') {
|
||||
Some((u, p)) => (u.to_string(), p.to_string()),
|
||||
None => (up.to_string(), String::new()),
|
||||
},
|
||||
None => ("postgres".to_string(), String::new()),
|
||||
};
|
||||
|
||||
let (host, port) = match hostport.rsplit_once(':') {
|
||||
Some((h, p)) => (
|
||||
h.to_string(),
|
||||
p.parse::<u16>().map_err(|_| format!("invalid port in DSN: {p}"))?,
|
||||
),
|
||||
None => (hostport.to_string(), 5432),
|
||||
};
|
||||
|
||||
if host.is_empty() {
|
||||
return Err("DSN has no host".into());
|
||||
}
|
||||
|
||||
Ok(DbConfig { host, port, database, user, password })
|
||||
}
|
||||
|
||||
/// Stream a Postgres table as Parquet bytes.
|
||||
///
|
||||
/// Returns the full Parquet payload + summary stats. The payload is
|
||||
/// written to an in-memory buffer via `ArrowWriter` so each batch gets
|
||||
/// columnar compression as it arrives — memory footprint stays roughly
|
||||
/// at one batch plus Parquet footer state.
|
||||
pub async fn stream_table_to_parquet(
|
||||
req: &PgStreamRequest,
|
||||
) -> Result<(bytes::Bytes, PgStreamResult), String> {
|
||||
let t0 = std::time::Instant::now();
|
||||
let config = parse_dsn(&req.dsn)?;
|
||||
let batch_size = req.batch_size.unwrap_or(10_000).max(1);
|
||||
|
||||
let (client, connection) = tokio_postgres::connect(&config.connection_string(), NoTls)
|
||||
.await
|
||||
.map_err(|e| format!("postgres connect: {e}"))?;
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = connection.await {
|
||||
tracing::error!("pg connection: {e}");
|
||||
}
|
||||
});
|
||||
|
||||
// Probe schema.
|
||||
let columns = probe_columns(&client, &req.table).await?;
|
||||
if columns.is_empty() {
|
||||
return Err(format!("table '{}' not found or has no columns", req.table));
|
||||
}
|
||||
|
||||
let arrow_fields: Vec<Field> = columns
|
||||
.iter()
|
||||
.map(|(name, pg)| Field::new(name, pg_type_to_arrow(pg), true))
|
||||
.collect();
|
||||
let schema = Arc::new(Schema::new(arrow_fields));
|
||||
let schema_report: Vec<String> = columns
|
||||
.iter()
|
||||
.map(|(n, t)| format!("{}:{}", n, t))
|
||||
.collect();
|
||||
|
||||
// Pagination key: user-specified or first column.
|
||||
let order_col = req.order_by.clone().unwrap_or_else(|| columns[0].0.clone());
|
||||
|
||||
// Stream batches into ArrowWriter.
|
||||
let mut buf: Vec<u8> = Vec::with_capacity(1024 * 1024);
|
||||
let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None)
|
||||
.map_err(|e| format!("ArrowWriter init: {e}"))?;
|
||||
|
||||
let mut total_rows: usize = 0;
|
||||
let mut batch_count: usize = 0;
|
||||
let row_cap = req.limit.unwrap_or(usize::MAX);
|
||||
|
||||
loop {
|
||||
let remaining = row_cap.saturating_sub(total_rows);
|
||||
if remaining == 0 { break; }
|
||||
let fetch = remaining.min(batch_size);
|
||||
|
||||
let sql = format!(
|
||||
"SELECT * FROM \"{}\" ORDER BY \"{}\" LIMIT {} OFFSET {}",
|
||||
req.table, order_col, fetch, total_rows,
|
||||
);
|
||||
|
||||
let rows = client
|
||||
.query(&sql, &[])
|
||||
.await
|
||||
.map_err(|e| format!("fetch batch at offset {total_rows}: {e}"))?;
|
||||
|
||||
if rows.is_empty() { break; }
|
||||
|
||||
let n = rows.len();
|
||||
let arrays: Vec<ArrayRef> = columns
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(idx, (_, pg))| rows_to_column(&rows, idx, pg))
|
||||
.collect::<Result<_, _>>()?;
|
||||
let batch = RecordBatch::try_new(schema.clone(), arrays)
|
||||
.map_err(|e| format!("RecordBatch: {e}"))?;
|
||||
writer.write(&batch).map_err(|e| format!("ArrowWriter::write: {e}"))?;
|
||||
|
||||
total_rows += n;
|
||||
batch_count += 1;
|
||||
tracing::info!(
|
||||
"pg stream '{}': fetched batch {} ({} rows, total {})",
|
||||
req.table, batch_count, n, total_rows,
|
||||
);
|
||||
|
||||
if n < fetch { break; } // short read = end of table
|
||||
}
|
||||
|
||||
writer.close().map_err(|e| format!("ArrowWriter::close: {e}"))?;
|
||||
let parquet_bytes = buf.len() as u64;
|
||||
let duration = t0.elapsed().as_secs_f32();
|
||||
|
||||
let result = PgStreamResult {
|
||||
table: req.table.clone(),
|
||||
rows: total_rows,
|
||||
batches: batch_count,
|
||||
columns: columns.len(),
|
||||
schema: schema_report,
|
||||
parquet_bytes,
|
||||
duration_secs: duration,
|
||||
};
|
||||
|
||||
Ok((bytes::Bytes::from(buf), result))
|
||||
}
|
||||
|
||||
async fn probe_columns(client: &Client, table: &str) -> Result<Vec<(String, PgType)>, String> {
|
||||
let stmt = client
|
||||
.prepare(&format!("SELECT * FROM \"{}\" LIMIT 0", table))
|
||||
.await
|
||||
.map_err(|e| format!("prepare schema probe: {e}"))?;
|
||||
Ok(stmt
|
||||
.columns()
|
||||
.iter()
|
||||
.map(|c| (c.name().to_string(), c.type_().clone()))
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn pg_type_to_arrow(pg: &PgType) -> DataType {
|
||||
match *pg {
|
||||
PgType::BOOL => DataType::Boolean,
|
||||
PgType::INT2 | PgType::INT4 => DataType::Int32,
|
||||
PgType::INT8 | PgType::OID => DataType::Int64,
|
||||
PgType::FLOAT4 | PgType::FLOAT8 | PgType::NUMERIC => DataType::Float64,
|
||||
_ => DataType::Utf8,
|
||||
}
|
||||
}
|
||||
|
||||
fn rows_to_column(
|
||||
rows: &[tokio_postgres::Row],
|
||||
idx: usize,
|
||||
pg: &PgType,
|
||||
) -> Result<ArrayRef, String> {
|
||||
match *pg {
|
||||
PgType::BOOL => {
|
||||
let v: Vec<Option<bool>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
|
||||
Ok(Arc::new(BooleanArray::from(v)))
|
||||
}
|
||||
PgType::INT2 => {
|
||||
let v: Vec<Option<i32>> = rows.iter()
|
||||
.map(|r| r.try_get::<_, i16>(idx).ok().map(|x| x as i32)).collect();
|
||||
Ok(Arc::new(Int32Array::from(v)))
|
||||
}
|
||||
PgType::INT4 => {
|
||||
let v: Vec<Option<i32>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
|
||||
Ok(Arc::new(Int32Array::from(v)))
|
||||
}
|
||||
PgType::INT8 | PgType::OID => {
|
||||
let v: Vec<Option<i64>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
|
||||
Ok(Arc::new(Int64Array::from(v)))
|
||||
}
|
||||
PgType::FLOAT4 => {
|
||||
let v: Vec<Option<f64>> = rows.iter()
|
||||
.map(|r| r.try_get::<_, f32>(idx).ok().map(|x| x as f64)).collect();
|
||||
Ok(Arc::new(Float64Array::from(v)))
|
||||
}
|
||||
PgType::FLOAT8 => {
|
||||
let v: Vec<Option<f64>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
|
||||
Ok(Arc::new(Float64Array::from(v)))
|
||||
}
|
||||
_ => {
|
||||
// Safe-default per ADR-010: everything else becomes string.
|
||||
let v: Vec<Option<String>> = rows.iter().map(|r| {
|
||||
r.try_get::<_, String>(idx).ok()
|
||||
.or_else(|| r.try_get::<_, serde_json::Value>(idx).ok().map(|j| j.to_string()))
|
||||
.or_else(|| {
|
||||
// Timestamps + uuid, rendered via Display.
|
||||
r.try_get::<_, chrono::DateTime<chrono::Utc>>(idx).ok()
|
||||
.map(|t| t.to_rfc3339())
|
||||
})
|
||||
.or_else(|| r.try_get::<_, uuid::Uuid>(idx).ok().map(|u| u.to_string()))
|
||||
.or(Some(String::new()))
|
||||
}).collect();
|
||||
Ok(Arc::new(StringArray::from(v)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parse_dsn_full() {
|
||||
let c = parse_dsn("postgresql://daisy:secret@db.example.com:6543/my_db").unwrap();
|
||||
assert_eq!(c.host, "db.example.com");
|
||||
assert_eq!(c.port, 6543);
|
||||
assert_eq!(c.user, "daisy");
|
||||
assert_eq!(c.password, "secret");
|
||||
assert_eq!(c.database, "my_db");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_dsn_minimal() {
|
||||
let c = parse_dsn("postgresql://localhost/knowledge_base").unwrap();
|
||||
assert_eq!(c.host, "localhost");
|
||||
assert_eq!(c.port, 5432);
|
||||
assert_eq!(c.user, "postgres");
|
||||
assert_eq!(c.password, "");
|
||||
assert_eq!(c.database, "knowledge_base");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_dsn_no_password() {
|
||||
let c = parse_dsn("postgres://postgres@127.0.0.1:5432/mydb").unwrap();
|
||||
assert_eq!(c.user, "postgres");
|
||||
assert_eq!(c.password, "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_dsn_rejects_non_pg() {
|
||||
assert!(parse_dsn("http://host/db").is_err());
|
||||
}
|
||||
}
|
||||
@@ -28,12 +28,29 @@ pub struct IngestResult {
|
||||
}
|
||||
|
||||
/// Full ingest pipeline: detect → parse → dedup → store → register.
|
||||
///
|
||||
/// Bucket name (federation layer 2) is recorded on the ObjectRef so the
|
||||
/// catalog knows which bucket holds the data — defaults to "primary" via
|
||||
/// the ingest service's resolve_bucket helper.
|
||||
pub async fn ingest_file(
|
||||
filename: &str,
|
||||
content: &[u8],
|
||||
dataset_name: Option<&str>,
|
||||
store: &Arc<dyn ObjectStore>,
|
||||
registry: &Registry,
|
||||
) -> Result<IngestResult, String> {
|
||||
ingest_file_to_bucket(filename, content, dataset_name, "primary", store, registry).await
|
||||
}
|
||||
|
||||
/// Same as `ingest_file` but with explicit target bucket name on the
|
||||
/// ObjectRef. Header-aware ingest endpoints call this directly.
|
||||
pub async fn ingest_file_to_bucket(
|
||||
filename: &str,
|
||||
content: &[u8],
|
||||
dataset_name: Option<&str>,
|
||||
bucket: &str,
|
||||
store: &Arc<dyn ObjectStore>,
|
||||
registry: &Registry,
|
||||
) -> Result<IngestResult, String> {
|
||||
// 1. Detect file type
|
||||
let file_type = detect_file_type(filename, content);
|
||||
@@ -98,7 +115,7 @@ pub async fn ingest_file(
|
||||
let schema_fp = fingerprint_schema(&schema);
|
||||
let now = chrono::Utc::now();
|
||||
let obj_ref = ObjectRef {
|
||||
bucket: "data".to_string(),
|
||||
bucket: bucket.to_string(),
|
||||
key: storage_key.clone(),
|
||||
size_bytes: parquet_size,
|
||||
created_at: now,
|
||||
@@ -147,6 +164,10 @@ pub async fn ingest_file(
|
||||
..Default::default()
|
||||
}).await;
|
||||
|
||||
// Phase C: if this dataset already had embeddings, they're now stale.
|
||||
// mark_embeddings_stale is a no-op for never-embedded datasets.
|
||||
let _ = registry.mark_embeddings_stale(&safe_name).await;
|
||||
|
||||
Ok(IngestResult {
|
||||
dataset_name: safe_name,
|
||||
file_type: format!("{:?}", file_type),
|
||||
|
||||
@@ -0,0 +1,626 @@
|
||||
//! Scheduled ingest — the production-substrate piece that turns the
|
||||
//! lakehouse from "manual API toolkit" into a system that runs itself.
|
||||
//!
|
||||
//! A `ScheduleDef` declares "what to ingest" + "when to fire." A
|
||||
//! background `Scheduler` task wakes periodically, asks each enabled
|
||||
//! schedule whether it's due, and runs the matching ingest path. Outcomes
|
||||
//! land back on the schedule so operators can see what happened without
|
||||
//! `journalctl` archaeology.
|
||||
//!
|
||||
//! Storage shape: one JSON file per schedule at
|
||||
//! `primary://_schedules/{id}.json`. Serializable via serde, rebuilt on
|
||||
//! startup. Same write-once shape as the rest of the catalog — no DB.
|
||||
//!
|
||||
//! Concurrency model: single-flight per schedule_id. The scheduler holds
|
||||
//! a `HashSet<String>` of currently-running ids; a tick that finds a
|
||||
//! schedule already in flight skips it (no queueing — a slow ingest
|
||||
//! shouldn't pile up runs). Different schedules run concurrently.
|
||||
//!
|
||||
//! What's deliberately not in scope here:
|
||||
//! - Cron expressions (the trigger enum has the variant but parsing is
|
||||
//! stubbed). Intervals cover 90% of operational scheduling; cron is
|
||||
//! easy to bolt on later.
|
||||
//! - Backoff / retry policies. A failed run records the failure and
|
||||
//! schedules `next_run_at` as if it succeeded — no exponential backoff.
|
||||
//! - Distributed coordination. Single-machine scheduler; no leader
|
||||
//! election. If you run two gateway instances they'll both fire.
|
||||
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use object_store::ObjectStore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use storaged::ops;
|
||||
|
||||
const PREFIX: &str = "_schedules";
|
||||
|
||||
// =================== Public types ===================
|
||||
|
||||
/// What kind of ingest the schedule fires.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ScheduleKind {
|
||||
/// Pull a MySQL/MariaDB table via the my_stream path.
|
||||
Mysql {
|
||||
/// `mysql://user:pass@host:port/db` — DSN with full credentials.
|
||||
/// Stored encrypted-at-rest is a future concern; for now it lives
|
||||
/// in plain JSON in the catalog. Don't put production root creds
|
||||
/// here without disk-level encryption.
|
||||
dsn: String,
|
||||
table: String,
|
||||
#[serde(default)]
|
||||
dataset_name: Option<String>,
|
||||
#[serde(default)]
|
||||
batch_size: Option<usize>,
|
||||
#[serde(default)]
|
||||
order_by: Option<String>,
|
||||
},
|
||||
/// Pull a Postgres table via the pg_stream path.
|
||||
Postgres {
|
||||
dsn: String,
|
||||
table: String,
|
||||
#[serde(default)]
|
||||
dataset_name: Option<String>,
|
||||
#[serde(default)]
|
||||
batch_size: Option<usize>,
|
||||
#[serde(default)]
|
||||
order_by: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
impl ScheduleKind {
|
||||
pub fn label(&self) -> String {
|
||||
match self {
|
||||
ScheduleKind::Mysql { table, .. } => format!("mysql:{table}"),
|
||||
ScheduleKind::Postgres { table, .. } => format!("postgres:{table}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// When the schedule fires.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ScheduleTrigger {
|
||||
/// Run every N seconds. Time anchor = compute_next_run_at decision.
|
||||
Interval { secs: u64 },
|
||||
/// Cron expression — parsing not implemented yet. Defining the
|
||||
/// variant now so the JSON shape is forward-compatible.
|
||||
Cron { expr: String },
|
||||
}
|
||||
|
||||
/// Outcome of a single run, kept on the schedule for observability.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct RunOutcome {
|
||||
pub at: DateTime<Utc>,
|
||||
pub success: bool,
|
||||
pub message: String,
|
||||
pub rows: Option<usize>,
|
||||
pub duration_secs: f32,
|
||||
}
|
||||
|
||||
/// One scheduled ingest definition.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ScheduleDef {
|
||||
pub id: String,
|
||||
pub kind: ScheduleKind,
|
||||
pub trigger: ScheduleTrigger,
|
||||
pub enabled: bool,
|
||||
pub created_at: DateTime<Utc>,
|
||||
#[serde(default)]
|
||||
pub created_by: String,
|
||||
/// When the next run is due. Recomputed after each completion via
|
||||
/// `compute_next_run_at`. On creation: set to "now" so the schedule
|
||||
/// fires on the next scheduler tick.
|
||||
pub next_run_at: DateTime<Utc>,
|
||||
#[serde(default)]
|
||||
pub last_run_at: Option<DateTime<Utc>>,
|
||||
#[serde(default)]
|
||||
pub last_outcome: Option<RunOutcome>,
|
||||
/// Total successful runs and total failures since creation. Useful
|
||||
/// for at-a-glance "is this schedule healthy?" checks.
|
||||
#[serde(default)]
|
||||
pub run_count: u64,
|
||||
#[serde(default)]
|
||||
pub failure_count: u64,
|
||||
}
|
||||
|
||||
impl ScheduleDef {
|
||||
/// True if this schedule should fire at `now`. Honors `enabled`.
|
||||
pub fn is_due(&self, now: DateTime<Utc>) -> bool {
|
||||
self.enabled && now >= self.next_run_at
|
||||
}
|
||||
}
|
||||
|
||||
// =================== Trigger semantics ===================
|
||||
//
|
||||
// J: THIS IS YOURS TO IMPLEMENT
|
||||
// =====================================================
|
||||
//
|
||||
// `compute_next_run_at` decides when a schedule fires next. The current
|
||||
// state is the "since last completion" semantics — simple but has real
|
||||
// trade-offs vs other strategies:
|
||||
//
|
||||
// 1. SINCE LAST COMPLETION (current default below):
|
||||
// next = completed_at + interval
|
||||
// Easy to reason about. A long-running ingest delays the next fire
|
||||
// by however long it took. Drift is bounded — the schedule never
|
||||
// "falls behind" trying to catch up.
|
||||
//
|
||||
// 2. ANCHORED INTERVAL:
|
||||
// next = previous_target + interval, regardless of when run finished
|
||||
// Keeps wall-clock alignment ("every 30min on the :00 and :30").
|
||||
// Risk: if ingest takes longer than the interval, the next run is
|
||||
// "due immediately" and you can fire back-to-back forever.
|
||||
//
|
||||
// 3. ANCHORED + SKIP:
|
||||
// Anchored, but skip a tick if the previous one is still running.
|
||||
// Best of both, slightly more code.
|
||||
//
|
||||
// 4. CRON:
|
||||
// Parse the expression, find the next match after `now`. Most
|
||||
// expressive but pulls in a cron parser dep.
|
||||
//
|
||||
// Cron is currently stubbed — falling back to "now + 1 hour" so a
|
||||
// mistakenly-configured Cron schedule doesn't melt the GPU.
|
||||
pub fn compute_next_run_at(
|
||||
trigger: &ScheduleTrigger,
|
||||
completed_at: DateTime<Utc>,
|
||||
_previous_next: DateTime<Utc>,
|
||||
) -> DateTime<Utc> {
|
||||
// TODO(J): pick a strategy. Starter = since-last-completion.
|
||||
match trigger {
|
||||
ScheduleTrigger::Interval { secs } => {
|
||||
completed_at + Duration::seconds(*secs as i64)
|
||||
}
|
||||
ScheduleTrigger::Cron { expr: _ } => {
|
||||
// Cron parsing not implemented — fall back to a safe 1h
|
||||
// window so a bad config can't fire-loop the system.
|
||||
completed_at + Duration::hours(1)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// =================== Persistence ===================
|
||||
|
||||
/// Storage path for a schedule's JSON record.
|
||||
fn schedule_key(id: &str) -> String {
|
||||
let safe: String = id
|
||||
.chars()
|
||||
.map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' { c } else { '_' })
|
||||
.collect();
|
||||
format!("{PREFIX}/{safe}.json")
|
||||
}
|
||||
|
||||
/// In-memory + on-disk schedule registry. The Scheduler reads from here;
|
||||
/// HTTP CRUD writes through here.
|
||||
#[derive(Clone)]
|
||||
pub struct ScheduleStore {
|
||||
/// Bucket where schedule JSON lives. Always primary today — a
|
||||
/// future enhancement could move to per-profile schedule buckets.
|
||||
store: Arc<dyn ObjectStore>,
|
||||
schedules: Arc<RwLock<HashMap<String, ScheduleDef>>>,
|
||||
}
|
||||
|
||||
impl ScheduleStore {
|
||||
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||
Self {
|
||||
store,
|
||||
schedules: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Load every persisted schedule into memory. Call on startup.
|
||||
pub async fn rebuild(&self) -> Result<usize, String> {
|
||||
let keys = ops::list(&self.store, Some(&format!("{PREFIX}/")))
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
let mut count = 0usize;
|
||||
for key in keys {
|
||||
if !key.ends_with(".json") { continue; }
|
||||
match ops::get(&self.store, &key).await {
|
||||
Ok(bytes) => match serde_json::from_slice::<ScheduleDef>(&bytes) {
|
||||
Ok(def) => {
|
||||
self.schedules.write().await.insert(def.id.clone(), def);
|
||||
count += 1;
|
||||
}
|
||||
Err(e) => tracing::warn!("schedule {key}: malformed: {e}"),
|
||||
},
|
||||
Err(e) => tracing::warn!("schedule {key}: read failed: {e}"),
|
||||
}
|
||||
}
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
pub async fn list(&self) -> Vec<ScheduleDef> {
|
||||
let mut out: Vec<ScheduleDef> = self.schedules.read().await.values().cloned().collect();
|
||||
out.sort_by(|a, b| a.id.cmp(&b.id));
|
||||
out
|
||||
}
|
||||
|
||||
pub async fn get(&self, id: &str) -> Option<ScheduleDef> {
|
||||
self.schedules.read().await.get(id).cloned()
|
||||
}
|
||||
|
||||
pub async fn put(&self, def: ScheduleDef) -> Result<ScheduleDef, String> {
|
||||
let json = serde_json::to_vec_pretty(&def).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &schedule_key(&def.id), json.into()).await?;
|
||||
self.schedules.write().await.insert(def.id.clone(), def.clone());
|
||||
Ok(def)
|
||||
}
|
||||
|
||||
pub async fn delete(&self, id: &str) -> Result<bool, String> {
|
||||
// Object_store's delete is idempotent — silently OKs missing keys.
|
||||
let _ = ops::delete(&self.store, &schedule_key(id)).await;
|
||||
Ok(self.schedules.write().await.remove(id).is_some())
|
||||
}
|
||||
|
||||
/// Update completion metadata + persist. Used by the scheduler
|
||||
/// after a run finishes (success or failure).
|
||||
pub async fn record_run(
|
||||
&self,
|
||||
id: &str,
|
||||
outcome: RunOutcome,
|
||||
) -> Result<(), String> {
|
||||
let mut guard = self.schedules.write().await;
|
||||
let Some(def) = guard.get_mut(id) else {
|
||||
return Err(format!("schedule '{id}' not registered"));
|
||||
};
|
||||
def.last_run_at = Some(outcome.at);
|
||||
if outcome.success {
|
||||
def.run_count += 1;
|
||||
} else {
|
||||
def.failure_count += 1;
|
||||
}
|
||||
def.next_run_at = compute_next_run_at(&def.trigger, outcome.at, def.next_run_at);
|
||||
def.last_outcome = Some(outcome);
|
||||
|
||||
let json = serde_json::to_vec_pretty(def).map_err(|e| e.to_string())?;
|
||||
let key = schedule_key(id);
|
||||
let store = self.store.clone();
|
||||
// Drop the write lock before async I/O.
|
||||
drop(guard);
|
||||
ops::put(&store, &key, json.into()).await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// =================== Scheduler task ===================
|
||||
|
||||
/// Long-running task that fires due schedules. Spawned at gateway
|
||||
/// startup. Stop semantics: drop the handle (the loop exits when its
|
||||
/// store is the only reference left, or sooner if you wire a stop signal).
|
||||
pub struct Scheduler {
|
||||
pub store: ScheduleStore,
|
||||
pub ingest: SchedulerIngestDeps,
|
||||
/// Polling cadence for the loop itself. The schedules' own intervals
|
||||
/// are independent — this just controls how granularly we check
|
||||
/// "is anything due?" Default 10s.
|
||||
pub tick_secs: u64,
|
||||
}
|
||||
|
||||
/// What the scheduler needs to actually run an ingest. We pass in only
|
||||
/// what's needed — keeps the dep set narrow.
|
||||
#[derive(Clone)]
|
||||
pub struct SchedulerIngestDeps {
|
||||
pub store: Arc<dyn ObjectStore>,
|
||||
pub registry: catalogd::registry::Registry,
|
||||
pub buckets: Arc<storaged::registry::BucketRegistry>,
|
||||
/// Optional Phase 16.5 hook — push a DatasetAppended event if any
|
||||
/// HNSW index is bound to the dataset that just got refreshed.
|
||||
pub agent_handle: Option<vectord::agent::AgentHandle>,
|
||||
pub index_registry: vectord::index_registry::IndexRegistry,
|
||||
}
|
||||
|
||||
impl Scheduler {
|
||||
/// Spawn the loop in the background. Returns immediately.
|
||||
pub fn spawn(self) {
|
||||
tokio::spawn(async move {
|
||||
self.run().await;
|
||||
});
|
||||
}
|
||||
|
||||
async fn run(self) {
|
||||
tracing::info!("scheduler started — tick={}s", self.tick_secs);
|
||||
let in_flight: Arc<RwLock<HashSet<String>>> = Arc::new(RwLock::new(HashSet::new()));
|
||||
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_secs(self.tick_secs)).await;
|
||||
let now = Utc::now();
|
||||
let due: Vec<ScheduleDef> = self.store.list().await
|
||||
.into_iter()
|
||||
.filter(|s| s.is_due(now))
|
||||
.collect();
|
||||
if due.is_empty() { continue; }
|
||||
|
||||
for def in due {
|
||||
{
|
||||
let in_flight_guard = in_flight.read().await;
|
||||
if in_flight_guard.contains(&def.id) {
|
||||
// Previous run still going — skip this tick.
|
||||
// We don't reschedule; next tick re-evaluates.
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
let store = self.store.clone();
|
||||
let ingest = self.ingest.clone();
|
||||
let in_flight = in_flight.clone();
|
||||
let id = def.id.clone();
|
||||
|
||||
in_flight.write().await.insert(id.clone());
|
||||
|
||||
tokio::spawn(async move {
|
||||
let t0 = std::time::Instant::now();
|
||||
let outcome = run_schedule(&def, &ingest).await;
|
||||
let elapsed = t0.elapsed().as_secs_f32();
|
||||
let outcome = RunOutcome {
|
||||
at: Utc::now(),
|
||||
success: outcome.is_ok(),
|
||||
message: match &outcome {
|
||||
Ok(s) => s.clone(),
|
||||
Err(e) => e.clone(),
|
||||
},
|
||||
rows: outcome.as_ref().ok().and_then(parse_rows_from_message),
|
||||
duration_secs: elapsed,
|
||||
};
|
||||
if let Err(e) = store.record_run(&id, outcome).await {
|
||||
tracing::warn!("scheduler: record_run({id}): {e}");
|
||||
}
|
||||
in_flight.write().await.remove(&id);
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Pluck a "rows=N" hint out of a success message — best-effort, no
|
||||
/// hard contract. Lets us surface row counts in the schedule status
|
||||
/// endpoint without forcing every ingest path to return structured
|
||||
/// stats here.
|
||||
fn parse_rows_from_message(msg: &String) -> Option<usize> {
|
||||
msg.split_whitespace()
|
||||
.find_map(|tok| tok.strip_prefix("rows=").and_then(|n| n.parse().ok()))
|
||||
}
|
||||
|
||||
/// Fire one schedule's ingest. Returns Ok(success_message) or Err(reason).
|
||||
pub async fn run_schedule(
|
||||
def: &ScheduleDef,
|
||||
deps: &SchedulerIngestDeps,
|
||||
) -> Result<String, String> {
|
||||
tracing::info!("scheduler: firing '{}' ({})", def.id, def.kind.label());
|
||||
match &def.kind {
|
||||
ScheduleKind::Mysql { dsn, table, dataset_name, batch_size, order_by } => {
|
||||
let req = crate::my_stream::MyStreamRequest {
|
||||
dsn: dsn.clone(),
|
||||
table: table.clone(),
|
||||
dataset_name: dataset_name.clone(),
|
||||
batch_size: *batch_size,
|
||||
order_by: order_by.clone(),
|
||||
limit: None,
|
||||
};
|
||||
run_mysql(req, deps).await
|
||||
}
|
||||
ScheduleKind::Postgres { dsn, table, dataset_name, batch_size, order_by } => {
|
||||
let req = crate::pg_stream::PgStreamRequest {
|
||||
dsn: dsn.clone(),
|
||||
table: table.clone(),
|
||||
dataset_name: dataset_name.clone(),
|
||||
batch_size: *batch_size,
|
||||
order_by: order_by.clone(),
|
||||
limit: None,
|
||||
};
|
||||
run_postgres(req, deps).await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_mysql(
|
||||
req: crate::my_stream::MyStreamRequest,
|
||||
deps: &SchedulerIngestDeps,
|
||||
) -> Result<String, String> {
|
||||
let (parquet, result) = crate::my_stream::stream_table_to_parquet(&req).await?;
|
||||
if result.rows == 0 {
|
||||
return Ok(format!("rows=0 — table empty"));
|
||||
}
|
||||
persist_and_register(
|
||||
&parquet, &result.table, req.dataset_name.as_deref(),
|
||||
result.rows, "mysql", &req.dsn, deps,
|
||||
).await
|
||||
}
|
||||
|
||||
async fn run_postgres(
|
||||
req: crate::pg_stream::PgStreamRequest,
|
||||
deps: &SchedulerIngestDeps,
|
||||
) -> Result<String, String> {
|
||||
let (parquet, result) = crate::pg_stream::stream_table_to_parquet(&req).await?;
|
||||
if result.rows == 0 {
|
||||
return Ok(format!("rows=0 — table empty"));
|
||||
}
|
||||
persist_and_register(
|
||||
&parquet, &result.table, req.dataset_name.as_deref(),
|
||||
result.rows, "postgresql", &req.dsn, deps,
|
||||
).await
|
||||
}
|
||||
|
||||
/// Shared write+register path. Mirrors the HTTP ingest handlers in
|
||||
/// service.rs so the scheduled path produces identical catalog output:
|
||||
/// PII detection, lineage with redacted DSN, mark-stale + agent trigger.
|
||||
async fn persist_and_register(
|
||||
parquet: &bytes::Bytes,
|
||||
table: &str,
|
||||
dataset_name_override: Option<&str>,
|
||||
rows: usize,
|
||||
source_system: &str,
|
||||
dsn: &str,
|
||||
deps: &SchedulerIngestDeps,
|
||||
) -> Result<String, String> {
|
||||
use shared::types::{ColumnMeta, Lineage, ObjectRef, SchemaFingerprint, Sensitivity};
|
||||
|
||||
let (schema, _) = shared::arrow_helpers::parquet_to_record_batches(parquet)
|
||||
.map_err(|e| format!("reparse parquet: {e}"))?;
|
||||
|
||||
let dataset_name = dataset_name_override.unwrap_or(table).to_string();
|
||||
let storage_key = format!("datasets/{}.parquet", dataset_name);
|
||||
let parquet_size = parquet.len() as u64;
|
||||
let bucket = "primary".to_string();
|
||||
|
||||
let target_store = deps.buckets.get(&bucket)?;
|
||||
ops::put(&target_store, &storage_key, parquet.clone()).await?;
|
||||
|
||||
let schema_fp = shared::arrow_helpers::fingerprint_schema(&schema);
|
||||
let now = Utc::now();
|
||||
deps.registry.register(
|
||||
dataset_name.clone(),
|
||||
SchemaFingerprint(schema_fp.0),
|
||||
vec![ObjectRef {
|
||||
bucket: bucket.clone(),
|
||||
key: storage_key.clone(),
|
||||
size_bytes: parquet_size,
|
||||
created_at: now,
|
||||
}],
|
||||
).await?;
|
||||
|
||||
let col_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
|
||||
let sensitivity = shared::pii::detect_dataset_sensitivity(&col_names);
|
||||
let columns: Vec<ColumnMeta> = schema.fields().iter().map(|f| {
|
||||
let sens = shared::pii::detect_sensitivity(f.name());
|
||||
ColumnMeta {
|
||||
name: f.name().clone(),
|
||||
data_type: f.data_type().to_string(),
|
||||
sensitivity: sens.clone(),
|
||||
description: String::new(),
|
||||
is_pii: matches!(sens, Some(Sensitivity::Pii)),
|
||||
}
|
||||
}).collect();
|
||||
|
||||
let lineage = Lineage {
|
||||
source_system: source_system.to_string(),
|
||||
source_file: format!("dsn: {}", redact_dsn(dsn)),
|
||||
ingest_job: format!("scheduled-{}-{}", source_system, now.timestamp_millis()),
|
||||
ingest_timestamp: now,
|
||||
parent_datasets: vec![],
|
||||
};
|
||||
|
||||
let _ = deps.registry.update_metadata(&dataset_name, catalogd::registry::MetadataUpdate {
|
||||
sensitivity,
|
||||
columns: Some(columns),
|
||||
lineage: Some(lineage),
|
||||
row_count: Some(rows as u64),
|
||||
..Default::default()
|
||||
}).await;
|
||||
|
||||
let _ = deps.registry.mark_embeddings_stale(&dataset_name).await;
|
||||
|
||||
if let Some(agent) = &deps.agent_handle {
|
||||
let bound = deps.index_registry.list(Some(&dataset_name), None).await;
|
||||
for meta in bound {
|
||||
let event = vectord::agent::TriggerEvent::dataset_appended(
|
||||
meta.index_name.clone(), &dataset_name,
|
||||
);
|
||||
let _ = agent.enqueue(event).await;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(format!("rows={rows} dataset={dataset_name} bytes={parquet_size}"))
|
||||
}
|
||||
|
||||
/// Same generic redactor used in service.rs — duplicated here to avoid
|
||||
/// pulling service::redact_dsn behind a `pub` visibility just for this.
|
||||
fn redact_dsn(dsn: &str) -> String {
|
||||
let scheme_end = match dsn.find("://") {
|
||||
Some(i) => i + 3,
|
||||
None => return dsn.to_string(),
|
||||
};
|
||||
let at_idx = match dsn.rfind('@') {
|
||||
Some(i) if i > scheme_end => i,
|
||||
_ => return dsn.to_string(),
|
||||
};
|
||||
let userpass = &dsn[scheme_end..at_idx];
|
||||
let colon_offset = match userpass.find(':') {
|
||||
Some(i) => i,
|
||||
None => return dsn.to_string(),
|
||||
};
|
||||
let colon_idx = scheme_end + colon_offset;
|
||||
format!("{}:***{}", &dsn[..colon_idx], &dsn[at_idx..])
|
||||
}
|
||||
|
||||
// =================== Tests ===================
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn mk(secs: u64) -> ScheduleDef {
|
||||
ScheduleDef {
|
||||
id: "test".into(),
|
||||
kind: ScheduleKind::Mysql {
|
||||
dsn: "mysql://x@host/db".into(),
|
||||
table: "t".into(),
|
||||
dataset_name: None,
|
||||
batch_size: None,
|
||||
order_by: None,
|
||||
},
|
||||
trigger: ScheduleTrigger::Interval { secs },
|
||||
enabled: true,
|
||||
created_at: Utc::now(),
|
||||
created_by: "tester".into(),
|
||||
next_run_at: Utc::now(),
|
||||
last_run_at: None,
|
||||
last_outcome: None,
|
||||
run_count: 0,
|
||||
failure_count: 0,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_due_respects_enabled() {
|
||||
let mut s = mk(60);
|
||||
s.next_run_at = Utc::now() - Duration::seconds(5);
|
||||
assert!(s.is_due(Utc::now()));
|
||||
s.enabled = false;
|
||||
assert!(!s.is_due(Utc::now()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_due_respects_future_time() {
|
||||
let mut s = mk(60);
|
||||
s.next_run_at = Utc::now() + Duration::seconds(60);
|
||||
assert!(!s.is_due(Utc::now()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn next_run_since_last_completion_advances() {
|
||||
let s = mk(120);
|
||||
let completed = Utc::now();
|
||||
let next = compute_next_run_at(&s.trigger, completed, completed);
|
||||
let delta = next - completed;
|
||||
assert_eq!(delta.num_seconds(), 120);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cron_falls_back_to_one_hour() {
|
||||
let trig = ScheduleTrigger::Cron { expr: "* * * * *".into() };
|
||||
let now = Utc::now();
|
||||
let next = compute_next_run_at(&trig, now, now);
|
||||
let delta = next - now;
|
||||
assert_eq!(delta.num_seconds(), 3600);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schedule_kind_label() {
|
||||
let k = ScheduleKind::Mysql {
|
||||
dsn: "x".into(), table: "customers".into(),
|
||||
dataset_name: None, batch_size: None, order_by: None,
|
||||
};
|
||||
assert_eq!(k.label(), "mysql:customers");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_dsn_handles_mysql_and_postgres() {
|
||||
assert_eq!(redact_dsn("mysql://u:secret@h/db"), "mysql://u:***@h/db");
|
||||
assert_eq!(redact_dsn("postgresql://u:secret@h:5432/db"), "postgresql://u:***@h:5432/db");
|
||||
assert_eq!(redact_dsn("mysql://u@h/db"), "mysql://u@h/db"); // no password
|
||||
}
|
||||
}
|
||||
@@ -1,9 +1,9 @@
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::{Multipart, Query, State},
|
||||
http::StatusCode,
|
||||
extract::{Multipart, Path, Query, State},
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::IntoResponse,
|
||||
routing::{get, post},
|
||||
routing::{delete, get, patch, post},
|
||||
};
|
||||
use bytes::Bytes;
|
||||
use object_store::ObjectStore;
|
||||
@@ -11,15 +11,74 @@ use serde::Deserialize;
|
||||
use std::sync::Arc;
|
||||
|
||||
use catalogd::registry::Registry;
|
||||
use crate::{db_ingest, pipeline};
|
||||
use crate::{db_ingest, my_stream, pg_stream, pipeline, schedule};
|
||||
use shared::arrow_helpers::record_batch_to_parquet;
|
||||
use shared::types::{ObjectRef, SchemaFingerprint};
|
||||
use storaged::ops;
|
||||
use storaged::registry::BucketRegistry;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct IngestState {
|
||||
pub store: Arc<dyn ObjectStore>,
|
||||
pub registry: Registry,
|
||||
/// Federation layer 2: lookup target bucket from request headers.
|
||||
pub buckets: Arc<BucketRegistry>,
|
||||
/// Phase 16.5: when ingest marks a dataset's embeddings stale, push
|
||||
/// a `DatasetAppended` trigger so the autotune agent can schedule
|
||||
/// re-trials. Holds the agent handle (no-op when agent disabled).
|
||||
pub agent_handle: vectord::agent::AgentHandle,
|
||||
/// Used to look up which HNSW indexes are attached to the
|
||||
/// just-ingested dataset. Each matching index gets one trigger.
|
||||
pub index_registry: vectord::index_registry::IndexRegistry,
|
||||
/// Scheduled-ingest registry. The scheduler task runs against this
|
||||
/// store; HTTP CRUD endpoints write through it.
|
||||
pub schedules: schedule::ScheduleStore,
|
||||
}
|
||||
|
||||
/// Push `DatasetAppended` triggers for every HNSW index bound to this
|
||||
/// dataset. Called after a successful ingest. Logs failures but never
|
||||
/// fails the ingest — the agent is best-effort, not load-bearing.
|
||||
async fn notify_agent_on_append(state: &IngestState, dataset_name: &str) {
|
||||
let indexes = state.index_registry.list(Some(dataset_name), None).await;
|
||||
if indexes.is_empty() {
|
||||
return;
|
||||
}
|
||||
for meta in indexes {
|
||||
let event = vectord::agent::TriggerEvent::dataset_appended(
|
||||
meta.index_name.clone(), dataset_name,
|
||||
);
|
||||
if let Err(e) = state.agent_handle.enqueue(event).await {
|
||||
tracing::warn!("agent enqueue failed for '{}': {}", meta.index_name, e);
|
||||
} else {
|
||||
tracing::info!(
|
||||
"agent: enqueued DatasetAppended({}) for index '{}'",
|
||||
dataset_name, meta.index_name,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve the target bucket from `X-Lakehouse-Bucket` header.
|
||||
/// Returns `(bucket_name, store_for_writes)`. Falls back to "primary"
|
||||
/// when the header is absent. Returns Err with the canonical bucket
|
||||
/// list when the header names an unknown bucket.
|
||||
fn resolve_bucket(
|
||||
headers: &HeaderMap,
|
||||
buckets: &BucketRegistry,
|
||||
) -> Result<(String, Arc<dyn ObjectStore>), (StatusCode, String)> {
|
||||
let target = headers
|
||||
.get("x-lakehouse-bucket")
|
||||
.and_then(|h| h.to_str().ok())
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_else(|| buckets.default_name().to_string());
|
||||
|
||||
match buckets.get(&target) {
|
||||
Ok(store) => Ok((target, store)),
|
||||
Err(_) => Err((
|
||||
StatusCode::NOT_FOUND,
|
||||
format!("unknown bucket '{}' — use GET /storage/buckets to list", target),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn router(state: IngestState) -> Router {
|
||||
@@ -28,6 +87,14 @@ pub fn router(state: IngestState) -> Router {
|
||||
.route("/file", post(ingest_file))
|
||||
.route("/postgres/tables", post(list_pg_tables))
|
||||
.route("/postgres/import", post(import_pg_table))
|
||||
.route("/db", post(ingest_db_stream))
|
||||
.route("/mysql", post(ingest_mysql_stream))
|
||||
// Phase E: scheduled ingest
|
||||
.route("/schedules", get(list_schedules).post(create_schedule))
|
||||
.route("/schedules/{id}", get(get_schedule)
|
||||
.patch(patch_schedule)
|
||||
.delete(delete_schedule))
|
||||
.route("/schedules/{id}/run-now", post(run_schedule_now))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
@@ -43,8 +110,11 @@ struct IngestQuery {
|
||||
async fn ingest_file(
|
||||
State(state): State<IngestState>,
|
||||
Query(query): Query<IngestQuery>,
|
||||
headers: HeaderMap,
|
||||
mut multipart: Multipart,
|
||||
) -> impl IntoResponse {
|
||||
let (bucket, store) = resolve_bucket(&headers, &state.buckets)?;
|
||||
|
||||
let field = match multipart.next_field().await {
|
||||
Ok(Some(f)) => f,
|
||||
Ok(None) => return Err((StatusCode::BAD_REQUEST, "no file uploaded".to_string())),
|
||||
@@ -55,12 +125,18 @@ async fn ingest_file(
|
||||
let content = field.bytes().await
|
||||
.map_err(|e| (StatusCode::BAD_REQUEST, format!("read error: {e}")))?;
|
||||
|
||||
tracing::info!("received file '{}' ({} bytes) for ingest", filename, content.len());
|
||||
tracing::info!(
|
||||
"ingest '{}' ({} bytes) -> bucket={}",
|
||||
filename, content.len(), bucket,
|
||||
);
|
||||
|
||||
let dataset_name = query.name.as_deref();
|
||||
|
||||
match pipeline::ingest_file(&filename, &content, dataset_name, &state.store, &state.registry).await {
|
||||
match pipeline::ingest_file_to_bucket(&filename, &content, dataset_name, &bucket, &store, &state.registry).await {
|
||||
Ok(result) => {
|
||||
if !result.deduplicated {
|
||||
notify_agent_on_append(&state, &result.dataset_name).await;
|
||||
}
|
||||
if result.deduplicated {
|
||||
Ok((StatusCode::OK, Json(result)))
|
||||
} else {
|
||||
@@ -188,3 +264,361 @@ async fn import_pg_table(
|
||||
"size_bytes": parquet_size,
|
||||
}))))
|
||||
}
|
||||
|
||||
/// Streaming DSN-based ingest. Paginates via ORDER BY + LIMIT/OFFSET so
|
||||
/// large tables don't blow up memory. This is the path the task spec calls
|
||||
/// `POST /ingest/db` with `{dsn, table, dataset_name}`.
|
||||
async fn ingest_db_stream(
|
||||
State(state): State<IngestState>,
|
||||
headers: HeaderMap,
|
||||
Json(req): Json<pg_stream::PgStreamRequest>,
|
||||
) -> Result<(StatusCode, Json<serde_json::Value>), (StatusCode, String)> {
|
||||
let (bucket, store) = resolve_bucket(&headers, &state.buckets)?;
|
||||
tracing::info!(
|
||||
"pg stream ingest: table='{}' dataset='{:?}' bucket='{}' batch_size={:?}",
|
||||
req.table, req.dataset_name, bucket, req.batch_size,
|
||||
);
|
||||
|
||||
// Stream from Postgres into Parquet bytes.
|
||||
let (parquet, stream_result) = pg_stream::stream_table_to_parquet(&req)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::BAD_GATEWAY, e))?;
|
||||
|
||||
if stream_result.rows == 0 {
|
||||
return Ok((StatusCode::OK, Json(serde_json::json!({
|
||||
"table": req.table,
|
||||
"rows": 0,
|
||||
"message": "table is empty",
|
||||
}))));
|
||||
}
|
||||
|
||||
// Recover schema from the parquet footer — keeps a single source of truth
|
||||
// for the schema fingerprint + column metadata.
|
||||
let (schema, _) = shared::arrow_helpers::parquet_to_record_batches(&parquet)
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("reparse parquet: {e}")))?;
|
||||
|
||||
let dataset_name = req.dataset_name.clone().unwrap_or_else(|| req.table.clone());
|
||||
let storage_key = format!("datasets/{}.parquet", dataset_name);
|
||||
let size_bytes = parquet.len() as u64;
|
||||
|
||||
ops::put(&store, &storage_key, parquet)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
|
||||
|
||||
let schema_fp = shared::arrow_helpers::fingerprint_schema(&schema);
|
||||
let now = chrono::Utc::now();
|
||||
state.registry.register(
|
||||
dataset_name.clone(),
|
||||
SchemaFingerprint(schema_fp.0),
|
||||
vec![ObjectRef {
|
||||
bucket: bucket.clone(),
|
||||
key: storage_key.clone(),
|
||||
size_bytes,
|
||||
created_at: now,
|
||||
}],
|
||||
).await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
|
||||
|
||||
// Rich metadata: PII auto-detect, lineage, row_count, columns.
|
||||
let col_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
|
||||
let sensitivity = shared::pii::detect_dataset_sensitivity(&col_names);
|
||||
let columns: Vec<shared::types::ColumnMeta> = schema.fields().iter().map(|f| {
|
||||
let sens = shared::pii::detect_sensitivity(f.name());
|
||||
shared::types::ColumnMeta {
|
||||
name: f.name().clone(),
|
||||
data_type: f.data_type().to_string(),
|
||||
sensitivity: sens.clone(),
|
||||
description: String::new(),
|
||||
is_pii: matches!(sens, Some(shared::types::Sensitivity::Pii)),
|
||||
}
|
||||
}).collect();
|
||||
|
||||
let lineage = shared::types::Lineage {
|
||||
source_system: "postgresql".to_string(),
|
||||
source_file: format!("dsn: {}", redact_dsn(&req.dsn)),
|
||||
ingest_job: format!("pg-stream-{}", now.timestamp_millis()),
|
||||
ingest_timestamp: now,
|
||||
parent_datasets: vec![],
|
||||
};
|
||||
|
||||
let _ = state.registry.update_metadata(&dataset_name, catalogd::registry::MetadataUpdate {
|
||||
sensitivity,
|
||||
columns: Some(columns),
|
||||
lineage: Some(lineage),
|
||||
row_count: Some(stream_result.rows as u64),
|
||||
..Default::default()
|
||||
}).await;
|
||||
|
||||
// Phase C: mark embeddings stale if the dataset already had a vector
|
||||
// index. No-op for newly-created datasets.
|
||||
let _ = state.registry.mark_embeddings_stale(&dataset_name).await;
|
||||
|
||||
// Phase 16.5: notify autotune agent. No-op if no indexes attached.
|
||||
notify_agent_on_append(&state, &dataset_name).await;
|
||||
|
||||
Ok((StatusCode::CREATED, Json(serde_json::json!({
|
||||
"dataset_name": dataset_name,
|
||||
"table": stream_result.table,
|
||||
"rows": stream_result.rows,
|
||||
"batches": stream_result.batches,
|
||||
"columns": stream_result.columns,
|
||||
"schema": stream_result.schema,
|
||||
"storage_key": storage_key,
|
||||
"size_bytes": size_bytes,
|
||||
"duration_secs": stream_result.duration_secs,
|
||||
}))))
|
||||
}
|
||||
|
||||
/// Streaming MySQL ingest. Mirrors `ingest_db_stream` — same pagination +
|
||||
/// Parquet-footer schema recovery + PII/lineage metadata. Different driver,
|
||||
/// otherwise identical.
|
||||
async fn ingest_mysql_stream(
|
||||
State(state): State<IngestState>,
|
||||
headers: HeaderMap,
|
||||
Json(req): Json<my_stream::MyStreamRequest>,
|
||||
) -> Result<(StatusCode, Json<serde_json::Value>), (StatusCode, String)> {
|
||||
let (bucket, store) = resolve_bucket(&headers, &state.buckets)?;
|
||||
tracing::info!(
|
||||
"mysql stream ingest: table='{}' dataset='{:?}' bucket='{}' batch_size={:?}",
|
||||
req.table, req.dataset_name, bucket, req.batch_size,
|
||||
);
|
||||
|
||||
let (parquet, stream_result) = my_stream::stream_table_to_parquet(&req)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::BAD_GATEWAY, e))?;
|
||||
|
||||
if stream_result.rows == 0 {
|
||||
return Ok((StatusCode::OK, Json(serde_json::json!({
|
||||
"table": req.table,
|
||||
"rows": 0,
|
||||
"message": "table is empty",
|
||||
}))));
|
||||
}
|
||||
|
||||
let (schema, _) = shared::arrow_helpers::parquet_to_record_batches(&parquet)
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("reparse parquet: {e}")))?;
|
||||
|
||||
let dataset_name = req.dataset_name.clone().unwrap_or_else(|| req.table.clone());
|
||||
let storage_key = format!("datasets/{}.parquet", dataset_name);
|
||||
let size_bytes = parquet.len() as u64;
|
||||
|
||||
ops::put(&store, &storage_key, parquet)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
|
||||
|
||||
let schema_fp = shared::arrow_helpers::fingerprint_schema(&schema);
|
||||
let now = chrono::Utc::now();
|
||||
state.registry.register(
|
||||
dataset_name.clone(),
|
||||
SchemaFingerprint(schema_fp.0),
|
||||
vec![ObjectRef {
|
||||
bucket: bucket.clone(),
|
||||
key: storage_key.clone(),
|
||||
size_bytes,
|
||||
created_at: now,
|
||||
}],
|
||||
).await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
|
||||
|
||||
let col_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
|
||||
let sensitivity = shared::pii::detect_dataset_sensitivity(&col_names);
|
||||
let columns: Vec<shared::types::ColumnMeta> = schema.fields().iter().map(|f| {
|
||||
let sens = shared::pii::detect_sensitivity(f.name());
|
||||
shared::types::ColumnMeta {
|
||||
name: f.name().clone(),
|
||||
data_type: f.data_type().to_string(),
|
||||
sensitivity: sens.clone(),
|
||||
description: String::new(),
|
||||
is_pii: matches!(sens, Some(shared::types::Sensitivity::Pii)),
|
||||
}
|
||||
}).collect();
|
||||
|
||||
let lineage = shared::types::Lineage {
|
||||
source_system: "mysql".to_string(),
|
||||
source_file: format!("dsn: {}", redact_dsn(&req.dsn)),
|
||||
ingest_job: format!("mysql-stream-{}", now.timestamp_millis()),
|
||||
ingest_timestamp: now,
|
||||
parent_datasets: vec![],
|
||||
};
|
||||
|
||||
let _ = state.registry.update_metadata(&dataset_name, catalogd::registry::MetadataUpdate {
|
||||
sensitivity,
|
||||
columns: Some(columns),
|
||||
lineage: Some(lineage),
|
||||
row_count: Some(stream_result.rows as u64),
|
||||
..Default::default()
|
||||
}).await;
|
||||
|
||||
// Phase C: mark stale + Phase 16.5: notify agent.
|
||||
let _ = state.registry.mark_embeddings_stale(&dataset_name).await;
|
||||
notify_agent_on_append(&state, &dataset_name).await;
|
||||
|
||||
Ok((StatusCode::CREATED, Json(serde_json::json!({
|
||||
"dataset_name": dataset_name,
|
||||
"table": stream_result.table,
|
||||
"rows": stream_result.rows,
|
||||
"batches": stream_result.batches,
|
||||
"columns": stream_result.columns,
|
||||
"schema": stream_result.schema,
|
||||
"storage_key": storage_key,
|
||||
"size_bytes": size_bytes,
|
||||
"duration_secs": stream_result.duration_secs,
|
||||
}))))
|
||||
}
|
||||
|
||||
/// Redact the password in any `scheme://user:pass@host/...` DSN. Works
|
||||
/// for postgres, mysql, and any other scheme using the same shape.
|
||||
fn redact_dsn(dsn: &str) -> String {
|
||||
let scheme_end = match dsn.find("://") {
|
||||
Some(i) => i + 3,
|
||||
None => return dsn.to_string(),
|
||||
};
|
||||
let at_idx = match dsn.rfind('@') {
|
||||
Some(i) if i > scheme_end => i,
|
||||
_ => return dsn.to_string(),
|
||||
};
|
||||
// Find the password separator (first `:` after scheme, before `@`).
|
||||
let userpass = &dsn[scheme_end..at_idx];
|
||||
let colon_offset = match userpass.find(':') {
|
||||
Some(i) => i,
|
||||
None => return dsn.to_string(), // no password in DSN
|
||||
};
|
||||
let colon_idx = scheme_end + colon_offset;
|
||||
format!("{}:***{}", &dsn[..colon_idx], &dsn[at_idx..])
|
||||
}
|
||||
|
||||
// --- Phase E: Scheduled ingest CRUD endpoints ---
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct CreateScheduleRequest {
|
||||
/// Stable id used in URLs and on disk. Required so callers can
|
||||
/// PATCH/DELETE deterministically. Sanitized on storage but the
|
||||
/// caller's id is what gets persisted as the canonical key.
|
||||
id: String,
|
||||
kind: schedule::ScheduleKind,
|
||||
trigger: schedule::ScheduleTrigger,
|
||||
#[serde(default = "default_enabled")]
|
||||
enabled: bool,
|
||||
#[serde(default)]
|
||||
created_by: String,
|
||||
}
|
||||
|
||||
fn default_enabled() -> bool { true }
|
||||
|
||||
async fn create_schedule(
|
||||
State(state): State<IngestState>,
|
||||
Json(req): Json<CreateScheduleRequest>,
|
||||
) -> impl IntoResponse {
|
||||
if state.schedules.get(&req.id).await.is_some() {
|
||||
return Err((StatusCode::CONFLICT, format!("schedule '{}' already exists", req.id)));
|
||||
}
|
||||
let now = chrono::Utc::now();
|
||||
let def = schedule::ScheduleDef {
|
||||
id: req.id,
|
||||
kind: req.kind,
|
||||
trigger: req.trigger,
|
||||
enabled: req.enabled,
|
||||
created_at: now,
|
||||
created_by: req.created_by,
|
||||
// Fire on the next tick — operators usually want to verify
|
||||
// their schedule works without waiting for the full interval.
|
||||
next_run_at: now,
|
||||
last_run_at: None,
|
||||
last_outcome: None,
|
||||
run_count: 0,
|
||||
failure_count: 0,
|
||||
};
|
||||
match state.schedules.put(def).await {
|
||||
Ok(d) => Ok((StatusCode::CREATED, Json(d))),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_schedules(State(state): State<IngestState>) -> impl IntoResponse {
|
||||
Json(state.schedules.list().await)
|
||||
}
|
||||
|
||||
async fn get_schedule(
|
||||
State(state): State<IngestState>,
|
||||
Path(id): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
match state.schedules.get(&id).await {
|
||||
Some(d) => Ok(Json(d)),
|
||||
None => Err((StatusCode::NOT_FOUND, format!("schedule '{id}' not found"))),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct PatchScheduleRequest {
|
||||
/// Toggle on/off without changing anything else.
|
||||
#[serde(default)]
|
||||
enabled: Option<bool>,
|
||||
/// Replace the trigger spec — useful for bumping intervals.
|
||||
#[serde(default)]
|
||||
trigger: Option<schedule::ScheduleTrigger>,
|
||||
}
|
||||
|
||||
async fn patch_schedule(
|
||||
State(state): State<IngestState>,
|
||||
Path(id): Path<String>,
|
||||
Json(req): Json<PatchScheduleRequest>,
|
||||
) -> impl IntoResponse {
|
||||
let Some(mut def) = state.schedules.get(&id).await else {
|
||||
return Err((StatusCode::NOT_FOUND, format!("schedule '{id}' not found")));
|
||||
};
|
||||
if let Some(e) = req.enabled { def.enabled = e; }
|
||||
if let Some(t) = req.trigger { def.trigger = t; }
|
||||
match state.schedules.put(def).await {
|
||||
Ok(d) => Ok(Json(d)),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_schedule(
|
||||
State(state): State<IngestState>,
|
||||
Path(id): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
match state.schedules.delete(&id).await {
|
||||
Ok(true) => Ok(Json(serde_json::json!({ "deleted": id }))),
|
||||
Ok(false) => Err((StatusCode::NOT_FOUND, format!("schedule '{id}' not found"))),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Force a schedule to fire right now, regardless of its next_run_at.
|
||||
/// Records the outcome and reschedules. Useful for "test my new schedule
|
||||
/// works without waiting an hour."
|
||||
async fn run_schedule_now(
|
||||
State(state): State<IngestState>,
|
||||
Path(id): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
let Some(def) = state.schedules.get(&id).await else {
|
||||
return Err((StatusCode::NOT_FOUND, format!("schedule '{id}' not found")));
|
||||
};
|
||||
let deps = schedule::SchedulerIngestDeps {
|
||||
store: state.store.clone(),
|
||||
registry: state.registry.clone(),
|
||||
buckets: state.buckets.clone(),
|
||||
agent_handle: Some(state.agent_handle.clone()),
|
||||
index_registry: state.index_registry.clone(),
|
||||
};
|
||||
|
||||
let t0 = std::time::Instant::now();
|
||||
let result = schedule::run_schedule(&def, &deps).await;
|
||||
let elapsed = t0.elapsed().as_secs_f32();
|
||||
let outcome = schedule::RunOutcome {
|
||||
at: chrono::Utc::now(),
|
||||
success: result.is_ok(),
|
||||
message: match &result {
|
||||
Ok(s) => s.clone(),
|
||||
Err(e) => e.clone(),
|
||||
},
|
||||
rows: result.as_ref().ok().and_then(|m: &String| {
|
||||
m.split_whitespace()
|
||||
.find_map(|tok| tok.strip_prefix("rows=").and_then(|n| n.parse().ok()))
|
||||
}),
|
||||
duration_secs: elapsed,
|
||||
};
|
||||
if let Err(e) = state.schedules.record_run(&id, outcome.clone()).await {
|
||||
tracing::warn!("schedule {id}: record_run failed: {e}");
|
||||
}
|
||||
Ok(Json(outcome))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
[package]
|
||||
name = "lance-bench"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
# Standalone pilot for Phase B (see docs/EXECUTION_PLAN.md).
|
||||
# Deliberately NOT sharing workspace deps — Lance 4.x pulls in its own
|
||||
# DataFusion and Arrow versions incompatible with the rest of the stack.
|
||||
# Isolating the pilot means we don't force a workspace-wide upgrade until
|
||||
# we've decided Lance is worth it.
|
||||
|
||||
[dependencies]
|
||||
# Only the features we actually need — the default brings in AWS/Azure/GCP/HF etc
|
||||
# which is ~200 extra crates we don't care about for a local pilot.
|
||||
lance = { version = "4.0", default-features = false }
|
||||
# Lance exposes DatasetIndexExt, IndexType, and IvfBuildParams through
|
||||
# its sub-crates which must be imported directly — lance itself doesn't
|
||||
# re-export them at a convenient path.
|
||||
lance-index = { version = "4.0", default-features = false }
|
||||
lance-linalg = { version = "4.0", default-features = false }
|
||||
|
||||
# Arrow re-exported by Lance; pin to a range Lance picks so types match.
|
||||
arrow = "57"
|
||||
arrow-array = "57"
|
||||
arrow-schema = "57"
|
||||
|
||||
# Also need to read the EXISTING Parquet vector files so we can compare.
|
||||
# These live in data/vectors/*.parquet. Lance's internal Parquet reading
|
||||
# might differ from ours; using our format's Arrow/Parquet versions for
|
||||
# the read side keeps the inputs identical.
|
||||
parquet = "57"
|
||||
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
futures = "0.3"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
anyhow = "1"
|
||||
bytes = "1"
|
||||
|
||||
[[bin]]
|
||||
name = "lance-bench"
|
||||
path = "src/main.rs"
|
||||
@@ -0,0 +1,633 @@
|
||||
//! Phase B: Lance pilot benchmark.
|
||||
//!
|
||||
//! Standalone binary that compares Lance vector storage against our
|
||||
//! Parquet-with-binary-blob + in-RAM HNSW approach. See
|
||||
//! docs/EXECUTION_PLAN.md for the decision rules this fuels.
|
||||
//!
|
||||
//! Inputs:
|
||||
//! data/vectors/resumes_100k_v2.parquet — existing 100K × 768d embeddings
|
||||
//!
|
||||
//! Output:
|
||||
//! A JSON report printed to stdout with measurements for:
|
||||
//! - Cold load time (parquet → arrow) vs Lance open + scan
|
||||
//! - Disk size
|
||||
//! - Vector search latency (p50 / p95 / p99)
|
||||
//! - Single-row random access
|
||||
//! - Append cost (adding 10K rows)
|
||||
//!
|
||||
//! Usage:
|
||||
//! cargo run --bin lance-bench -- \
|
||||
//! --parquet data/vectors/resumes_100k_v2.parquet \
|
||||
//! --lance-out /tmp/lance_resumes_100k_v2 \
|
||||
//! --json-out /tmp/lance_bench.json
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use arrow_array::{Array, ArrayRef, BinaryArray, FixedSizeListArray, Float32Array, RecordBatch, RecordBatchIterator};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use serde::Serialize;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct BenchReport {
|
||||
vectors: usize,
|
||||
dimensions: usize,
|
||||
parquet_path: String,
|
||||
lance_path: String,
|
||||
|
||||
// Parquet baseline
|
||||
parquet_disk_bytes: u64,
|
||||
parquet_cold_load_secs: f32,
|
||||
|
||||
// Lance numbers
|
||||
lance_write_secs: f32,
|
||||
lance_disk_bytes: u64,
|
||||
lance_cold_open_secs: f32,
|
||||
|
||||
// Index + search
|
||||
lance_index_build_secs: Option<f32>,
|
||||
lance_index_disk_bytes: Option<u64>,
|
||||
lance_search_p50_us: Option<f32>,
|
||||
lance_search_p95_us: Option<f32>,
|
||||
lance_search_p99_us: Option<f32>,
|
||||
|
||||
// Architectural features Parquet+sidecar can't cheaply do
|
||||
lance_random_row_access_us: Option<f32>, // fetch one row by row_id
|
||||
parquet_random_row_access_us: Option<f32>, // for comparison — full scan cost
|
||||
lance_append_10k_secs: Option<f32>, // add 10K new rows
|
||||
lance_append_disk_bytes_added: Option<u64>,
|
||||
|
||||
// Head-to-head reference (from our own measurements)
|
||||
reference_hnsw_p50_us: f32,
|
||||
reference_hnsw_p95_us: f32,
|
||||
reference_brute_force_us: f32,
|
||||
reference_hnsw_build_secs: f32,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
// Simple positional args: parquet_in, lance_out.
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
let parquet_path = args
|
||||
.get(1)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| "data/vectors/resumes_100k_v2.parquet".to_string());
|
||||
let lance_path = args
|
||||
.get(2)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| "/tmp/lance_bench_dataset".to_string());
|
||||
|
||||
eprintln!("=== Phase B Lance pilot ===");
|
||||
eprintln!("input parquet: {}", parquet_path);
|
||||
eprintln!("output lance: {}", lance_path);
|
||||
|
||||
// --- 1. Cold-load the existing Parquet vector index into memory
|
||||
eprintln!("\n[1/4] reading Parquet baseline...");
|
||||
let t0 = Instant::now();
|
||||
let (schema, batches, total_rows) = read_parquet_vectors(&parquet_path)
|
||||
.context("read parquet")?;
|
||||
let parquet_cold_load_secs = t0.elapsed().as_secs_f32();
|
||||
let parquet_disk_bytes = std::fs::metadata(&parquet_path)?.len();
|
||||
|
||||
let dims = detect_vector_dims(&batches)?;
|
||||
eprintln!(
|
||||
" loaded {} rows, {} columns, vectors={}d, disk={:.1} MB, cold load={:.2}s",
|
||||
total_rows,
|
||||
schema.fields().len(),
|
||||
dims,
|
||||
parquet_disk_bytes as f64 / 1_000_000.0,
|
||||
parquet_cold_load_secs,
|
||||
);
|
||||
|
||||
// --- 2. Convert from binary-blob-of-f32 to Lance's FixedSizeList<Float32>
|
||||
eprintln!("\n[2/4] converting binary-blob vectors to Arrow FixedSizeList...");
|
||||
let t0 = Instant::now();
|
||||
let (lance_schema, lance_batches) = convert_to_fixed_size_list(&schema, batches, dims)?;
|
||||
eprintln!(" conversion took {:.2}s", t0.elapsed().as_secs_f32());
|
||||
|
||||
// --- 3. Write as Lance dataset
|
||||
eprintln!("\n[3/4] writing Lance dataset...");
|
||||
let t0 = Instant::now();
|
||||
// Clean up any prior run
|
||||
let _ = std::fs::remove_dir_all(&lance_path);
|
||||
write_lance_dataset(&lance_path, lance_schema.clone(), lance_batches).await?;
|
||||
let lance_write_secs = t0.elapsed().as_secs_f32();
|
||||
let lance_disk_bytes = dir_size_bytes(&lance_path);
|
||||
eprintln!(
|
||||
" write took {:.2}s, disk={:.1} MB",
|
||||
lance_write_secs,
|
||||
lance_disk_bytes as f64 / 1_000_000.0,
|
||||
);
|
||||
|
||||
// --- 4. Cold open + scan the Lance dataset
|
||||
eprintln!("\n[4/6] cold-opening Lance dataset...");
|
||||
let t0 = Instant::now();
|
||||
let scanned_rows = cold_open_and_scan_lance(&lance_path).await?;
|
||||
let lance_cold_open_secs = t0.elapsed().as_secs_f32();
|
||||
eprintln!(
|
||||
" open + full scan: {} rows in {:.2}s",
|
||||
scanned_rows, lance_cold_open_secs,
|
||||
);
|
||||
|
||||
// --- 5. Build a vector index on the Lance dataset
|
||||
eprintln!("\n[5/6] building Lance vector index (IVF_PQ)...");
|
||||
let t0 = Instant::now();
|
||||
let index_built = build_lance_vector_index(&lance_path, dims).await;
|
||||
let (lance_index_build_secs, lance_index_disk_bytes) = match index_built {
|
||||
Ok(()) => {
|
||||
let secs = t0.elapsed().as_secs_f32();
|
||||
let disk = dir_size_bytes(&lance_path) - lance_disk_bytes;
|
||||
eprintln!(" built in {:.2}s, index adds {:.1} MB on disk", secs, disk as f64 / 1e6);
|
||||
(Some(secs), Some(disk))
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!(" index build failed: {e:#}");
|
||||
(None, None)
|
||||
}
|
||||
};
|
||||
|
||||
// --- 6. Run search queries, measure latency
|
||||
eprintln!("\n[6/6] running vector search benchmarks...");
|
||||
let search_stats = if lance_index_build_secs.is_some() {
|
||||
run_search_benchmarks(&lance_path, dims).await.ok()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let (lance_search_p50, lance_search_p95, lance_search_p99) = match search_stats {
|
||||
Some((p50, p95, p99)) => {
|
||||
eprintln!(" p50={:.0}us p95={:.0}us p99={:.0}us", p50, p95, p99);
|
||||
(Some(p50), Some(p95), Some(p99))
|
||||
}
|
||||
None => (None, None, None),
|
||||
};
|
||||
|
||||
// --- Random access comparison
|
||||
eprintln!("\n[7/8] random row access — Lance vs full-scan Parquet...");
|
||||
let lance_random = measure_random_access_lance(&lance_path).await.ok();
|
||||
let parquet_random = measure_random_access_parquet(&parquet_path).ok();
|
||||
if let Some(us) = lance_random {
|
||||
eprintln!(" Lance random-fetch avg: {:.0}us", us);
|
||||
}
|
||||
if let Some(us) = parquet_random {
|
||||
eprintln!(" Parquet full-scan-to-row avg: {:.0}us", us);
|
||||
}
|
||||
|
||||
// --- Append cost
|
||||
eprintln!("\n[8/8] append 10K new rows to existing dataset...");
|
||||
let t0 = Instant::now();
|
||||
let pre_append_bytes = dir_size_bytes(&lance_path);
|
||||
let append_result = append_10k_rows(&lance_path, dims).await;
|
||||
let (lance_append_secs, lance_append_bytes) = match append_result {
|
||||
Ok(()) => {
|
||||
let secs = t0.elapsed().as_secs_f32();
|
||||
let bytes = dir_size_bytes(&lance_path).saturating_sub(pre_append_bytes);
|
||||
eprintln!(" append took {:.2}s, added {:.1} MB", secs, bytes as f64 / 1e6);
|
||||
(Some(secs), Some(bytes))
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!(" append failed: {e:#}");
|
||||
(None, None)
|
||||
}
|
||||
};
|
||||
|
||||
// --- Report
|
||||
let report = BenchReport {
|
||||
vectors: total_rows,
|
||||
dimensions: dims,
|
||||
parquet_path: parquet_path.clone(),
|
||||
lance_path: lance_path.clone(),
|
||||
parquet_disk_bytes,
|
||||
parquet_cold_load_secs,
|
||||
lance_write_secs,
|
||||
lance_disk_bytes,
|
||||
lance_cold_open_secs,
|
||||
lance_index_build_secs,
|
||||
lance_index_disk_bytes,
|
||||
lance_search_p50_us: lance_search_p50,
|
||||
lance_search_p95_us: lance_search_p95,
|
||||
lance_search_p99_us: lance_search_p99,
|
||||
lance_random_row_access_us: lance_random,
|
||||
parquet_random_row_access_us: parquet_random,
|
||||
lance_append_10k_secs: lance_append_secs,
|
||||
lance_append_disk_bytes_added: lance_append_bytes,
|
||||
// From our Phase 15 trial on the SAME index (ec=80 es=30, recall=1.00):
|
||||
reference_hnsw_p50_us: 873.0,
|
||||
reference_hnsw_p95_us: 1413.0,
|
||||
reference_brute_force_us: 43983.0,
|
||||
reference_hnsw_build_secs: 230.0,
|
||||
};
|
||||
|
||||
let json = serde_json::to_string_pretty(&report)?;
|
||||
println!("{}", json);
|
||||
|
||||
eprintln!("\n=== Summary ===");
|
||||
eprintln!(" Parquet cold load: {:.2}s", report.parquet_cold_load_secs);
|
||||
eprintln!(" Lance cold open: {:.2}s ({})",
|
||||
report.lance_cold_open_secs,
|
||||
format_ratio(report.parquet_cold_load_secs, report.lance_cold_open_secs));
|
||||
eprintln!(" Parquet disk: {:.1} MB", report.parquet_disk_bytes as f64 / 1e6);
|
||||
eprintln!(" Lance disk: {:.1} MB ({})",
|
||||
report.lance_disk_bytes as f64 / 1e6,
|
||||
format_ratio(report.parquet_disk_bytes as f32, report.lance_disk_bytes as f32));
|
||||
if let (Some(p50), Some(p95)) = (report.lance_search_p50_us, report.lance_search_p95_us) {
|
||||
eprintln!(" Lance search p50: {:.0}us vs our HNSW {:.0}us ({})",
|
||||
p50, report.reference_hnsw_p50_us,
|
||||
format_ratio(report.reference_hnsw_p50_us, p50));
|
||||
eprintln!(" Lance search p95: {:.0}us vs our HNSW {:.0}us ({})",
|
||||
p95, report.reference_hnsw_p95_us,
|
||||
format_ratio(report.reference_hnsw_p95_us, p95));
|
||||
eprintln!(" Speedup vs brute force: {:.1}× (Lance) vs {:.1}× (HNSW)",
|
||||
report.reference_brute_force_us / p50,
|
||||
report.reference_brute_force_us / report.reference_hnsw_p50_us);
|
||||
}
|
||||
if let Some(build) = report.lance_index_build_secs {
|
||||
eprintln!(" Index build: {:.1}s (Lance IVF_PQ) vs {:.0}s (our HNSW ec=80) ({}× faster)",
|
||||
build, report.reference_hnsw_build_secs, report.reference_hnsw_build_secs / build);
|
||||
}
|
||||
if let (Some(lance_us), Some(parquet_us)) = (report.lance_random_row_access_us, report.parquet_random_row_access_us) {
|
||||
eprintln!(" Random row access: {:.0}us (Lance) vs {:.0}us (Parquet scan) ({})",
|
||||
lance_us, parquet_us, format_ratio(parquet_us, lance_us));
|
||||
}
|
||||
if let Some(append_secs) = report.lance_append_10k_secs {
|
||||
eprintln!(" Append 10K rows: {:.2}s (Lance native) [Parquet would require full rewrite]",
|
||||
append_secs);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn format_ratio(baseline: f32, candidate: f32) -> String {
|
||||
if candidate == 0.0 { return "inf".into(); }
|
||||
let ratio = baseline / candidate;
|
||||
if ratio >= 1.0 {
|
||||
format!("{:.2}× faster/smaller", ratio)
|
||||
} else {
|
||||
format!("{:.2}× slower/larger", 1.0 / ratio)
|
||||
}
|
||||
}
|
||||
|
||||
fn dir_size_bytes(path: &str) -> u64 {
|
||||
fn recurse(p: &std::path::Path) -> u64 {
|
||||
let Ok(meta) = std::fs::metadata(p) else { return 0; };
|
||||
if meta.is_file() { return meta.len(); }
|
||||
let Ok(entries) = std::fs::read_dir(p) else { return 0; };
|
||||
entries
|
||||
.filter_map(|e| e.ok())
|
||||
.map(|e| recurse(&e.path()))
|
||||
.sum()
|
||||
}
|
||||
recurse(std::path::Path::new(path))
|
||||
}
|
||||
|
||||
/// Read the existing vector Parquet (binary-blob format: source, doc_id,
|
||||
/// chunk_idx, chunk_text, vector as Binary bytes).
|
||||
fn read_parquet_vectors(path: &str) -> Result<(Arc<Schema>, Vec<RecordBatch>, usize)> {
|
||||
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
|
||||
use std::fs::File;
|
||||
|
||||
let file = File::open(path).with_context(|| format!("open {path}"))?;
|
||||
let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
|
||||
let schema = builder.schema().clone();
|
||||
let reader = builder.build()?;
|
||||
let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>()?;
|
||||
let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
|
||||
Ok((schema, batches, rows))
|
||||
}
|
||||
|
||||
fn detect_vector_dims(batches: &[RecordBatch]) -> Result<usize> {
|
||||
for batch in batches {
|
||||
let vector_col_idx = batch
|
||||
.schema()
|
||||
.index_of("vector")
|
||||
.context("no 'vector' column in parquet")?;
|
||||
let col = batch.column(vector_col_idx);
|
||||
if let Some(binary) = col.as_any().downcast_ref::<BinaryArray>() {
|
||||
for i in 0..binary.len() {
|
||||
if !binary.is_null(i) {
|
||||
let bytes = binary.value(i);
|
||||
return Ok(bytes.len() / 4); // f32 = 4 bytes
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
anyhow::bail!("could not determine vector dimensions")
|
||||
}
|
||||
|
||||
/// Convert our binary-blob vector representation into Arrow's native
|
||||
/// FixedSizeList<Float32> — that's what Lance expects for vector columns.
|
||||
fn convert_to_fixed_size_list(
|
||||
schema: &Arc<Schema>,
|
||||
batches: Vec<RecordBatch>,
|
||||
dims: usize,
|
||||
) -> Result<(Arc<Schema>, Vec<RecordBatch>)> {
|
||||
// New schema keeps everything identical but replaces the vector column
|
||||
// with a FixedSizeList<Float32, dims>.
|
||||
let new_fields: Vec<Arc<Field>> = schema
|
||||
.fields()
|
||||
.iter()
|
||||
.map(|f| {
|
||||
if f.name() == "vector" {
|
||||
Arc::new(Field::new(
|
||||
"vector",
|
||||
DataType::FixedSizeList(
|
||||
Arc::new(Field::new("item", DataType::Float32, true)),
|
||||
dims as i32,
|
||||
),
|
||||
false,
|
||||
))
|
||||
} else {
|
||||
f.clone()
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let new_schema = Arc::new(Schema::new(new_fields));
|
||||
|
||||
let mut new_batches = Vec::with_capacity(batches.len());
|
||||
for batch in batches {
|
||||
let vector_idx = batch.schema().index_of("vector")?;
|
||||
let mut new_arrays: Vec<ArrayRef> = Vec::with_capacity(batch.num_columns());
|
||||
for (i, col) in batch.columns().iter().enumerate() {
|
||||
if i == vector_idx {
|
||||
let binary = col
|
||||
.as_any()
|
||||
.downcast_ref::<BinaryArray>()
|
||||
.context("vector column must be Binary")?;
|
||||
let fsl = binary_to_fixed_size_list(binary, dims)?;
|
||||
new_arrays.push(Arc::new(fsl));
|
||||
} else {
|
||||
new_arrays.push(col.clone());
|
||||
}
|
||||
}
|
||||
new_batches.push(RecordBatch::try_new(new_schema.clone(), new_arrays)?);
|
||||
}
|
||||
|
||||
Ok((new_schema, new_batches))
|
||||
}
|
||||
|
||||
fn binary_to_fixed_size_list(binary: &BinaryArray, dims: usize) -> Result<FixedSizeListArray> {
|
||||
let n = binary.len();
|
||||
let mut all_floats: Vec<f32> = Vec::with_capacity(n * dims);
|
||||
for i in 0..n {
|
||||
if binary.is_null(i) {
|
||||
all_floats.extend(std::iter::repeat(0.0).take(dims));
|
||||
continue;
|
||||
}
|
||||
let bytes = binary.value(i);
|
||||
if bytes.len() != dims * 4 {
|
||||
anyhow::bail!(
|
||||
"row {} has {} bytes, expected {} ({} × f32)",
|
||||
i, bytes.len(), dims * 4, dims,
|
||||
);
|
||||
}
|
||||
for chunk in bytes.chunks_exact(4) {
|
||||
all_floats.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
|
||||
}
|
||||
}
|
||||
let values = Float32Array::from(all_floats);
|
||||
let field = Arc::new(Field::new("item", DataType::Float32, true));
|
||||
FixedSizeListArray::try_new(field, dims as i32, Arc::new(values), None)
|
||||
.context("build FixedSizeListArray")
|
||||
}
|
||||
|
||||
/// Write batches into a Lance dataset at the given path.
|
||||
async fn write_lance_dataset(
|
||||
path: &str,
|
||||
schema: Arc<Schema>,
|
||||
batches: Vec<RecordBatch>,
|
||||
) -> Result<()> {
|
||||
use lance::dataset::{Dataset, WriteParams};
|
||||
|
||||
let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema);
|
||||
Dataset::write(reader, path, Some(WriteParams::default()))
|
||||
.await
|
||||
.context("Dataset::write")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Open a Lance dataset cold (from disk) and scan it fully — measuring the
|
||||
/// equivalent of our "load embeddings from Parquet" cost.
|
||||
async fn cold_open_and_scan_lance(path: &str) -> Result<usize> {
|
||||
use futures::StreamExt;
|
||||
use lance::dataset::Dataset;
|
||||
|
||||
let dataset = Dataset::open(path).await.context("Dataset::open")?;
|
||||
let scanner = dataset.scan();
|
||||
let mut stream = scanner.try_into_stream().await?;
|
||||
let mut total = 0usize;
|
||||
while let Some(batch) = stream.next().await {
|
||||
let batch = batch?;
|
||||
total += batch.num_rows();
|
||||
}
|
||||
Ok(total)
|
||||
}
|
||||
|
||||
/// Build an IVF_PQ vector index on the `vector` column. IVF_PQ (Inverted File
|
||||
/// with Product Quantization) is Lance's native ANN index — comparable to
|
||||
/// HNSW in intent, but on-disk and compatible with Lance's random-access
|
||||
/// model.
|
||||
async fn build_lance_vector_index(path: &str, _dims: usize) -> Result<()> {
|
||||
use lance::dataset::Dataset;
|
||||
use lance::index::vector::VectorIndexParams;
|
||||
use lance_index::{DatasetIndexExt, IndexType};
|
||||
use lance_linalg::distance::MetricType;
|
||||
|
||||
let mut dataset = Dataset::open(path).await?;
|
||||
|
||||
// IVF_PQ with ~sqrt(N) partitions is a reasonable default for 100K.
|
||||
// num_sub_vectors must divide dims evenly: 768/48 = 16 dims per subvector.
|
||||
// num_bits = 8 gives 256 codes per subvector (good recall/size trade).
|
||||
// max_iterations = 50 is plenty for this scale.
|
||||
let params = VectorIndexParams::ivf_pq(
|
||||
316, // num_partitions (~sqrt(100000))
|
||||
8, // num_bits
|
||||
48, // num_sub_vectors
|
||||
MetricType::Cosine,
|
||||
50, // max_iterations
|
||||
);
|
||||
|
||||
dataset
|
||||
.create_index(
|
||||
&["vector"],
|
||||
IndexType::Vector,
|
||||
Some("vec_idx".into()),
|
||||
¶ms,
|
||||
true,
|
||||
)
|
||||
.await
|
||||
.context("create_index")?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Run N vector searches against the Lance dataset and return (p50, p95, p99) latencies in us.
|
||||
/// Uses a handful of random rows as queries — same pattern as our harness::synthetic_from_chunks.
|
||||
async fn run_search_benchmarks(path: &str, _dims: usize) -> Result<(f32, f32, f32)> {
|
||||
use futures::StreamExt;
|
||||
use lance::dataset::Dataset;
|
||||
|
||||
let dataset = Dataset::open(path).await?;
|
||||
|
||||
// Pick 20 representative query vectors from the data itself.
|
||||
// (Synthetic — same pattern as our existing harness.)
|
||||
let query_vectors = sample_query_vectors(&dataset, 20).await?;
|
||||
|
||||
let mut latencies_us: Vec<f32> = Vec::with_capacity(query_vectors.len());
|
||||
for (i, qv) in query_vectors.iter().enumerate() {
|
||||
let qarr = Arc::new(Float32Array::from(qv.clone())) as ArrayRef;
|
||||
|
||||
let t0 = Instant::now();
|
||||
let mut scanner = dataset.scan();
|
||||
scanner
|
||||
.nearest("vector", qarr.as_any().downcast_ref::<Float32Array>().unwrap(), 10)
|
||||
.context("scanner.nearest")?;
|
||||
let mut stream = scanner.try_into_stream().await?;
|
||||
let mut hits = 0;
|
||||
while let Some(batch) = stream.next().await {
|
||||
let batch = batch?;
|
||||
hits += batch.num_rows();
|
||||
}
|
||||
let us = t0.elapsed().as_micros() as f32;
|
||||
latencies_us.push(us);
|
||||
if i == 0 {
|
||||
eprintln!(" first query: {} hits in {:.0}us (includes any lazy init)", hits, us);
|
||||
}
|
||||
}
|
||||
|
||||
latencies_us.sort_by(|a, b| a.partial_cmp(b).unwrap());
|
||||
let p = |pct: f32| -> f32 {
|
||||
let idx = ((latencies_us.len() as f32 - 1.0) * pct).round() as usize;
|
||||
latencies_us[idx.min(latencies_us.len() - 1)]
|
||||
};
|
||||
Ok((p(0.50), p(0.95), p(0.99)))
|
||||
}
|
||||
|
||||
/// Random row access via Lance's `take` — fetch 20 random rows by index, measure avg latency.
|
||||
async fn measure_random_access_lance(path: &str) -> Result<f32> {
|
||||
use lance::dataset::Dataset;
|
||||
let dataset = Dataset::open(path).await?;
|
||||
let n = dataset.count_rows(None).await?;
|
||||
let indices: Vec<u64> = (0..20).map(|i| ((i as u64) * (n as u64 / 23)) % (n as u64)).collect();
|
||||
|
||||
// Full-schema projection — Lance's Schema implements Into<ProjectionRequest>.
|
||||
let schema = dataset.schema().clone();
|
||||
let mut total_us: u128 = 0;
|
||||
for idx in &indices {
|
||||
let t0 = Instant::now();
|
||||
let _batch = dataset.take(&[*idx], schema.clone()).await?;
|
||||
total_us += t0.elapsed().as_micros();
|
||||
}
|
||||
Ok(total_us as f32 / indices.len() as f32)
|
||||
}
|
||||
|
||||
/// Random row access for Parquet — full scan + filter. There's no random-access
|
||||
/// primitive in vanilla Parquet, so this is the cost of finding one specific row.
|
||||
/// This is the cost our current design pays for "get doc X's full text for RAG."
|
||||
fn measure_random_access_parquet(path: &str) -> Result<f32> {
|
||||
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
|
||||
use std::fs::File;
|
||||
|
||||
// We simulate 5 lookups — full scan each time. 20 would be painful.
|
||||
let iters = 5;
|
||||
let mut total_us: u128 = 0;
|
||||
for _ in 0..iters {
|
||||
let t0 = Instant::now();
|
||||
let file = File::open(path)?;
|
||||
let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
|
||||
let reader = builder.build()?;
|
||||
// Iterate until we've conceptually found a row — we stop early if
|
||||
// we wanted row 50000, but we have to at least read its batch.
|
||||
let mut seen = 0usize;
|
||||
for b in reader {
|
||||
let b = b?;
|
||||
seen += b.num_rows();
|
||||
if seen > 50000 { break; }
|
||||
}
|
||||
total_us += t0.elapsed().as_micros();
|
||||
}
|
||||
Ok(total_us as f32 / iters as f32)
|
||||
}
|
||||
|
||||
/// Append 10K new rows to the existing Lance dataset.
|
||||
/// Measures the "ingest delta" cost without full rewrite.
|
||||
async fn append_10k_rows(path: &str, dims: usize) -> Result<()> {
|
||||
use lance::dataset::{Dataset, WriteMode, WriteParams};
|
||||
|
||||
let dataset = Dataset::open(path).await?;
|
||||
let schema = dataset.schema();
|
||||
let arrow_schema: Arc<Schema> = Arc::new(schema.into());
|
||||
|
||||
// Build a 10K row batch with random-ish data matching the existing schema.
|
||||
let n = 10_000;
|
||||
let arrays: Vec<ArrayRef> = arrow_schema
|
||||
.fields()
|
||||
.iter()
|
||||
.map(|f| -> Result<ArrayRef> {
|
||||
match f.data_type() {
|
||||
DataType::Utf8 => {
|
||||
let vals: Vec<String> = (0..n).map(|i| format!("appended-{}", i)).collect();
|
||||
Ok(Arc::new(arrow_array::StringArray::from(vals)))
|
||||
}
|
||||
DataType::Int32 => {
|
||||
let vals: Vec<i32> = (0..n as i32).collect();
|
||||
Ok(Arc::new(arrow_array::Int32Array::from(vals)))
|
||||
}
|
||||
DataType::FixedSizeList(_, _) => {
|
||||
let floats: Vec<f32> = (0..n * dims).map(|i| (i as f32).sin()).collect();
|
||||
let values = Float32Array::from(floats);
|
||||
let field = Arc::new(Field::new("item", DataType::Float32, true));
|
||||
let fsl = FixedSizeListArray::try_new(field, dims as i32, Arc::new(values), None)?;
|
||||
Ok(Arc::new(fsl))
|
||||
}
|
||||
other => anyhow::bail!("unsupported append column type: {:?}", other),
|
||||
}
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
|
||||
let batch = RecordBatch::try_new(arrow_schema.clone(), arrays)?;
|
||||
let reader = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), arrow_schema);
|
||||
let params = WriteParams { mode: WriteMode::Append, ..Default::default() };
|
||||
Dataset::write(reader, path, Some(params)).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Grab a few existing vectors from the dataset to use as self-similar queries.
|
||||
async fn sample_query_vectors(
|
||||
dataset: &lance::dataset::Dataset,
|
||||
count: usize,
|
||||
) -> Result<Vec<Vec<f32>>> {
|
||||
use futures::StreamExt;
|
||||
|
||||
// Just take the first `count` rows; good enough for latency measurement.
|
||||
let scanner = dataset.scan();
|
||||
let mut scanner = scanner;
|
||||
scanner.limit(Some(count as i64), None)?;
|
||||
scanner.project(&["vector"])?;
|
||||
let mut stream = scanner.try_into_stream().await?;
|
||||
|
||||
let mut out = Vec::with_capacity(count);
|
||||
while let Some(batch) = stream.next().await {
|
||||
let batch = batch?;
|
||||
let vector_col = batch
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<FixedSizeListArray>()
|
||||
.context("vector column must be FixedSizeList")?;
|
||||
|
||||
for row in 0..vector_col.len() {
|
||||
if out.len() >= count { break; }
|
||||
let values = vector_col.value(row);
|
||||
let f32_arr = values
|
||||
.as_any()
|
||||
.downcast_ref::<Float32Array>()
|
||||
.context("inner array must be Float32")?;
|
||||
let mut v = Vec::with_capacity(f32_arr.len());
|
||||
for i in 0..f32_arr.len() {
|
||||
v.push(f32_arr.value(i));
|
||||
}
|
||||
out.push(v);
|
||||
}
|
||||
if out.len() >= count { break; }
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
@@ -15,6 +15,7 @@ tracing = { workspace = true }
|
||||
datafusion = { workspace = true }
|
||||
object_store = { workspace = true }
|
||||
arrow = { workspace = true }
|
||||
parquet = { workspace = true }
|
||||
bytes = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
url = { workspace = true }
|
||||
|
||||
+211
-16
@@ -5,6 +5,7 @@ use datafusion::datasource::file_format::parquet::ParquetFormat;
|
||||
use datafusion::prelude::*;
|
||||
use object_store::ObjectStore;
|
||||
use std::sync::Arc;
|
||||
use storaged::registry::BucketRegistry;
|
||||
use url::Url;
|
||||
|
||||
use crate::cache::MemCache;
|
||||
@@ -12,27 +13,61 @@ use crate::delta;
|
||||
|
||||
const STORE_SCHEME: &str = "lakehouse";
|
||||
|
||||
/// URL scheme used to register a bucket's object store with DataFusion.
|
||||
/// Bucket names safely embed in URL host components after sanitizing
|
||||
/// `:` (used in `profile:user`) to `-`.
|
||||
fn bucket_scheme(bucket: &str) -> String {
|
||||
format!("{STORE_SCHEME}-{}", bucket.replace(':', "-"))
|
||||
}
|
||||
|
||||
/// Build a `ListingTableUrl`-shaped string for an object in a bucket.
|
||||
fn bucket_object_url(bucket: &str, key: &str) -> String {
|
||||
format!("{}://data/{}", bucket_scheme(bucket), key)
|
||||
}
|
||||
|
||||
/// Build the registration base URL for a bucket.
|
||||
fn bucket_base_url(bucket: &str) -> String {
|
||||
format!("{}://data/", bucket_scheme(bucket))
|
||||
}
|
||||
|
||||
/// Query engine with in-memory cache and delta merge support.
|
||||
/// Federation layer 2: holds a `BucketRegistry` so cross-bucket queries
|
||||
/// route reads to the right ObjectStore based on `ObjectRef.bucket`.
|
||||
#[derive(Clone)]
|
||||
pub struct QueryEngine {
|
||||
registry: Registry,
|
||||
buckets: Arc<BucketRegistry>,
|
||||
/// Primary bucket store cached for backward-compat callers (compact,
|
||||
/// workspace ops) that haven't been migrated to bucket-aware paths.
|
||||
store: Arc<dyn ObjectStore>,
|
||||
cache: MemCache,
|
||||
}
|
||||
|
||||
impl QueryEngine {
|
||||
pub fn new(registry: Registry, store: Arc<dyn ObjectStore>, cache: MemCache) -> Self {
|
||||
Self { registry, store, cache }
|
||||
pub fn new(registry: Registry, buckets: Arc<BucketRegistry>, cache: MemCache) -> Self {
|
||||
let store = buckets.default_store();
|
||||
Self { registry, buckets, store, cache }
|
||||
}
|
||||
|
||||
pub fn cache(&self) -> &MemCache {
|
||||
&self.cache
|
||||
}
|
||||
|
||||
/// Backward-compat: returns the primary bucket's store. Bucket-aware
|
||||
/// callers should use `engine.buckets()` directly.
|
||||
pub fn store(&self) -> &Arc<dyn ObjectStore> {
|
||||
&self.store
|
||||
}
|
||||
|
||||
pub fn buckets(&self) -> &Arc<BucketRegistry> {
|
||||
&self.buckets
|
||||
}
|
||||
|
||||
/// Catalog registry — used by compact/tombstone integration.
|
||||
pub fn catalog(&self) -> &Registry {
|
||||
&self.registry
|
||||
}
|
||||
|
||||
/// Execute a SQL query. Uses cache for hot data, falls back to Parquet.
|
||||
pub async fn query(&self, sql: &str) -> Result<Vec<arrow::array::RecordBatch>, String> {
|
||||
let ctx = self.build_context().await?;
|
||||
@@ -43,11 +78,16 @@ impl QueryEngine {
|
||||
|
||||
/// Pin a dataset into the memory cache.
|
||||
pub async fn pin_dataset(&self, name: &str) -> Result<(), String> {
|
||||
// Read from Parquet
|
||||
let ctx = SessionContext::new();
|
||||
let base_url = Url::parse(&format!("{STORE_SCHEME}://data/"))
|
||||
.map_err(|e| format!("invalid store url: {e}"))?;
|
||||
ctx.runtime_env().register_object_store(&base_url, self.store.clone());
|
||||
|
||||
// Register every bucket so a multi-bucket dataset can be pinned.
|
||||
for info in self.buckets.list().await {
|
||||
let url = Url::parse(&bucket_base_url(&info.name))
|
||||
.map_err(|e| format!("invalid store url: {e}"))?;
|
||||
let store = self.buckets.get(&info.name)
|
||||
.map_err(|e| format!("registry inconsistency: {e}"))?;
|
||||
ctx.runtime_env().register_object_store(&url, store);
|
||||
}
|
||||
|
||||
let dataset = self.registry.get_by_name(name).await
|
||||
.ok_or_else(|| format!("dataset not found: {name}"))?;
|
||||
@@ -58,7 +98,10 @@ impl QueryEngine {
|
||||
|
||||
let opts = ListingOptions::new(Arc::new(ParquetFormat::default()));
|
||||
let table_paths: Vec<ListingTableUrl> = dataset.objects.iter()
|
||||
.filter_map(|o| ListingTableUrl::parse(&format!("{STORE_SCHEME}://data/{}", o.key)).ok())
|
||||
.filter_map(|o| {
|
||||
let bucket = if o.bucket.is_empty() { "primary" } else { &o.bucket };
|
||||
ListingTableUrl::parse(&bucket_object_url(bucket, &o.key)).ok()
|
||||
})
|
||||
.collect();
|
||||
|
||||
let schema = opts.infer_schema(&ctx.state(), &table_paths[0]).await
|
||||
@@ -79,9 +122,41 @@ impl QueryEngine {
|
||||
async fn build_context(&self) -> Result<SessionContext, String> {
|
||||
let ctx = SessionContext::new();
|
||||
|
||||
let base_url = Url::parse(&format!("{STORE_SCHEME}://data/"))
|
||||
.map_err(|e| format!("invalid store url: {e}"))?;
|
||||
ctx.runtime_env().register_object_store(&base_url, self.store.clone());
|
||||
// Phase E: snapshot tombstones by dataset before registering tables
|
||||
// so we can wrap tombstoned tables in a filter view. The underlying
|
||||
// base table is registered under an internal name `__raw__{dataset}`
|
||||
// and the public `{dataset}` name becomes the filtered view.
|
||||
let all_dataset_names: Vec<String> = self.registry.list().await
|
||||
.iter().map(|d| d.name.clone()).collect();
|
||||
let tombstones_by_dataset = self
|
||||
.registry
|
||||
.tombstones()
|
||||
.all_grouped(&all_dataset_names)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
// Federation layer 2: register every configured bucket as its own
|
||||
// DataFusion ObjectStore under a distinct URL scheme. Each
|
||||
// dataset's ObjectRef.bucket determines which store DataFusion
|
||||
// routes to at scan time.
|
||||
let bucket_infos = self.buckets.list().await;
|
||||
for info in &bucket_infos {
|
||||
let url_str = bucket_base_url(&info.name);
|
||||
let url = Url::parse(&url_str)
|
||||
.map_err(|e| format!("invalid store url '{url_str}': {e}"))?;
|
||||
// unknown-bucket here would be a config invariant violation,
|
||||
// since we just listed them — propagate as an error to surface it.
|
||||
let store = self.buckets.get(&info.name)
|
||||
.map_err(|e| format!("registry inconsistency: {e}"))?;
|
||||
ctx.runtime_env().register_object_store(&url, store);
|
||||
}
|
||||
|
||||
// Backward-compat: also register the legacy `lakehouse://data/`
|
||||
// URL pointing at primary, in case any code path still constructs
|
||||
// ListingTableUrls without a bucket prefix.
|
||||
let legacy_url = Url::parse(&format!("{STORE_SCHEME}://data/"))
|
||||
.map_err(|e| format!("invalid legacy store url: {e}"))?;
|
||||
ctx.runtime_env().register_object_store(&legacy_url, self.store.clone());
|
||||
|
||||
let datasets = self.registry.list().await;
|
||||
for dataset in &datasets {
|
||||
@@ -89,9 +164,9 @@ impl QueryEngine {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Check cache first
|
||||
// Check cache first — cached batches are bucket-agnostic since
|
||||
// they're already materialized in memory.
|
||||
if let Some(cached) = self.cache.get(&dataset.name).await {
|
||||
// Load any delta files and merge
|
||||
let delta_batches = delta::load_deltas(&self.store, &dataset.name).await.unwrap_or_default();
|
||||
|
||||
let mut all_batches = cached.batches;
|
||||
@@ -106,10 +181,15 @@ impl QueryEngine {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Fall back to Parquet file
|
||||
// Build ListingTable URLs that include the per-object bucket
|
||||
// scheme. DataFusion routes each scan to the right store
|
||||
// automatically based on the URL.
|
||||
let opts = ListingOptions::new(Arc::new(ParquetFormat::default()));
|
||||
let table_paths: Vec<ListingTableUrl> = dataset.objects.iter()
|
||||
.filter_map(|o| ListingTableUrl::parse(&format!("{STORE_SCHEME}://data/{}", o.key)).ok())
|
||||
.filter_map(|o| {
|
||||
let bucket = if o.bucket.is_empty() { "primary" } else { &o.bucket };
|
||||
ListingTableUrl::parse(&bucket_object_url(bucket, &o.key)).ok()
|
||||
})
|
||||
.collect();
|
||||
|
||||
if table_paths.is_empty() {
|
||||
@@ -124,8 +204,123 @@ impl QueryEngine {
|
||||
let table = ListingTable::try_new(config)
|
||||
.map_err(|e| format!("table creation failed for {}: {e}", dataset.name))?;
|
||||
|
||||
ctx.register_table(&dataset.name, Arc::new(table))
|
||||
.map_err(|e| format!("table registration failed for {}: {e}", dataset.name))?;
|
||||
// Decide the registration name: if this dataset has any
|
||||
// tombstones, the raw table gets an internal name and the
|
||||
// public name becomes a filtered view.
|
||||
let tombstone_entry = tombstones_by_dataset.get(&dataset.name);
|
||||
let register_name = if tombstone_entry.is_some() {
|
||||
format!("__raw__{}", dataset.name)
|
||||
} else {
|
||||
dataset.name.clone()
|
||||
};
|
||||
|
||||
if let Err(e) = ctx.register_table(register_name.as_str(), Arc::new(table)) {
|
||||
let msg = e.to_string();
|
||||
if msg.contains("already exists") {
|
||||
tracing::debug!("skip duplicate manifest registration: {}", register_name);
|
||||
continue;
|
||||
} else {
|
||||
return Err(format!("table registration failed for {}: {}", register_name, msg));
|
||||
}
|
||||
}
|
||||
|
||||
// If there are tombstones, register the public name as a
|
||||
// filtered view that excludes tombstoned row_key_values.
|
||||
if let Some((key_col, values)) = tombstone_entry {
|
||||
// Build WHERE NOT IN (...) — quote values to be SQL-safe.
|
||||
// For string keys this is a literal list; for integer keys
|
||||
// CAST(col AS VARCHAR) makes the comparison unambiguous.
|
||||
let quoted: Vec<String> = values.iter()
|
||||
.map(|v| format!("'{}'", v.replace('\'', "''")))
|
||||
.collect();
|
||||
let sql = format!(
|
||||
"SELECT * FROM \"{}\" WHERE CAST(\"{}\" AS VARCHAR) NOT IN ({})",
|
||||
register_name, key_col, quoted.join(", "),
|
||||
);
|
||||
tracing::debug!(
|
||||
"tombstone filter for '{}': {} row_keys excluded",
|
||||
dataset.name, values.len(),
|
||||
);
|
||||
match ctx.sql(&sql).await {
|
||||
Ok(df) => {
|
||||
if let Err(e) = ctx.register_table(dataset.name.as_str(), df.into_view()) {
|
||||
let msg = e.to_string();
|
||||
if msg.contains("already exists") {
|
||||
tracing::debug!("skip duplicate tombstone view: {}", dataset.name);
|
||||
} else {
|
||||
tracing::warn!("tombstone view registration failed for {}: {msg}", dataset.name);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("tombstone view SQL failed for '{}': {e}", dataset.name);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Phase D: register every AiView as a DataFusion view that wraps
|
||||
// the base table with the safe projection + filter + redactions.
|
||||
for view in self.registry.list_views().await {
|
||||
// Build the SELECT clause with column whitelist and any
|
||||
// redaction expressions. Each column is either the raw column
|
||||
// name or an expression aliased back to the column name so
|
||||
// downstream queries see the same shape.
|
||||
let select_cols: Vec<String> = view.columns.iter().map(|col| {
|
||||
match view.column_redactions.get(col) {
|
||||
None => format!("\"{}\"", col),
|
||||
Some(shared::types::Redaction::Null) => {
|
||||
format!("CAST(NULL AS VARCHAR) AS \"{}\"", col)
|
||||
}
|
||||
Some(shared::types::Redaction::Hash) => {
|
||||
// DataFusion has digest('sha256', value) — but column
|
||||
// type might not be string. Cast first to be safe.
|
||||
format!("digest(CAST(\"{}\" AS VARCHAR), 'sha256') AS \"{}\"", col, col)
|
||||
}
|
||||
Some(shared::types::Redaction::Mask { keep_prefix, keep_suffix }) => {
|
||||
// Keep first N + last M chars, replace middle with stars.
|
||||
// CASE NULL through with COALESCE for null-safe behavior.
|
||||
format!(
|
||||
"CASE WHEN \"{c}\" IS NULL THEN NULL ELSE \
|
||||
concat(\
|
||||
substr(CAST(\"{c}\" AS VARCHAR), 1, {p}), \
|
||||
repeat('*', greatest(0, length(CAST(\"{c}\" AS VARCHAR)) - {p} - {s})), \
|
||||
substr(CAST(\"{c}\" AS VARCHAR), greatest(1, length(CAST(\"{c}\" AS VARCHAR)) - {s} + 1))\
|
||||
) END AS \"{c}\"",
|
||||
c = col, p = keep_prefix, s = keep_suffix,
|
||||
)
|
||||
}
|
||||
}
|
||||
}).collect();
|
||||
|
||||
let where_clause = view.row_filter
|
||||
.as_deref()
|
||||
.map(|f| format!(" WHERE {}", f))
|
||||
.unwrap_or_default();
|
||||
let view_sql = format!(
|
||||
"SELECT {} FROM \"{}\"{}",
|
||||
select_cols.join(", "),
|
||||
view.base_dataset,
|
||||
where_clause,
|
||||
);
|
||||
|
||||
tracing::debug!("registering AiView '{}': {}", view.name, view_sql);
|
||||
|
||||
let df = match ctx.sql(&view_sql).await {
|
||||
Ok(df) => df,
|
||||
Err(e) => {
|
||||
tracing::warn!("view '{}' SQL invalid, skipping: {e}", view.name);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
if let Err(e) = ctx.register_table(view.name.as_str(), df.into_view()) {
|
||||
let msg = e.to_string();
|
||||
if msg.contains("already exists") {
|
||||
tracing::debug!("skip duplicate view registration: {}", view.name);
|
||||
} else {
|
||||
tracing::warn!("view '{}' registration failed: {e}", view.name);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(ctx)
|
||||
|
||||
+160
-22
@@ -3,10 +3,15 @@
|
||||
/// we write small delta files. At query time, deltas are merged with the base.
|
||||
/// Periodic compaction merges deltas into the base file.
|
||||
|
||||
use arrow::array::RecordBatch;
|
||||
use arrow::datatypes::SchemaRef;
|
||||
use arrow::array::{Array, BooleanArray, Int32Array, Int64Array, RecordBatch, StringArray};
|
||||
use arrow::compute::filter_record_batch;
|
||||
use bytes::Bytes;
|
||||
use object_store::ObjectStore;
|
||||
use parquet::arrow::ArrowWriter;
|
||||
use parquet::basic::Compression;
|
||||
use parquet::file::properties::WriterProperties;
|
||||
use shared::types::Tombstone;
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
|
||||
use shared::arrow_helpers::{parquet_to_record_batches, record_batch_to_parquet};
|
||||
@@ -58,74 +63,207 @@ pub async fn load_deltas(
|
||||
}
|
||||
|
||||
/// Compact: merge base Parquet + all deltas into a single new base file.
|
||||
/// Optionally deduplicates by a primary key column.
|
||||
/// Optionally filters out tombstoned rows — this is the physical-deletion
|
||||
/// half of Phase E (the query-time filter handles the "immediate hide",
|
||||
/// compaction handles the "actually remove from disk" requirement).
|
||||
///
|
||||
/// Tombstones must share one `row_key_column` (enforced at tombstone write).
|
||||
/// If non-empty, every row in the merged output whose row_key_column value
|
||||
/// matches a tombstone is dropped before the new base is written.
|
||||
pub async fn compact(
|
||||
store: &Arc<dyn ObjectStore>,
|
||||
dataset_name: &str,
|
||||
base_key: &str,
|
||||
primary_key_col: Option<&str>,
|
||||
tombstones: &[Tombstone],
|
||||
) -> Result<CompactResult, String> {
|
||||
// Load base
|
||||
let base_data = ops::get(store, base_key).await?;
|
||||
let (schema, mut base_batches) = parquet_to_record_batches(&base_data)?;
|
||||
let (_schema, mut base_batches) = parquet_to_record_batches(&base_data)?;
|
||||
|
||||
// Load deltas
|
||||
let delta_batches = load_deltas(store, dataset_name).await?;
|
||||
let delta_count = delta_batches.len();
|
||||
|
||||
if delta_batches.is_empty() {
|
||||
let has_tombstones = !tombstones.is_empty();
|
||||
let nothing_to_do = delta_batches.is_empty() && !has_tombstones;
|
||||
if nothing_to_do {
|
||||
return Ok(CompactResult {
|
||||
base_rows: base_batches.iter().map(|b| b.num_rows()).sum(),
|
||||
delta_rows: 0,
|
||||
final_rows: base_batches.iter().map(|b| b.num_rows()).sum(),
|
||||
deltas_merged: 0,
|
||||
tombstones_applied: 0,
|
||||
rows_dropped_by_tombstones: 0,
|
||||
});
|
||||
}
|
||||
|
||||
base_batches.extend(delta_batches);
|
||||
|
||||
let base_rows: usize = base_batches.iter().map(|b| b.num_rows()).sum();
|
||||
let pre_filter_rows: usize = base_batches.iter().map(|b| b.num_rows()).sum();
|
||||
|
||||
// If primary key specified, deduplicate (keep last occurrence)
|
||||
let final_batches = if let Some(_pk) = primary_key_col {
|
||||
// For now, just concatenate. Full dedup requires sorting by PK
|
||||
// and keeping the last row per key — this is a simplification.
|
||||
// TODO: implement proper merge with dedup
|
||||
let merged_batches = if let Some(_pk) = primary_key_col {
|
||||
// Current simplification — full dedup requires sort by PK; leave
|
||||
// that work for a follow-up. Tombstone filter still runs below.
|
||||
base_batches
|
||||
} else {
|
||||
base_batches
|
||||
};
|
||||
|
||||
// Tombstone filter: drop rows whose row_key matches.
|
||||
let (final_batches, dropped) = if has_tombstones {
|
||||
apply_tombstone_filter(merged_batches, tombstones)?
|
||||
} else {
|
||||
(merged_batches, 0)
|
||||
};
|
||||
|
||||
let final_rows: usize = final_batches.iter().map(|b| b.num_rows()).sum();
|
||||
|
||||
// Write merged base
|
||||
let mut merged_parquet = Vec::new();
|
||||
for batch in &final_batches {
|
||||
let pq = record_batch_to_parquet(batch)?;
|
||||
merged_parquet.extend_from_slice(&pq);
|
||||
}
|
||||
ops::put(store, base_key, Bytes::from(merged_parquet)).await?;
|
||||
// Write merged base atomically: build ONE Parquet file (single
|
||||
// footer!) via ArrowWriter, stage it under a temporary key, read
|
||||
// back to verify row count matches what we intended, THEN swap
|
||||
// into base_key. If anything goes wrong between build and swap,
|
||||
// the original base_key is untouched.
|
||||
//
|
||||
// Why this matters: concatenating N record_batch_to_parquet()
|
||||
// outputs produces N independent Parquet files glued together —
|
||||
// readers see only the first footer. That silently corrupted the
|
||||
// candidates table before this code was fixed.
|
||||
let schema_ref = final_batches
|
||||
.first()
|
||||
.map(|b| b.schema())
|
||||
.ok_or_else(|| "no batches to write".to_string())?;
|
||||
|
||||
// Delete delta files
|
||||
// Use Snappy to match the compression of files written by ingest —
|
||||
// otherwise compaction would inflate file size 3× on each rewrite.
|
||||
let props = WriterProperties::builder()
|
||||
.set_compression(Compression::SNAPPY)
|
||||
.build();
|
||||
|
||||
let mut buf: Vec<u8> = Vec::with_capacity(16 * 1024 * 1024);
|
||||
{
|
||||
let mut writer = ArrowWriter::try_new(&mut buf, schema_ref.clone(), Some(props))
|
||||
.map_err(|e| format!("ArrowWriter init: {e}"))?;
|
||||
for batch in &final_batches {
|
||||
writer.write(batch).map_err(|e| format!("write batch: {e}"))?;
|
||||
}
|
||||
writer.close().map_err(|e| format!("close writer: {e}"))?;
|
||||
}
|
||||
|
||||
// Verify before committing — if the written bytes don't parse as
|
||||
// the row count we expect, DON'T overwrite the base file.
|
||||
let (_, verify_batches) = parquet_to_record_batches(&buf)
|
||||
.map_err(|e| format!("verify parse: {e}"))?;
|
||||
let verify_rows: usize = verify_batches.iter().map(|b| b.num_rows()).sum();
|
||||
if verify_rows != final_rows {
|
||||
return Err(format!(
|
||||
"compact verification failed: wrote {} bytes that parse as {} rows (expected {}); base_key '{}' untouched",
|
||||
buf.len(), verify_rows, final_rows, base_key,
|
||||
));
|
||||
}
|
||||
|
||||
// Stage under a temp key first, then promote via a second put
|
||||
// over base_key. Object stores don't have atomic rename across
|
||||
// arbitrary backends; this is "write twice + delete temp" which
|
||||
// is the widely-adopted pattern.
|
||||
let temp_key = format!("{base_key}.compact-{}.tmp", chrono::Utc::now().timestamp_millis());
|
||||
ops::put(store, &temp_key, Bytes::from(buf.clone())).await?;
|
||||
ops::put(store, base_key, Bytes::from(buf)).await?;
|
||||
let _ = ops::delete(store, &temp_key).await;
|
||||
|
||||
// Only now that base is durably updated, delete delta files.
|
||||
let delta_keys = list_deltas(store, dataset_name).await?;
|
||||
for key in &delta_keys {
|
||||
let _ = ops::delete(store, key).await;
|
||||
}
|
||||
|
||||
tracing::info!("compacted '{}': {} deltas merged, {} → {} rows", dataset_name, delta_count, base_rows, final_rows);
|
||||
tracing::info!(
|
||||
"compacted '{}': {} deltas merged, {} tombstones applied ({} rows dropped), {} → {} rows",
|
||||
dataset_name, delta_count, tombstones.len(), dropped, pre_filter_rows, final_rows,
|
||||
);
|
||||
|
||||
Ok(CompactResult {
|
||||
base_rows,
|
||||
delta_rows: final_rows - base_rows + delta_count, // approximate
|
||||
base_rows: pre_filter_rows - delta_count, // rough base-before-deltas
|
||||
delta_rows: delta_count,
|
||||
final_rows,
|
||||
deltas_merged: delta_count,
|
||||
tombstones_applied: tombstones.len(),
|
||||
rows_dropped_by_tombstones: dropped,
|
||||
})
|
||||
}
|
||||
|
||||
/// Filter every batch in `batches` to exclude rows whose `row_key_column`
|
||||
/// value appears in the tombstone set. All tombstones must share one
|
||||
/// row_key_column — enforced at tombstone write time.
|
||||
fn apply_tombstone_filter(
|
||||
batches: Vec<RecordBatch>,
|
||||
tombstones: &[Tombstone],
|
||||
) -> Result<(Vec<RecordBatch>, usize), String> {
|
||||
let col_name = &tombstones[0].row_key_column;
|
||||
let bad_values: HashSet<String> = tombstones
|
||||
.iter()
|
||||
.map(|t| t.row_key_value.clone())
|
||||
.collect();
|
||||
|
||||
let mut out = Vec::with_capacity(batches.len());
|
||||
let mut total_dropped = 0usize;
|
||||
|
||||
for batch in batches {
|
||||
let col_idx = match batch.schema().index_of(col_name) {
|
||||
Ok(i) => i,
|
||||
Err(_) => {
|
||||
// Column not in this batch — nothing to drop, pass through.
|
||||
out.push(batch);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let col = batch.column(col_idx);
|
||||
|
||||
// Build a "keep" mask: true where row is NOT tombstoned.
|
||||
let mask: Vec<bool> = if let Some(s) = col.as_any().downcast_ref::<StringArray>() {
|
||||
(0..s.len()).map(|i| {
|
||||
s.is_null(i) || !bad_values.contains(s.value(i))
|
||||
}).collect()
|
||||
} else if let Some(a) = col.as_any().downcast_ref::<Int32Array>() {
|
||||
(0..a.len()).map(|i| {
|
||||
a.is_null(i) || !bad_values.contains(&a.value(i).to_string())
|
||||
}).collect()
|
||||
} else if let Some(a) = col.as_any().downcast_ref::<Int64Array>() {
|
||||
(0..a.len()).map(|i| {
|
||||
a.is_null(i) || !bad_values.contains(&a.value(i).to_string())
|
||||
}).collect()
|
||||
} else {
|
||||
// Key column of an unsupported type — log and pass through
|
||||
// rather than panic. Operator has to handle this manually.
|
||||
tracing::warn!(
|
||||
"tombstone filter: column '{}' has unsupported type {:?}; passing rows through",
|
||||
col_name, col.data_type(),
|
||||
);
|
||||
out.push(batch);
|
||||
continue;
|
||||
};
|
||||
|
||||
let kept = mask.iter().filter(|b| **b).count();
|
||||
let dropped = mask.len() - kept;
|
||||
total_dropped += dropped;
|
||||
if dropped == 0 {
|
||||
out.push(batch);
|
||||
continue;
|
||||
}
|
||||
let bool_arr = BooleanArray::from(mask);
|
||||
let filtered = filter_record_batch(&batch, &bool_arr)
|
||||
.map_err(|e| format!("filter batch: {e}"))?;
|
||||
out.push(filtered);
|
||||
}
|
||||
Ok((out, total_dropped))
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct CompactResult {
|
||||
pub base_rows: usize,
|
||||
pub delta_rows: usize,
|
||||
pub final_rows: usize,
|
||||
pub deltas_merged: usize,
|
||||
pub tombstones_applied: usize,
|
||||
pub rows_dropped_by_tombstones: usize,
|
||||
}
|
||||
|
||||
@@ -184,13 +184,31 @@ async fn compact_dataset(
|
||||
State(state): State<QueryState>,
|
||||
Json(req): Json<CompactRequest>,
|
||||
) -> impl IntoResponse {
|
||||
// Phase E: pull tombstones for this dataset and let compact physically
|
||||
// drop those rows. After a successful rewrite, clear the tombstone log
|
||||
// — the rows are gone from disk, the tombstones have done their job.
|
||||
let tombstones = state
|
||||
.engine
|
||||
.catalog()
|
||||
.list_tombstones(&req.dataset)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
match delta::compact(
|
||||
state.engine.store(),
|
||||
&req.dataset,
|
||||
&req.base_key,
|
||||
req.primary_key.as_deref(),
|
||||
&tombstones,
|
||||
).await {
|
||||
Ok(result) => Ok(Json(result)),
|
||||
Ok(result) => {
|
||||
if result.rows_dropped_by_tombstones > 0 {
|
||||
if let Err(e) = state.engine.catalog().tombstones().clear(&req.dataset).await {
|
||||
tracing::warn!("post-compact tombstone clear failed: {e}");
|
||||
}
|
||||
}
|
||||
Ok(Json(result))
|
||||
}
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -15,3 +15,5 @@ bytes = { workspace = true }
|
||||
sha2 = { workspace = true }
|
||||
toml = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
async-trait = "0.1"
|
||||
|
||||
@@ -16,8 +16,46 @@ pub struct Config {
|
||||
pub auth: AuthConfig,
|
||||
#[serde(default)]
|
||||
pub observability: ObservabilityConfig,
|
||||
#[serde(default)]
|
||||
pub agent: AgentSettings,
|
||||
}
|
||||
|
||||
/// Phase 16.2 — background autotune agent settings.
|
||||
///
|
||||
/// Duplicated from `vectord::agent::AgentConfig` because `shared` can't
|
||||
/// depend on `vectord` (vectord already depends on shared). The gateway
|
||||
/// copies these into the vectord config at startup.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct AgentSettings {
|
||||
#[serde(default)]
|
||||
pub enabled: bool,
|
||||
#[serde(default = "default_cycle_interval_secs")]
|
||||
pub cycle_interval_secs: u64,
|
||||
#[serde(default = "default_cooldown_secs")]
|
||||
pub cooldown_between_trials_secs: u64,
|
||||
#[serde(default = "default_min_recall")]
|
||||
pub min_recall: f32,
|
||||
#[serde(default = "default_max_trials_per_hour")]
|
||||
pub max_trials_per_hour: u32,
|
||||
}
|
||||
|
||||
impl Default for AgentSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: false,
|
||||
cycle_interval_secs: default_cycle_interval_secs(),
|
||||
cooldown_between_trials_secs: default_cooldown_secs(),
|
||||
min_recall: default_min_recall(),
|
||||
max_trials_per_hour: default_max_trials_per_hour(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn default_cycle_interval_secs() -> u64 { 60 }
|
||||
fn default_cooldown_secs() -> u64 { 30 }
|
||||
fn default_min_recall() -> f32 { 0.9 }
|
||||
fn default_max_trials_per_hour() -> u32 { 30 }
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct GatewayConfig {
|
||||
#[serde(default = "default_host")]
|
||||
@@ -28,8 +66,39 @@ pub struct GatewayConfig {
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct StorageConfig {
|
||||
/// Legacy single-backend root. If `buckets` is empty, this is used to
|
||||
/// create an implicit `primary` bucket at this path — preserves the
|
||||
/// pre-federation config shape.
|
||||
#[serde(default = "default_storage_root")]
|
||||
pub root: String,
|
||||
|
||||
/// Where profile buckets are rooted when auto-provisioned.
|
||||
#[serde(default = "default_profile_root")]
|
||||
pub profile_root: String,
|
||||
|
||||
/// Name of the bucket used for read fallback when a target bucket is
|
||||
/// unreachable. If `None`, no fallback — reads fail hard.
|
||||
#[serde(default)]
|
||||
pub rescue_bucket: Option<String>,
|
||||
|
||||
/// Explicitly configured buckets. Empty = backward-compat single-bucket
|
||||
/// mode driven by `root`.
|
||||
#[serde(default)]
|
||||
pub buckets: Vec<BucketConfig>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct BucketConfig {
|
||||
pub name: String,
|
||||
pub backend: String, // "local" | "s3"
|
||||
/// Local filesystem root (for backend = "local")
|
||||
pub root: Option<String>,
|
||||
/// S3 bucket name (for backend = "s3")
|
||||
pub bucket: Option<String>,
|
||||
pub region: Option<String>,
|
||||
pub endpoint: Option<String>,
|
||||
/// Handle for the secrets provider — never the literal credential.
|
||||
pub secret_ref: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Default)]
|
||||
@@ -78,6 +147,7 @@ pub struct ObservabilityConfig {
|
||||
fn default_host() -> String { "0.0.0.0".to_string() }
|
||||
fn default_gateway_port() -> u16 { 3100 }
|
||||
fn default_storage_root() -> String { "./data".to_string() }
|
||||
fn default_profile_root() -> String { "./data/_profiles".to_string() }
|
||||
fn default_manifest_prefix() -> String { "_catalog/manifests".to_string() }
|
||||
fn default_sidecar_url() -> String { "http://localhost:3200".to_string() }
|
||||
fn default_embed_model() -> String { "nomic-embed-text".to_string() }
|
||||
@@ -115,13 +185,19 @@ impl Default for Config {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
gateway: GatewayConfig { host: default_host(), port: default_gateway_port() },
|
||||
storage: StorageConfig { root: default_storage_root() },
|
||||
storage: StorageConfig {
|
||||
root: default_storage_root(),
|
||||
profile_root: default_profile_root(),
|
||||
rescue_bucket: None,
|
||||
buckets: Vec::new(),
|
||||
},
|
||||
catalog: CatalogConfig::default(),
|
||||
query: QueryConfig::default(),
|
||||
sidecar: SidecarConfig { url: default_sidecar_url() },
|
||||
ai: AiConfig::default(),
|
||||
auth: AuthConfig::default(),
|
||||
observability: ObservabilityConfig::default(),
|
||||
agent: AgentSettings::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,3 +3,4 @@ pub mod errors;
|
||||
pub mod arrow_helpers;
|
||||
pub mod config;
|
||||
pub mod pii;
|
||||
pub mod secrets;
|
||||
|
||||
@@ -0,0 +1,144 @@
|
||||
/// Secrets layer for bucket credentials.
|
||||
///
|
||||
/// Credentials never appear in `lakehouse.toml`. Each bucket entry has an
|
||||
/// opaque `secret_ref` handle, resolved at startup by a `SecretsProvider`.
|
||||
///
|
||||
/// MVP ships `FileSecretsProvider` — reads a TOML file at a path given by
|
||||
/// `LAKEHOUSE_SECRETS` env var, defaulting to `/etc/lakehouse/secrets.toml`.
|
||||
/// The file should be root-owned, mode 0600, never in git.
|
||||
///
|
||||
/// Future providers (Vault, SOPS, OS keyring) implement the same trait.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BucketCredentials {
|
||||
pub access_key: String,
|
||||
pub secret_key: String,
|
||||
}
|
||||
|
||||
// Redact credentials when debug-printing — NEVER let these end up in logs.
|
||||
impl BucketCredentials {
|
||||
pub fn redacted(&self) -> String {
|
||||
let prefix: String = self.access_key.chars().take(4).collect();
|
||||
format!("access_key={}... (redacted)", prefix)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait SecretsProvider: Send + Sync {
|
||||
/// Resolve a handle to credentials. Returns Err if the handle is unknown.
|
||||
async fn resolve(&self, handle: &str) -> Result<BucketCredentials, String>;
|
||||
|
||||
/// List known handles (for diagnostics only, never values).
|
||||
async fn list_handles(&self) -> Vec<String>;
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct RawSecretEntry {
|
||||
access_key: String,
|
||||
secret_key: String,
|
||||
}
|
||||
|
||||
/// File-backed secrets provider.
|
||||
pub struct FileSecretsProvider {
|
||||
path: PathBuf,
|
||||
cache: tokio::sync::RwLock<HashMap<String, BucketCredentials>>,
|
||||
}
|
||||
|
||||
impl FileSecretsProvider {
|
||||
pub fn new(path: impl Into<PathBuf>) -> Self {
|
||||
Self {
|
||||
path: path.into(),
|
||||
cache: tokio::sync::RwLock::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Default path: env var `LAKEHOUSE_SECRETS` or `/etc/lakehouse/secrets.toml`.
|
||||
pub fn from_env() -> Self {
|
||||
let path = std::env::var("LAKEHOUSE_SECRETS")
|
||||
.unwrap_or_else(|_| "/etc/lakehouse/secrets.toml".to_string());
|
||||
Self::new(path)
|
||||
}
|
||||
|
||||
/// Arc wrapper for convenient injection.
|
||||
pub fn shared(path: impl Into<PathBuf>) -> Arc<dyn SecretsProvider> {
|
||||
Arc::new(Self::new(path))
|
||||
}
|
||||
|
||||
/// Load and populate the cache. Call once at startup.
|
||||
pub async fn load(&self) -> Result<usize, String> {
|
||||
if !self.path.exists() {
|
||||
tracing::info!(
|
||||
"secrets file {} not present — running with zero credentials configured",
|
||||
self.path.display()
|
||||
);
|
||||
return Ok(0);
|
||||
}
|
||||
check_permissions(&self.path)?;
|
||||
let content = std::fs::read_to_string(&self.path)
|
||||
.map_err(|e| format!("read secrets file: {e}"))?;
|
||||
let raw: HashMap<String, RawSecretEntry> = toml::from_str(&content)
|
||||
.map_err(|e| format!("parse secrets file: {e}"))?;
|
||||
|
||||
let mut cache = self.cache.write().await;
|
||||
for (handle, entry) in raw {
|
||||
cache.insert(handle.clone(), BucketCredentials {
|
||||
access_key: entry.access_key,
|
||||
secret_key: entry.secret_key,
|
||||
});
|
||||
}
|
||||
let count = cache.len();
|
||||
tracing::info!("secrets: loaded {count} handles from {}", self.path.display());
|
||||
Ok(count)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SecretsProvider for FileSecretsProvider {
|
||||
async fn resolve(&self, handle: &str) -> Result<BucketCredentials, String> {
|
||||
let cache = self.cache.read().await;
|
||||
cache
|
||||
.get(handle)
|
||||
.cloned()
|
||||
.ok_or_else(|| format!("secret handle not found: {handle}"))
|
||||
}
|
||||
|
||||
async fn list_handles(&self) -> Vec<String> {
|
||||
self.cache.read().await.keys().cloned().collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Reject world-readable secrets files. Fail closed.
|
||||
fn check_permissions(path: &Path) -> Result<(), String> {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
let meta = std::fs::metadata(path)
|
||||
.map_err(|e| format!("stat secrets file: {e}"))?;
|
||||
let mode = meta.permissions().mode() & 0o777;
|
||||
if mode & 0o044 != 0 {
|
||||
return Err(format!(
|
||||
"secrets file {} is group/world-readable (mode {:o}); chmod 600 to proceed",
|
||||
path.display(), mode
|
||||
));
|
||||
}
|
||||
}
|
||||
let _ = path;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// A zero-handles provider for tests / purely-local setups.
|
||||
pub struct EmptySecretsProvider;
|
||||
|
||||
#[async_trait]
|
||||
impl SecretsProvider for EmptySecretsProvider {
|
||||
async fn resolve(&self, handle: &str) -> Result<BucketCredentials, String> {
|
||||
Err(format!("no secrets provider configured (wanted handle '{handle}')"))
|
||||
}
|
||||
async fn list_handles(&self) -> Vec<String> { Vec::new() }
|
||||
}
|
||||
@@ -116,4 +116,237 @@ pub struct DatasetManifest {
|
||||
/// Row count (updated on ingest/compact)
|
||||
#[serde(default)]
|
||||
pub row_count: Option<u64>,
|
||||
|
||||
// --- Embedding freshness (Phase C) ---
|
||||
|
||||
/// When the attached vector index was last refreshed. `None` means this
|
||||
/// dataset has no vector index yet, or has never been embedded.
|
||||
#[serde(default)]
|
||||
pub last_embedded_at: Option<chrono::DateTime<chrono::Utc>>,
|
||||
|
||||
/// When data was written that hasn't yet been embedded. `Some(t)` means
|
||||
/// the vector index is out-of-date as of timestamp `t`. Cleared on refresh.
|
||||
#[serde(default)]
|
||||
pub embedding_stale_since: Option<chrono::DateTime<chrono::Utc>>,
|
||||
|
||||
/// How this dataset wants stale embeddings handled.
|
||||
#[serde(default)]
|
||||
pub embedding_refresh_policy: Option<RefreshPolicy>,
|
||||
}
|
||||
|
||||
/// Controls what happens when new data is written to a dataset with an
|
||||
/// attached vector index.
|
||||
///
|
||||
/// - `Manual` (default): data writes set `embedding_stale_since`; nothing
|
||||
/// embeds until an operator or agent calls `/vectors/refresh/{dataset}`.
|
||||
/// - `OnAppend`: ingest fires a background refresh immediately after writing.
|
||||
/// Suitable for datasets where vector freshness matters more than ingest
|
||||
/// latency.
|
||||
/// - `Scheduled(cron)`: a timer or external scheduler triggers refresh at
|
||||
/// the named cadence. The scheduler itself is not in this ADR scope —
|
||||
/// this just declares the intent so operators can see which policy a
|
||||
/// dataset expects.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||
pub enum RefreshPolicy {
|
||||
Manual,
|
||||
OnAppend,
|
||||
Scheduled { cron: String },
|
||||
}
|
||||
|
||||
impl Default for RefreshPolicy {
|
||||
fn default() -> Self { Self::Manual }
|
||||
}
|
||||
|
||||
/// AI-safe view of a dataset (Phase D).
|
||||
///
|
||||
/// A view is a named projection over a base dataset that:
|
||||
/// - whitelists the columns visible (everything else is invisible),
|
||||
/// - optionally applies a row filter (a SQL WHERE clause fragment),
|
||||
/// - optionally redacts column values (e.g. mask the last 4 of an SSN).
|
||||
///
|
||||
/// Queries that reference the view name see only the safe projection.
|
||||
/// The base table remains separately accessible by its raw name — this
|
||||
/// is a *capability surface*, not a data classification fence. Use
|
||||
/// access control (Phase 13) on top to enforce who can query which view.
|
||||
///
|
||||
/// Stored as JSON at `_catalog/views/{name}.json` in the primary bucket
|
||||
/// alongside dataset manifests. The view itself doesn't store data —
|
||||
/// queryd materializes it as a DataFusion view at SessionContext build
|
||||
/// time.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AiView {
|
||||
pub name: String,
|
||||
/// Source dataset name. Must exist in the catalog.
|
||||
pub base_dataset: String,
|
||||
/// Whitelisted columns. Empty list is invalid (would expose nothing).
|
||||
pub columns: Vec<String>,
|
||||
/// Optional WHERE clause fragment (without the `WHERE` keyword), e.g.
|
||||
/// `status != 'blocked' AND city = 'Chicago'`. SQL injection here is
|
||||
/// the operator's problem — only trust this from authenticated admins.
|
||||
#[serde(default)]
|
||||
pub row_filter: Option<String>,
|
||||
/// Per-column redactions: `{column_name -> Redaction}`.
|
||||
#[serde(default)]
|
||||
pub column_redactions: std::collections::HashMap<String, Redaction>,
|
||||
pub created_at: chrono::DateTime<chrono::Utc>,
|
||||
#[serde(default)]
|
||||
pub created_by: String,
|
||||
#[serde(default)]
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
/// Per-model configuration for HNSW builds. Declared inline so the PRD
|
||||
/// invariant 9 ("every reader gets its own profile") can carry tuning
|
||||
/// preferences without forcing a round-trip to vectord's trial system.
|
||||
/// Mirrors `vectord::trial::HnswConfig` to avoid a cross-crate dep cycle.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct ProfileHnswConfig {
|
||||
pub ef_construction: usize,
|
||||
pub ef_search: usize,
|
||||
#[serde(default)]
|
||||
pub seed: Option<u64>,
|
||||
}
|
||||
|
||||
impl Default for ProfileHnswConfig {
|
||||
/// Matches the production HNSW default from Phase 15 trial sweep
|
||||
/// (see docs/ADR and vectord::trial::HnswConfig::default).
|
||||
fn default() -> Self {
|
||||
Self { ef_construction: 80, ef_search: 30, seed: None }
|
||||
}
|
||||
}
|
||||
|
||||
/// A named client of the substrate — a local LLM, an AI agent, or a
|
||||
/// human role that should see a specific scoped view of the data.
|
||||
/// Profiles are the first-class "every reader gets its own profile"
|
||||
/// concept from PRD invariant 9.
|
||||
///
|
||||
/// Bound datasets can be raw catalog tables OR `AiView` names — either
|
||||
/// works at query time since both register as DataFusion tables with
|
||||
/// the same name. Mixing raw and views lets you compose access rules:
|
||||
/// bind a model to `candidates_safe` (PII-redacted view) and `placements`
|
||||
/// (raw table, no PII), and the model gets exactly what it needs.
|
||||
///
|
||||
/// Stored at `_catalog/profiles/{id}.json`. Activation (`/profile/{id}/
|
||||
/// activate`) pre-loads the embedding cache and builds HNSW with the
|
||||
/// profile's config — first search after activate is warm.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ModelProfile {
|
||||
/// Unique id — slug form (alphanumeric + `-`, `_`). Used in URLs,
|
||||
/// bucket names (`profile:{id}`), and trial journal prefixes.
|
||||
pub id: String,
|
||||
/// Ollama model tag (e.g. "qwen2.5:7b", "nomic-embed-text:latest").
|
||||
/// Used when this profile is the "generate" or "embed" caller.
|
||||
pub ollama_name: String,
|
||||
/// Free-text description surfaced to operators. Who/what is this
|
||||
/// profile for?
|
||||
#[serde(default)]
|
||||
pub description: String,
|
||||
/// Names of datasets or AiViews this profile can read. A search
|
||||
/// against an index whose `source` isn't in this list is refused.
|
||||
pub bound_datasets: Vec<String>,
|
||||
/// Per-profile HNSW config, used when activate builds the graph.
|
||||
#[serde(default)]
|
||||
pub hnsw_config: ProfileHnswConfig,
|
||||
/// Which embedding model this profile expects. Defaults to
|
||||
/// `nomic-embed-text` (the only one currently installed).
|
||||
#[serde(default = "default_embed_model")]
|
||||
pub embed_model: String,
|
||||
pub created_at: chrono::DateTime<chrono::Utc>,
|
||||
#[serde(default)]
|
||||
pub created_by: String,
|
||||
/// Federation layer 2: which bucket this profile's artifacts
|
||||
/// (trial journals, promotions for its indexes) live in. If set and
|
||||
/// not yet registered, `POST /vectors/profile/{id}/activate` auto-
|
||||
/// provisions a local bucket under `storage.profile_root` and
|
||||
/// registers it before warming indexes. `None` = share the primary
|
||||
/// bucket (pre-federation behavior).
|
||||
#[serde(default)]
|
||||
pub bucket: Option<String>,
|
||||
/// ADR-019 hybrid: which storage backend serves this profile's
|
||||
/// indexes. Default Parquet so pre-existing profiles unchanged.
|
||||
/// New profiles can opt into Lance for append-heavy / random-access
|
||||
/// / 5M+ vector workloads.
|
||||
#[serde(default)]
|
||||
pub vector_backend: VectorBackend,
|
||||
}
|
||||
|
||||
fn default_embed_model() -> String { "nomic-embed-text".to_string() }
|
||||
|
||||
/// Per-profile / per-index storage backend choice.
|
||||
///
|
||||
/// `Parquet` (default) — vectors as binary-blob in Parquet + in-RAM HNSW
|
||||
/// for ANN. Best when:
|
||||
/// - Index ≤ ~5M vectors (fits in RAM)
|
||||
/// - Search latency dominates (HNSW p50 < 1ms at 100K)
|
||||
/// - Workload is mostly read, not append
|
||||
///
|
||||
/// `Lance` — vectors as native FixedSizeList in a Lance dataset, IVF_PQ
|
||||
/// on-disk index. Best when:
|
||||
/// - Index will exceed RAM ceiling (5M+ vectors per index)
|
||||
/// - Heavy random-row access for RAG (Lance: 311us; Parquet: ~35ms full
|
||||
/// scan per fetch)
|
||||
/// - Append-heavy ingest (Lance: 0.08s for 10K rows; Parquet: full
|
||||
/// ~330MB rewrite per change)
|
||||
/// - Need to rebuild index frequently (Lance IVF_PQ: 16s; HNSW: 230s
|
||||
/// at the same scale)
|
||||
///
|
||||
/// Defaults to Parquet for backward compat with pre-ADR-019 indexes.
|
||||
/// See `docs/ADR-019-vector-storage.md` for the full scorecard.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum VectorBackend {
|
||||
Parquet,
|
||||
Lance,
|
||||
}
|
||||
|
||||
impl Default for VectorBackend {
|
||||
fn default() -> Self { VectorBackend::Parquet }
|
||||
}
|
||||
|
||||
/// Soft-delete marker (Phase E).
|
||||
///
|
||||
/// Tombstones live beside the dataset in `_catalog/tombstones/{dataset}/`
|
||||
/// as append-log JSONL. Query-time filter in queryd reads all tombstones
|
||||
/// for each dataset and wraps the base table in a DataFusion view that
|
||||
/// excludes tombstoned rows. Physical deletion happens later (compaction),
|
||||
/// so the row count immediately reflects the delete but data is still on
|
||||
/// disk until compact runs. That's deliberate — it gives a reversal
|
||||
/// window and keeps the event journal audit trail intact.
|
||||
///
|
||||
/// All tombstones for a given dataset must use the same `row_key_column`
|
||||
/// (enforced on write); otherwise the query-time filter can't be built
|
||||
/// as a single WHERE clause.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Tombstone {
|
||||
pub dataset: String,
|
||||
/// Column name that identifies the row (e.g. "candidate_id").
|
||||
pub row_key_column: String,
|
||||
/// Value of that column for the tombstoned row.
|
||||
pub row_key_value: String,
|
||||
pub deleted_at: chrono::DateTime<chrono::Utc>,
|
||||
/// Human / system actor responsible for the delete (audit).
|
||||
#[serde(default)]
|
||||
pub actor: String,
|
||||
/// Why (e.g. "GDPR request #1234", "user-requested erasure").
|
||||
#[serde(default)]
|
||||
pub reason: String,
|
||||
}
|
||||
|
||||
/// How a column's values should be transformed before being returned.
|
||||
/// `Mask` is the most common — keeps a few visible chars, replaces the
|
||||
/// rest with `*`. `Hash` returns SHA-256 of the value for join keys you
|
||||
/// want to preserve referential integrity on without exposing the value.
|
||||
/// `Null` returns NULL unconditionally.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||
pub enum Redaction {
|
||||
Null,
|
||||
Hash,
|
||||
Mask {
|
||||
/// Number of leading characters left visible.
|
||||
keep_prefix: usize,
|
||||
/// Number of trailing characters left visible.
|
||||
keep_suffix: usize,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -13,3 +13,4 @@ tracing = { workspace = true }
|
||||
object_store = { workspace = true }
|
||||
bytes = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
|
||||
@@ -0,0 +1,239 @@
|
||||
/// Write-once batched append log.
|
||||
///
|
||||
/// Problem we're fixing: the error journal and HNSW trial journal both
|
||||
/// previously did read-modify-write of their whole JSONL file on every
|
||||
/// event. That's O(N²) cumulative work and generates huge churn at scale.
|
||||
/// It's exactly the pattern llms3.com flags as the "small-file /
|
||||
/// rewrite-amplification" anti-pattern.
|
||||
///
|
||||
/// This helper implements the pattern object storage actually wants:
|
||||
///
|
||||
/// - Events accumulate in an **in-memory buffer** (reads see them immediately).
|
||||
/// - When the buffer hits a threshold, or `flush()` is called, the buffer is
|
||||
/// written **as one new object** with a timestamp-sorted key.
|
||||
/// - Existing objects are never rewritten.
|
||||
/// - Reads enumerate all batch files, sort by key, and concat in order.
|
||||
/// - An explicit `compact()` reads every batch file, writes one consolidated
|
||||
/// file, and deletes the originals — the LSM-tree compaction idea applied
|
||||
/// to small JSONL events.
|
||||
///
|
||||
/// Storage layout:
|
||||
/// ```
|
||||
/// {prefix}/
|
||||
/// batch_0001776319628000123.jsonl
|
||||
/// batch_0001776319745987654.jsonl
|
||||
/// batch_compacted_00001776319800000000.jsonl (after compact())
|
||||
/// ```
|
||||
/// Key format: the zero-padded epoch microsecond of the write, so lexical
|
||||
/// sort == chronological sort.
|
||||
|
||||
use bytes::Bytes;
|
||||
use chrono::Utc;
|
||||
use object_store::ObjectStore;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::ops;
|
||||
|
||||
const DEFAULT_FLUSH_THRESHOLD: usize = 32;
|
||||
|
||||
pub struct AppendLog {
|
||||
store: Arc<dyn ObjectStore>,
|
||||
prefix: String,
|
||||
buffer: Mutex<Vec<Vec<u8>>>,
|
||||
flush_threshold: usize,
|
||||
}
|
||||
|
||||
impl AppendLog {
|
||||
/// Create a new append log rooted at `prefix` in the given object store.
|
||||
/// Events auto-flush when buffer reaches `flush_threshold` (default 32).
|
||||
pub fn new(store: Arc<dyn ObjectStore>, prefix: impl Into<String>) -> Self {
|
||||
Self {
|
||||
store,
|
||||
prefix: prefix.into(),
|
||||
buffer: Mutex::new(Vec::new()),
|
||||
flush_threshold: DEFAULT_FLUSH_THRESHOLD,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_flush_threshold(mut self, threshold: usize) -> Self {
|
||||
self.flush_threshold = threshold.max(1);
|
||||
self
|
||||
}
|
||||
|
||||
/// Add an event. The returned future completes either immediately
|
||||
/// (buffered) or after a flush, depending on whether the buffer hit the
|
||||
/// threshold. Callers don't need to care either way.
|
||||
pub async fn append(&self, line: Vec<u8>) -> Result<(), String> {
|
||||
let should_flush = {
|
||||
let mut buf = self.buffer.lock().await;
|
||||
buf.push(line);
|
||||
buf.len() >= self.flush_threshold
|
||||
};
|
||||
if should_flush {
|
||||
self.flush().await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Force-flush the in-memory buffer to object storage as a single new
|
||||
/// batch file. Safe to call anytime (idempotent no-op when buffer empty).
|
||||
pub async fn flush(&self) -> Result<(), String> {
|
||||
let batch = {
|
||||
let mut buf = self.buffer.lock().await;
|
||||
if buf.is_empty() { return Ok(()); }
|
||||
std::mem::take(&mut *buf)
|
||||
};
|
||||
|
||||
let ts_us = Utc::now().timestamp_micros().max(0) as u128;
|
||||
let key = format!("{}/batch_{:019}.jsonl", self.prefix, ts_us);
|
||||
|
||||
let mut body = Vec::with_capacity(batch.iter().map(|b| b.len() + 1).sum());
|
||||
for line in batch {
|
||||
body.extend_from_slice(&line);
|
||||
if !body.ends_with(b"\n") { body.push(b'\n'); }
|
||||
}
|
||||
ops::put(&self.store, &key, Bytes::from(body)).await
|
||||
}
|
||||
|
||||
/// Read every event across all batch files + unflushed in-memory buffer.
|
||||
/// Events are returned in chronological order.
|
||||
pub async fn read_all(&self) -> Result<Vec<Vec<u8>>, String> {
|
||||
let mut keys = self.list_batch_keys().await?;
|
||||
keys.sort();
|
||||
|
||||
let mut out = Vec::new();
|
||||
for key in keys {
|
||||
let bytes = ops::get(&self.store, &key).await?;
|
||||
for line in bytes.split(|b| *b == b'\n') {
|
||||
if !line.is_empty() {
|
||||
out.push(line.to_vec());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Include unflushed events so callers see the latest state
|
||||
// whether or not someone ran flush() recently.
|
||||
let buf = self.buffer.lock().await;
|
||||
for line in buf.iter() {
|
||||
out.push(line.clone());
|
||||
}
|
||||
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Consolidate all current batch files into one compacted file, then
|
||||
/// delete the originals. Safe to call while appends are in flight:
|
||||
/// new batches written during compaction get a higher timestamp and
|
||||
/// survive. Fails closed — if anything goes wrong mid-delete, the
|
||||
/// compacted file coexists with originals and next read sees duplicates
|
||||
/// (which the dedup caller must handle) rather than data loss.
|
||||
pub async fn compact(&self) -> Result<CompactStats, String> {
|
||||
// Snapshot which files to compact BEFORE we write the new one.
|
||||
let mut originals = self.list_batch_keys().await?;
|
||||
originals.sort();
|
||||
|
||||
if originals.len() < 2 {
|
||||
return Ok(CompactStats { merged_files: originals.len(), events: 0, new_key: None });
|
||||
}
|
||||
|
||||
// Gather all existing events.
|
||||
let mut events = Vec::new();
|
||||
for key in &originals {
|
||||
let bytes = ops::get(&self.store, key).await?;
|
||||
for line in bytes.split(|b| *b == b'\n') {
|
||||
if !line.is_empty() {
|
||||
events.push(line.to_vec());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let total_events = events.len();
|
||||
if total_events == 0 {
|
||||
// Clean up empty files without writing a new one.
|
||||
for key in &originals {
|
||||
let _ = ops::delete(&self.store, key).await;
|
||||
}
|
||||
return Ok(CompactStats { merged_files: originals.len(), events: 0, new_key: None });
|
||||
}
|
||||
|
||||
// Name the compacted file with the SAME `batch_{ts}` format so it
|
||||
// sorts chronologically with future batches. Using a distinct prefix
|
||||
// ("batch_compacted_") would break lex ordering: later `batch_N`
|
||||
// files would sort BEFORE the compacted file because 'c' > digits.
|
||||
// Timestamp = now, so any appends arriving during compaction (which
|
||||
// get the current wall-clock time) sort AFTER this file.
|
||||
let ts_us = Utc::now().timestamp_micros().max(0) as u128;
|
||||
let new_key = format!("{}/batch_{:019}.jsonl", self.prefix, ts_us);
|
||||
|
||||
let mut body = Vec::new();
|
||||
for line in &events {
|
||||
body.extend_from_slice(line);
|
||||
body.push(b'\n');
|
||||
}
|
||||
ops::put(&self.store, &new_key, Bytes::from(body)).await?;
|
||||
|
||||
// Only delete originals once the consolidated file is persisted.
|
||||
let mut failures = 0;
|
||||
for key in &originals {
|
||||
if ops::delete(&self.store, key).await.is_err() {
|
||||
failures += 1;
|
||||
}
|
||||
}
|
||||
if failures > 0 {
|
||||
tracing::warn!(
|
||||
"compact '{}': {} original files failed to delete — consolidated file {} has the data",
|
||||
self.prefix, failures, new_key,
|
||||
);
|
||||
}
|
||||
|
||||
Ok(CompactStats {
|
||||
merged_files: originals.len(),
|
||||
events: total_events,
|
||||
new_key: Some(new_key),
|
||||
})
|
||||
}
|
||||
|
||||
/// How many batch files exist on disk right now.
|
||||
pub async fn file_count(&self) -> Result<usize, String> {
|
||||
Ok(self.list_batch_keys().await?.len())
|
||||
}
|
||||
|
||||
async fn list_batch_keys(&self) -> Result<Vec<String>, String> {
|
||||
let prefix_with_slash = format!("{}/", self.prefix);
|
||||
// list the prefix then filter for keys that match our naming scheme;
|
||||
// unrelated files at the same prefix won't be touched.
|
||||
let raw = ops::list(&self.store, Some(&prefix_with_slash)).await?;
|
||||
Ok(raw
|
||||
.into_iter()
|
||||
.filter(|k| {
|
||||
let basename = k.rsplit('/').next().unwrap_or(k);
|
||||
basename.starts_with("batch_") && basename.ends_with(".jsonl")
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct CompactStats {
|
||||
pub merged_files: usize,
|
||||
pub events: usize,
|
||||
pub new_key: Option<String>,
|
||||
}
|
||||
|
||||
// Proactively flush on drop, best-effort.
|
||||
// We can't `.await` in Drop; we spawn the flush on the tokio runtime if one
|
||||
// is available. If the runtime is already shutting down the flush is lost —
|
||||
// acceptable for these observability journals, which are hints not records.
|
||||
impl Drop for AppendLog {
|
||||
fn drop(&mut self) {
|
||||
let buf_len = self.buffer.try_lock().map(|b| b.len()).unwrap_or(0);
|
||||
if buf_len == 0 { return; }
|
||||
// Can't spawn from sync Drop on every runtime shape; log + move on.
|
||||
tracing::debug!(
|
||||
"append_log '{}' dropping with {} unflushed events",
|
||||
self.prefix, buf_len,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,200 @@
|
||||
/// Bucket operation error journal.
|
||||
///
|
||||
/// Every bucket op failure and every rescue fallback lands here. Goal:
|
||||
/// answering "is anything broken?" with one HTTP call.
|
||||
///
|
||||
/// Storage: batched write-once files under `_errors/bucket_errors/` in the
|
||||
/// primary bucket. Uses the shared `AppendLog` helper so we never rewrite
|
||||
/// existing files — see `append_log.rs` for the full pattern rationale.
|
||||
///
|
||||
/// In-memory ring buffer holds the last N events for fast response; on
|
||||
/// startup `load_recent` hydrates it from all batch files.
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use object_store::ObjectStore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::append_log::{AppendLog, CompactStats};
|
||||
|
||||
const JOURNAL_PREFIX: &str = "_errors/bucket_errors";
|
||||
const RING_CAPACITY: usize = 2000;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum BucketOp {
|
||||
Read,
|
||||
Write,
|
||||
Delete,
|
||||
List,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct BucketErrorEvent {
|
||||
pub ts: DateTime<Utc>,
|
||||
pub op: BucketOp,
|
||||
pub target: String,
|
||||
pub key: String,
|
||||
pub error: String,
|
||||
#[serde(default)]
|
||||
pub rescued: bool,
|
||||
}
|
||||
|
||||
impl BucketErrorEvent {
|
||||
pub fn new_read(target: &str, key: &str, error: &str) -> Self {
|
||||
Self { ts: Utc::now(), op: BucketOp::Read, target: target.into(), key: key.into(), error: error.into(), rescued: false }
|
||||
}
|
||||
pub fn new_write(target: &str, key: &str, error: &str) -> Self {
|
||||
Self { ts: Utc::now(), op: BucketOp::Write, target: target.into(), key: key.into(), error: error.into(), rescued: false }
|
||||
}
|
||||
pub fn new_delete(target: &str, key: &str, error: &str) -> Self {
|
||||
Self { ts: Utc::now(), op: BucketOp::Delete, target: target.into(), key: key.into(), error: error.into(), rescued: false }
|
||||
}
|
||||
pub fn new_list(target: &str, prefix: &str, error: &str) -> Self {
|
||||
Self { ts: Utc::now(), op: BucketOp::List, target: target.into(), key: prefix.into(), error: error.into(), rescued: false }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ErrorJournal {
|
||||
log: Arc<AppendLog>,
|
||||
ring: Arc<RwLock<VecDeque<BucketErrorEvent>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct HealthReport {
|
||||
pub period_minutes: i64,
|
||||
pub total_errors: usize,
|
||||
pub per_bucket: std::collections::HashMap<String, usize>,
|
||||
pub unhealthy_buckets: Vec<String>,
|
||||
}
|
||||
|
||||
impl ErrorJournal {
|
||||
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||
// Keep flush threshold lowish — operators checking /storage/errors
|
||||
// after a recent incident want to see fresh rows even without an
|
||||
// explicit flush.
|
||||
let log = Arc::new(
|
||||
AppendLog::new(store, JOURNAL_PREFIX).with_flush_threshold(8),
|
||||
);
|
||||
Self {
|
||||
log,
|
||||
ring: Arc::new(RwLock::new(VecDeque::with_capacity(RING_CAPACITY))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Hydrate the ring buffer from existing batch files. Call once at
|
||||
/// startup. Tolerates malformed lines (skipped with a warning) and
|
||||
/// missing files (returns 0).
|
||||
pub async fn load_recent(&self) -> Result<usize, String> {
|
||||
let lines = self.log.read_all().await.unwrap_or_default();
|
||||
let mut ring = self.ring.write().await;
|
||||
for line in lines {
|
||||
match serde_json::from_slice::<BucketErrorEvent>(&line) {
|
||||
Ok(ev) => {
|
||||
if ring.len() >= RING_CAPACITY { ring.pop_front(); }
|
||||
ring.push_back(ev);
|
||||
}
|
||||
Err(e) => tracing::warn!("error journal: skip malformed line ({e})"),
|
||||
}
|
||||
}
|
||||
Ok(ring.len())
|
||||
}
|
||||
|
||||
/// Append an event. In-memory ring updated immediately; persistence
|
||||
/// happens in batches via the underlying AppendLog.
|
||||
pub async fn append(&self, event: BucketErrorEvent) {
|
||||
{
|
||||
let mut ring = self.ring.write().await;
|
||||
if ring.len() >= RING_CAPACITY { ring.pop_front(); }
|
||||
ring.push_back(event.clone());
|
||||
}
|
||||
|
||||
match serde_json::to_vec(&event) {
|
||||
Ok(line) => {
|
||||
if let Err(e) = self.log.append(line).await {
|
||||
tracing::error!("error journal persist failed: {e}");
|
||||
}
|
||||
}
|
||||
Err(e) => tracing::error!("error journal serialize failed: {e}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Mark the most recent matching in-memory event as rescued.
|
||||
/// Only updates the ring buffer — the JSONL line already persisted
|
||||
/// records the failure fact; rescue status travels in response headers.
|
||||
pub async fn mark_rescued_last(&self, target: &str, key: &str) {
|
||||
let mut ring = self.ring.write().await;
|
||||
for ev in ring.iter_mut().rev() {
|
||||
if matches!(ev.op, BucketOp::Read) && ev.target == target && ev.key == key && !ev.rescued {
|
||||
ev.rescued = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn recent(&self, limit: usize) -> Vec<BucketErrorEvent> {
|
||||
let ring = self.ring.read().await;
|
||||
let start = ring.len().saturating_sub(limit);
|
||||
ring.iter().skip(start).cloned().collect()
|
||||
}
|
||||
|
||||
pub async fn filter(
|
||||
&self,
|
||||
bucket: Option<&str>,
|
||||
since: Option<DateTime<Utc>>,
|
||||
limit: usize,
|
||||
) -> Vec<BucketErrorEvent> {
|
||||
let ring = self.ring.read().await;
|
||||
ring.iter()
|
||||
.rev()
|
||||
.filter(|ev| bucket.map_or(true, |b| ev.target == b))
|
||||
.filter(|ev| since.map_or(true, |s| ev.ts >= s))
|
||||
.take(limit)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>()
|
||||
.into_iter()
|
||||
.rev()
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Summarize errors in the last N minutes.
|
||||
pub async fn health(&self, period_minutes: i64) -> HealthReport {
|
||||
use std::collections::HashMap;
|
||||
let cutoff = Utc::now() - chrono::Duration::minutes(period_minutes);
|
||||
let ring = self.ring.read().await;
|
||||
let recent: Vec<_> = ring.iter().filter(|ev| ev.ts >= cutoff).collect();
|
||||
let mut per_bucket: HashMap<String, usize> = HashMap::new();
|
||||
for ev in &recent {
|
||||
*per_bucket.entry(ev.target.clone()).or_insert(0) += 1;
|
||||
}
|
||||
let unhealthy_buckets: Vec<String> = per_bucket
|
||||
.iter()
|
||||
.filter(|(_, c)| **c >= 3)
|
||||
.map(|(k, _)| k.clone())
|
||||
.collect();
|
||||
HealthReport {
|
||||
period_minutes,
|
||||
total_errors: recent.len(),
|
||||
per_bucket,
|
||||
unhealthy_buckets,
|
||||
}
|
||||
}
|
||||
|
||||
/// Force an immediate flush of buffered events to object storage.
|
||||
pub async fn flush(&self) -> Result<(), String> {
|
||||
self.log.flush().await
|
||||
}
|
||||
|
||||
/// Consolidate all batch files into one. Operator cleanup.
|
||||
pub async fn compact(&self) -> Result<CompactStats, String> {
|
||||
self.log.compact().await
|
||||
}
|
||||
|
||||
/// How many JSONL batch files currently exist.
|
||||
pub async fn file_count(&self) -> Result<usize, String> {
|
||||
self.log.file_count().await
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,179 @@
|
||||
/// HTTP surface for federation operator endpoints.
|
||||
///
|
||||
/// - `GET /buckets` — configured buckets with reachability status
|
||||
/// - `GET /errors` — recent bucket op failures (filterable)
|
||||
/// - `GET /bucket-health` — aggregated errors-per-bucket in the last 5 minutes
|
||||
///
|
||||
/// Mounted under `/storage` by the gateway. `/storage/health` is already
|
||||
/// claimed by the existing storage router for service liveness, so the
|
||||
/// aggregated bucket health lives at `/storage/bucket-health`.
|
||||
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::{Path, Query, State},
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
routing::{delete, get, post, put},
|
||||
};
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::Deserialize;
|
||||
use shared::config::BucketConfig;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::error_journal::BucketErrorEvent;
|
||||
use crate::registry::{BucketInfo, BucketRegistry};
|
||||
|
||||
pub fn router(registry: Arc<BucketRegistry>) -> Router {
|
||||
Router::new()
|
||||
.route("/buckets", get(list_buckets).post(add_bucket))
|
||||
.route("/buckets/{name}", delete(remove_bucket))
|
||||
.route("/errors", get(list_errors))
|
||||
.route("/errors/flush", post(flush_errors))
|
||||
.route("/errors/compact", post(compact_errors))
|
||||
.route("/bucket-health", get(get_health))
|
||||
.route("/buckets/{bucket}/objects/{*key}", put(put_bucket_object))
|
||||
.route("/buckets/{bucket}/objects/{*key}", get(get_bucket_object))
|
||||
.with_state(registry)
|
||||
}
|
||||
|
||||
/// Provision + register a bucket at runtime. Body is a `BucketConfig`.
|
||||
///
|
||||
/// Federation layer 2: profile buckets (`profile:alice`) and tenant
|
||||
/// buckets can be added after startup without a service restart.
|
||||
async fn add_bucket(
|
||||
State(reg): State<Arc<BucketRegistry>>,
|
||||
Json(bc): Json<BucketConfig>,
|
||||
) -> impl IntoResponse {
|
||||
// Safety net: local backends must land somewhere under the configured
|
||||
// profile_root (for profile:*) or be otherwise explicitly rooted.
|
||||
// Refuse empty / missing root to avoid "bucket created at ./" surprises.
|
||||
if bc.backend == "local" {
|
||||
match bc.root.as_deref() {
|
||||
Some(r) if !r.trim().is_empty() => {}
|
||||
_ => return Err((StatusCode::BAD_REQUEST,
|
||||
"local bucket requires a non-empty 'root' path".to_string())),
|
||||
}
|
||||
}
|
||||
match reg.add_bucket(bc).await {
|
||||
Ok(info) => Ok((StatusCode::CREATED, Json(info))),
|
||||
Err(e) => {
|
||||
let code = if e.contains("already registered") {
|
||||
StatusCode::CONFLICT
|
||||
} else {
|
||||
StatusCode::BAD_REQUEST
|
||||
};
|
||||
Err((code, e))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Unregister a bucket. Refused for primary / rescue / unknown names.
|
||||
async fn remove_bucket(
|
||||
State(reg): State<Arc<BucketRegistry>>,
|
||||
Path(name): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
match reg.remove_bucket(&name) {
|
||||
Ok(()) => Ok((StatusCode::OK, format!("unregistered: {name}"))),
|
||||
Err(e) => {
|
||||
let code = if e.contains("cannot remove") {
|
||||
StatusCode::FORBIDDEN
|
||||
} else if e.contains("not registered") {
|
||||
StatusCode::NOT_FOUND
|
||||
} else {
|
||||
StatusCode::BAD_REQUEST
|
||||
};
|
||||
Err((code, e))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_buckets(State(reg): State<Arc<BucketRegistry>>) -> Json<Vec<BucketInfo>> {
|
||||
Json(reg.list().await)
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ErrorQuery {
|
||||
bucket: Option<String>,
|
||||
since: Option<DateTime<Utc>>,
|
||||
#[serde(default = "default_limit")]
|
||||
limit: usize,
|
||||
}
|
||||
|
||||
fn default_limit() -> usize { 50 }
|
||||
|
||||
async fn list_errors(
|
||||
State(reg): State<Arc<BucketRegistry>>,
|
||||
Query(q): Query<ErrorQuery>,
|
||||
) -> Json<Vec<BucketErrorEvent>> {
|
||||
Json(reg.journal().filter(q.bucket.as_deref(), q.since, q.limit).await)
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct HealthQuery {
|
||||
#[serde(default = "default_period")]
|
||||
minutes: i64,
|
||||
}
|
||||
|
||||
fn default_period() -> i64 { 5 }
|
||||
|
||||
async fn get_health(
|
||||
State(reg): State<Arc<BucketRegistry>>,
|
||||
Query(q): Query<HealthQuery>,
|
||||
) -> impl IntoResponse {
|
||||
Json(reg.journal().health(q.minutes).await)
|
||||
}
|
||||
|
||||
/// Bucket-aware write. Hard-fails if the target bucket is unreachable;
|
||||
/// never falls back to rescue for writes.
|
||||
async fn put_bucket_object(
|
||||
State(reg): State<Arc<BucketRegistry>>,
|
||||
Path((bucket, key)): Path<(String, String)>,
|
||||
body: bytes::Bytes,
|
||||
) -> impl IntoResponse {
|
||||
match reg.write_smart(&bucket, &key, body).await {
|
||||
Ok(()) => (StatusCode::CREATED, format!("stored: {bucket}/{key}")).into_response(),
|
||||
Err(e) => (StatusCode::SERVICE_UNAVAILABLE, e).into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Bucket-aware read. Falls through to the rescue bucket if the target
|
||||
/// is unreachable or the key is missing; emits observability headers so
|
||||
/// callers can detect the fallback.
|
||||
async fn flush_errors(State(reg): State<Arc<BucketRegistry>>) -> impl IntoResponse {
|
||||
match reg.journal().flush().await {
|
||||
Ok(()) => (StatusCode::OK, "flushed").into_response(),
|
||||
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e).into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn compact_errors(State(reg): State<Arc<BucketRegistry>>) -> impl IntoResponse {
|
||||
match reg.journal().compact().await {
|
||||
Ok(stats) => Json(stats).into_response(),
|
||||
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e).into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_bucket_object(
|
||||
State(reg): State<Arc<BucketRegistry>>,
|
||||
Path((bucket, key)): Path<(String, String)>,
|
||||
) -> impl IntoResponse {
|
||||
match reg.read_smart(&bucket, &key).await {
|
||||
Ok(outcome) => {
|
||||
let mut resp = outcome.data.into_response();
|
||||
if outcome.rescued {
|
||||
let headers = resp.headers_mut();
|
||||
headers.insert("x-lakehouse-rescue-used", "true".parse().unwrap());
|
||||
headers.insert(
|
||||
"x-lakehouse-original-bucket",
|
||||
outcome.original_bucket.parse().unwrap(),
|
||||
);
|
||||
headers.insert(
|
||||
"x-lakehouse-served-by",
|
||||
outcome.served_by.parse().unwrap(),
|
||||
);
|
||||
}
|
||||
resp
|
||||
}
|
||||
Err(e) => (StatusCode::NOT_FOUND, e).into_response(),
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,7 @@
|
||||
pub mod append_log;
|
||||
pub mod backend;
|
||||
pub mod error_journal;
|
||||
pub mod federation_service;
|
||||
pub mod ops;
|
||||
pub mod registry;
|
||||
pub mod service;
|
||||
|
||||
@@ -0,0 +1,377 @@
|
||||
/// Multi-backend bucket registry — the federation foundation.
|
||||
///
|
||||
/// Federation rule: every `ObjectRef` belongs to exactly one named bucket.
|
||||
/// The registry resolves bucket names to `object_store` backends, handles
|
||||
/// rescue-bucket fallback on read failure, writes every failure to the
|
||||
/// error journal, and exposes a health summary for operators.
|
||||
///
|
||||
/// Existing call sites can keep using `ops::*` with `registry.get(name)`.
|
||||
/// New bucket-aware call sites use `registry.read_smart` / `write_smart`
|
||||
/// which handle fallback + journaling automatically.
|
||||
|
||||
use object_store::ObjectStore;
|
||||
use object_store::local::LocalFileSystem;
|
||||
use serde::Serialize;
|
||||
use shared::config::{BucketConfig, StorageConfig};
|
||||
use shared::secrets::{BucketCredentials, SecretsProvider};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, RwLock};
|
||||
|
||||
use crate::error_journal::{BucketErrorEvent, ErrorJournal};
|
||||
|
||||
/// A registered bucket — the store handle + its configuration.
|
||||
pub struct BucketEntry {
|
||||
pub name: String,
|
||||
pub backend: String,
|
||||
pub store: Arc<dyn ObjectStore>,
|
||||
pub config: BucketConfig,
|
||||
}
|
||||
|
||||
/// Read outcome — may have been rescued.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ReadOutcome {
|
||||
pub data: bytes::Bytes,
|
||||
pub rescued: bool,
|
||||
pub original_bucket: String,
|
||||
pub served_by: String,
|
||||
}
|
||||
|
||||
/// Summary entry for GET /storage/buckets.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct BucketInfo {
|
||||
pub name: String,
|
||||
pub backend: String,
|
||||
pub reachable: bool,
|
||||
pub role: BucketRole,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, PartialEq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum BucketRole {
|
||||
Primary,
|
||||
Rescue,
|
||||
Profile,
|
||||
Tenant,
|
||||
}
|
||||
|
||||
pub struct BucketRegistry {
|
||||
/// RwLock because federation layer 2 adds runtime bucket lifecycle
|
||||
/// (`POST /storage/buckets` / `DELETE /storage/buckets/{name}`).
|
||||
/// Almost all accesses are reads — writes happen at provision /
|
||||
/// deactivate time only.
|
||||
buckets: RwLock<HashMap<String, Arc<BucketEntry>>>,
|
||||
default: String,
|
||||
rescue: Option<String>,
|
||||
profile_root: String,
|
||||
/// Held so runtime `add_bucket` can resolve secret_ref handles.
|
||||
secrets: Arc<dyn SecretsProvider>,
|
||||
journal: ErrorJournal,
|
||||
}
|
||||
|
||||
impl BucketRegistry {
|
||||
/// Build the registry from storage config + secrets provider.
|
||||
/// Back-compat: if `buckets` is empty, synthesize a `primary` bucket from
|
||||
/// the legacy `root` field so pre-federation configs keep working.
|
||||
pub async fn from_config(
|
||||
cfg: &StorageConfig,
|
||||
secrets: Arc<dyn SecretsProvider>,
|
||||
) -> Result<Self, String> {
|
||||
let mut buckets: HashMap<String, Arc<BucketEntry>> = HashMap::new();
|
||||
|
||||
let bucket_configs: Vec<BucketConfig> = if cfg.buckets.is_empty() {
|
||||
vec![BucketConfig {
|
||||
name: "primary".to_string(),
|
||||
backend: "local".to_string(),
|
||||
root: Some(cfg.root.clone()),
|
||||
bucket: None,
|
||||
region: None,
|
||||
endpoint: None,
|
||||
secret_ref: None,
|
||||
}]
|
||||
} else {
|
||||
cfg.buckets.clone()
|
||||
};
|
||||
|
||||
for bc in bucket_configs {
|
||||
let store = build_store(&bc, secrets.as_ref()).await?;
|
||||
let entry = Arc::new(BucketEntry {
|
||||
name: bc.name.clone(),
|
||||
backend: bc.backend.clone(),
|
||||
store,
|
||||
config: bc.clone(),
|
||||
});
|
||||
buckets.insert(bc.name.clone(), entry);
|
||||
}
|
||||
|
||||
// Ensure `primary` always exists — it's where error journals live.
|
||||
if !buckets.contains_key("primary") {
|
||||
return Err("no bucket named 'primary' configured — required as error-journal home".into());
|
||||
}
|
||||
|
||||
// Rescue bucket is optional but, if named, must exist.
|
||||
if let Some(r) = &cfg.rescue_bucket {
|
||||
if !buckets.contains_key(r) {
|
||||
return Err(format!("rescue_bucket '{r}' not found among configured buckets"));
|
||||
}
|
||||
}
|
||||
|
||||
let journal = ErrorJournal::new(buckets.get("primary").unwrap().store.clone());
|
||||
let _ = journal.load_recent().await;
|
||||
|
||||
Ok(Self {
|
||||
buckets: RwLock::new(buckets),
|
||||
default: "primary".to_string(),
|
||||
rescue: cfg.rescue_bucket.clone(),
|
||||
profile_root: cfg.profile_root.clone(),
|
||||
secrets,
|
||||
journal,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn default_name(&self) -> &str { &self.default }
|
||||
|
||||
pub fn rescue_name(&self) -> Option<&str> { self.rescue.as_deref() }
|
||||
|
||||
pub fn profile_root(&self) -> &str { &self.profile_root }
|
||||
|
||||
pub fn journal(&self) -> &ErrorJournal { &self.journal }
|
||||
|
||||
/// Resolve a bucket name to its object store. Existing call sites use
|
||||
/// this as a drop-in replacement for the old single-store pattern.
|
||||
///
|
||||
/// Uses `std::sync::RwLock` — every caller holds the guard for
|
||||
/// microseconds (clones an Arc and drops), so there's no async-
|
||||
/// blocking concern. Never hold across an await.
|
||||
pub fn get(&self, bucket: &str) -> Result<Arc<dyn ObjectStore>, String> {
|
||||
let guard = self.buckets.read().expect("bucket registry lock poisoned");
|
||||
guard
|
||||
.get(bucket)
|
||||
.map(|e| e.store.clone())
|
||||
.ok_or_else(|| format!("unknown bucket: {bucket}"))
|
||||
}
|
||||
|
||||
/// The default bucket's store — use for code paths that don't yet know
|
||||
/// about buckets.
|
||||
pub fn default_store(&self) -> Arc<dyn ObjectStore> {
|
||||
let guard = self.buckets.read().expect("bucket registry lock poisoned");
|
||||
guard.get(&self.default).unwrap().store.clone()
|
||||
}
|
||||
|
||||
/// True if this bucket name is registered.
|
||||
pub fn contains(&self, bucket: &str) -> bool {
|
||||
let guard = self.buckets.read().expect("bucket registry lock poisoned");
|
||||
guard.contains_key(bucket)
|
||||
}
|
||||
|
||||
/// List all registered buckets. Checks reachability by doing a trivial
|
||||
/// `list` with limit 1 on each.
|
||||
pub async fn list(&self) -> Vec<BucketInfo> {
|
||||
// Snapshot (name, backend, store_clone) under the lock, then probe
|
||||
// outside — probing can be slow (network), don't hold the lock.
|
||||
let snapshot: Vec<(String, String, Arc<dyn ObjectStore>)> = {
|
||||
let guard = self.buckets.read().expect("bucket registry lock poisoned");
|
||||
guard
|
||||
.iter()
|
||||
.map(|(n, e)| (n.clone(), e.backend.clone(), e.store.clone()))
|
||||
.collect()
|
||||
};
|
||||
|
||||
let mut out = Vec::with_capacity(snapshot.len());
|
||||
for (name, backend, store) in snapshot {
|
||||
let reachable = probe(&store).await;
|
||||
let role = self.classify(&name);
|
||||
out.push(BucketInfo { name, backend, reachable, role });
|
||||
}
|
||||
out.sort_by(|a, b| a.name.cmp(&b.name));
|
||||
out
|
||||
}
|
||||
|
||||
fn classify(&self, name: &str) -> BucketRole {
|
||||
if name == self.default { BucketRole::Primary }
|
||||
else if Some(name) == self.rescue.as_deref() { BucketRole::Rescue }
|
||||
else if name.starts_with("profile:") { BucketRole::Profile }
|
||||
else { BucketRole::Tenant }
|
||||
}
|
||||
|
||||
/// Provision + register a bucket at runtime (federation layer 2).
|
||||
/// Returns Err if the name is already registered. Local backends get
|
||||
/// their root directory created automatically.
|
||||
pub async fn add_bucket(&self, bc: BucketConfig) -> Result<BucketInfo, String> {
|
||||
if self.contains(&bc.name) {
|
||||
return Err(format!("bucket '{}' already registered", bc.name));
|
||||
}
|
||||
let store = build_store(&bc, self.secrets.as_ref()).await?;
|
||||
let reachable = probe(&store).await;
|
||||
let entry = Arc::new(BucketEntry {
|
||||
name: bc.name.clone(),
|
||||
backend: bc.backend.clone(),
|
||||
store,
|
||||
config: bc.clone(),
|
||||
});
|
||||
let role = self.classify(&bc.name);
|
||||
let info = BucketInfo {
|
||||
name: bc.name.clone(),
|
||||
backend: bc.backend.clone(),
|
||||
reachable,
|
||||
role,
|
||||
};
|
||||
{
|
||||
let mut guard = self.buckets.write().expect("bucket registry lock poisoned");
|
||||
guard.insert(bc.name.clone(), entry);
|
||||
}
|
||||
tracing::info!("registered bucket '{}' backend={} reachable={}",
|
||||
bc.name, bc.backend, reachable);
|
||||
Ok(info)
|
||||
}
|
||||
|
||||
/// Unregister a bucket. Refuses to remove `primary` or the configured
|
||||
/// rescue bucket — those are load-bearing. Caller is responsible for
|
||||
/// any final flush of bucket-local state before calling this.
|
||||
pub fn remove_bucket(&self, name: &str) -> Result<(), String> {
|
||||
if name == self.default {
|
||||
return Err("cannot remove primary bucket".into());
|
||||
}
|
||||
if Some(name) == self.rescue.as_deref() {
|
||||
return Err(format!("cannot remove rescue bucket '{name}'"));
|
||||
}
|
||||
let removed = {
|
||||
let mut guard = self.buckets.write().expect("bucket registry lock poisoned");
|
||||
guard.remove(name)
|
||||
};
|
||||
if removed.is_none() {
|
||||
return Err(format!("bucket '{name}' not registered"));
|
||||
}
|
||||
tracing::info!("unregistered bucket '{}'", name);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read with rescue-bucket fallback. If the target bucket fails and a
|
||||
/// rescue is configured, retries against rescue. Records every failure
|
||||
/// in the error journal.
|
||||
pub async fn read_smart(&self, bucket: &str, key: &str) -> Result<ReadOutcome, String> {
|
||||
let target_store = {
|
||||
let guard = self.buckets.read().expect("bucket registry lock poisoned");
|
||||
guard.get(bucket).map(|e| e.store.clone())
|
||||
.ok_or_else(|| format!("unknown bucket: {bucket}"))?
|
||||
};
|
||||
|
||||
match crate::ops::get(&target_store, key).await {
|
||||
Ok(data) => Ok(ReadOutcome {
|
||||
data, rescued: false,
|
||||
original_bucket: bucket.to_string(),
|
||||
served_by: bucket.to_string(),
|
||||
}),
|
||||
Err(err) => {
|
||||
// Record failure regardless of what happens next.
|
||||
self.journal.append(BucketErrorEvent::new_read(bucket, key, &err)).await;
|
||||
|
||||
// Try rescue, if any.
|
||||
if let Some(rescue_name) = &self.rescue {
|
||||
if rescue_name != bucket {
|
||||
let rescue_store = {
|
||||
let guard = self.buckets.read().expect("bucket registry lock poisoned");
|
||||
guard.get(rescue_name).map(|e| e.store.clone())
|
||||
};
|
||||
if let Some(rescue_store) = rescue_store {
|
||||
match crate::ops::get(&rescue_store, key).await {
|
||||
Ok(data) => {
|
||||
self.journal.mark_rescued_last(bucket, key).await;
|
||||
return Ok(ReadOutcome {
|
||||
data, rescued: true,
|
||||
original_bucket: bucket.to_string(),
|
||||
served_by: rescue_name.clone(),
|
||||
});
|
||||
}
|
||||
Err(rescue_err) => {
|
||||
return Err(format!(
|
||||
"read '{key}' failed in '{bucket}' ({err}); rescue '{rescue_name}' also failed ({rescue_err})"
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(format!("read '{key}' failed in '{bucket}': {err}"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Write always goes to target. No rescue fallback for writes — writes
|
||||
/// that silently vanish are the worst possible failure.
|
||||
pub async fn write_smart(
|
||||
&self,
|
||||
bucket: &str,
|
||||
key: &str,
|
||||
data: bytes::Bytes,
|
||||
) -> Result<(), String> {
|
||||
let target_store = {
|
||||
let guard = self.buckets.read().expect("bucket registry lock poisoned");
|
||||
guard.get(bucket).map(|e| e.store.clone())
|
||||
.ok_or_else(|| format!("unknown bucket: {bucket}"))?
|
||||
};
|
||||
match crate::ops::put(&target_store, key, data).await {
|
||||
Ok(()) => Ok(()),
|
||||
Err(err) => {
|
||||
self.journal.append(BucketErrorEvent::new_write(bucket, key, &err)).await;
|
||||
Err(format!("write '{key}' failed in '{bucket}': {err}"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Trivial reachability check — try to list with limit 0.
|
||||
async fn probe(store: &Arc<dyn ObjectStore>) -> bool {
|
||||
use futures::StreamExt;
|
||||
let mut stream = store.list(None);
|
||||
// Pulling the first item confirms the store responds. Empty bucket = ok.
|
||||
match stream.next().await {
|
||||
Some(Ok(_)) => true,
|
||||
None => true, // empty but reachable
|
||||
Some(Err(_)) => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a concrete ObjectStore from a BucketConfig.
|
||||
async fn build_store(
|
||||
bc: &BucketConfig,
|
||||
secrets: &dyn SecretsProvider,
|
||||
) -> Result<Arc<dyn ObjectStore>, String> {
|
||||
match bc.backend.as_str() {
|
||||
"local" => {
|
||||
let root = bc.root.as_deref()
|
||||
.ok_or_else(|| format!("bucket '{}' is backend=local but has no root", bc.name))?;
|
||||
std::fs::create_dir_all(root)
|
||||
.map_err(|e| format!("create bucket dir '{root}': {e}"))?;
|
||||
let fs = LocalFileSystem::new_with_prefix(root)
|
||||
.map_err(|e| format!("init local bucket '{}': {e}", bc.name))?;
|
||||
Ok(Arc::new(fs))
|
||||
}
|
||||
"s3" => {
|
||||
let handle = bc.secret_ref.as_deref()
|
||||
.ok_or_else(|| format!("s3 bucket '{}' has no secret_ref", bc.name))?;
|
||||
let creds: BucketCredentials = secrets.resolve(handle).await?;
|
||||
let s3_bucket = bc.bucket.as_deref()
|
||||
.ok_or_else(|| format!("s3 bucket '{}' has no `bucket` name", bc.name))?;
|
||||
let region = bc.region.as_deref().unwrap_or("us-east-1");
|
||||
|
||||
let mut builder = object_store::aws::AmazonS3Builder::new()
|
||||
.with_bucket_name(s3_bucket)
|
||||
.with_region(region)
|
||||
.with_access_key_id(&creds.access_key)
|
||||
.with_secret_access_key(&creds.secret_key);
|
||||
if let Some(endpoint) = &bc.endpoint {
|
||||
builder = builder.with_endpoint(endpoint);
|
||||
// MinIO and other S3-compatible services often run on plain
|
||||
// HTTP. object_store refuses HTTP by default — opt in when
|
||||
// a custom endpoint is configured (TLS endpoints work either way).
|
||||
if endpoint.starts_with("http://") {
|
||||
builder = builder.with_allow_http(true);
|
||||
}
|
||||
}
|
||||
let s3 = builder.build()
|
||||
.map_err(|e| format!("init s3 bucket '{}': {e}", bc.name))?;
|
||||
Ok(Arc::new(s3))
|
||||
}
|
||||
other => Err(format!("unknown backend '{other}' for bucket '{}'", bc.name)),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
[package]
|
||||
name = "vectord-lance"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
# Firewall crate: the Lance stack (Arrow 57, DataFusion 52) is isolated
|
||||
# from the rest of vectord (Arrow 55, DataFusion 47). Public API uses
|
||||
# only std types — no Arrow types cross the crate boundary — so no
|
||||
# version conflict propagates outward.
|
||||
#
|
||||
# See docs/ADR-019-vector-storage.md for the rationale: "productionizing
|
||||
# will need either workspace-wide upgrade or a firewall via a dedicated
|
||||
# vectord-lance crate." This is that firewall.
|
||||
|
||||
[dependencies]
|
||||
# S3 support: Lance delegates to its internal object_store crate when
|
||||
# given s3:// URIs. The "dynamodb" feature enables DynamoDB-based
|
||||
# commit locking for multi-writer S3; we don't need that (single-writer)
|
||||
# so just the base AWS/S3 feature is enough.
|
||||
# Lance 4.0 feature "aws" enables S3-compatible storage via its internal
|
||||
# object_store + opendal crates. Reads AWS_* env vars for credentials.
|
||||
lance = { version = "4.0", default-features = false, features = ["aws"] }
|
||||
lance-index = { version = "4.0", default-features = false }
|
||||
lance-linalg = { version = "4.0", default-features = false }
|
||||
|
||||
# These Arrow/Parquet versions MUST match Lance's expectations — Lance
|
||||
# re-exports their types through its API so any mismatch is a hard
|
||||
# compile error. Keep in sync with lance-bench.
|
||||
arrow = "57"
|
||||
arrow-array = "57"
|
||||
arrow-schema = "57"
|
||||
parquet = "57"
|
||||
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
futures = "0.3"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
bytes = "1"
|
||||
@@ -0,0 +1,585 @@
|
||||
//! Production Lance vector backend (ADR-019 — hybrid architecture).
|
||||
//!
|
||||
//! This is the firewall crate. It owns its own Arrow 57 / DataFusion 52
|
||||
//! / Lance 4 dependency tree. The public API uses only std types
|
||||
//! (`Vec<f32>`, `Vec<String>`, `String`, `bool`) so nothing Arrow-shaped
|
||||
//! crosses the crate boundary. That keeps vectord (Arrow 55) from
|
||||
//! picking up an incompatible dep.
|
||||
//!
|
||||
//! Responsibilities:
|
||||
//! - Migrate an existing binary-blob Parquet vector file into a Lance
|
||||
//! dataset with `FixedSizeList<Float32, dims>`. One-time cost.
|
||||
//! - Append new rows natively (no full-file rewrite — Lance's structural win).
|
||||
//! - Build an IVF_PQ ANN index on the vector column.
|
||||
//! - Vector search (`search`) using the IVF_PQ index when present,
|
||||
//! falling back to full scan otherwise.
|
||||
//! - Random-access row fetch by `doc_id` (`get_by_doc_id`) — the O(1)
|
||||
//! lookup that Parquet-on-object-store can't cheaply do.
|
||||
//! - Cheap count + stats introspection.
|
||||
|
||||
use arrow_array::{
|
||||
Array, ArrayRef, BinaryArray, FixedSizeListArray, Float32Array, Int32Array,
|
||||
RecordBatch, RecordBatchIterator, StringArray,
|
||||
};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use futures::StreamExt;
|
||||
use serde::Serialize;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
// ================= Public types =================
|
||||
|
||||
/// One search result. Mirrors vectord's existing `SearchResult` shape
|
||||
/// structurally but carries simpler types so this crate stays firewalled.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct Hit {
|
||||
pub doc_id: String,
|
||||
pub chunk_text: String,
|
||||
pub score: f32,
|
||||
/// Optional — set by search_with_vector, not by index-only search.
|
||||
pub distance: Option<f32>,
|
||||
}
|
||||
|
||||
/// A fully-hydrated row fetched by `get_by_doc_id` — includes the vector
|
||||
/// so callers can do downstream work (rerank, cite, etc.) without a
|
||||
/// second round trip.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct Row {
|
||||
pub doc_id: String,
|
||||
pub chunk_text: String,
|
||||
pub vector: Vec<f32>,
|
||||
pub source: Option<String>,
|
||||
pub chunk_idx: Option<i32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct MigrationStats {
|
||||
pub rows_written: usize,
|
||||
pub dimensions: usize,
|
||||
pub disk_bytes: u64,
|
||||
pub duration_secs: f32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct AppendStats {
|
||||
pub rows_appended: usize,
|
||||
pub disk_bytes_added: u64,
|
||||
pub duration_secs: f32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct IndexStats {
|
||||
pub name: String,
|
||||
pub num_partitions: u32,
|
||||
pub num_bits: u32,
|
||||
pub num_sub_vectors: u32,
|
||||
pub build_time_secs: f32,
|
||||
pub disk_bytes_added: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct ScalarIndexStats {
|
||||
pub name: String,
|
||||
pub column: String,
|
||||
pub build_time_secs: f32,
|
||||
pub disk_bytes_added: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct DatasetStats {
|
||||
pub path: String,
|
||||
pub rows: usize,
|
||||
pub disk_bytes: u64,
|
||||
pub has_vector_index: bool,
|
||||
pub has_doc_id_index: bool,
|
||||
}
|
||||
|
||||
// ================= The backend =================
|
||||
|
||||
/// Thin wrapper around a Lance dataset path. Lance handles the heavy
|
||||
/// lifting — we just expose a narrow API.
|
||||
#[derive(Clone)]
|
||||
pub struct LanceVectorStore {
|
||||
/// Local filesystem path or object-store URI (file:///..., s3://...).
|
||||
/// Lance's internal URI parsing handles both.
|
||||
path: String,
|
||||
}
|
||||
|
||||
impl LanceVectorStore {
|
||||
pub fn new(path: impl Into<String>) -> Self {
|
||||
Self { path: path.into() }
|
||||
}
|
||||
|
||||
pub fn path(&self) -> &str { &self.path }
|
||||
|
||||
/// Row count via Lance's fast metadata path (no scan).
|
||||
pub async fn count(&self) -> Result<usize, String> {
|
||||
let dataset = open_or_err(&self.path).await?;
|
||||
let n = dataset.count_rows(None).await.map_err(e)?;
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
/// True if the on-disk dataset exists AND has at least one `vector`
|
||||
/// column index attached.
|
||||
pub async fn has_vector_index(&self) -> Result<bool, String> {
|
||||
use lance_index::DatasetIndexExt;
|
||||
let dataset = match lance::dataset::Dataset::open(&self.path).await {
|
||||
Ok(d) => d,
|
||||
Err(_) => return Ok(false),
|
||||
};
|
||||
let indexes = dataset.load_indices().await.map_err(e)?;
|
||||
Ok(indexes.iter().any(|ix| {
|
||||
ix.fields.iter().any(|fid| {
|
||||
dataset.schema().field_by_id(*fid)
|
||||
.map(|f| f.name == "vector")
|
||||
.unwrap_or(false)
|
||||
})
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn stats(&self) -> Result<DatasetStats, String> {
|
||||
let rows = self.count().await.unwrap_or(0);
|
||||
let disk_bytes = dir_size_bytes(&strip_file_uri(&self.path));
|
||||
let has_vector_index = self.has_vector_index().await.unwrap_or(false);
|
||||
let has_doc_id_index = self.has_scalar_index("doc_id").await.unwrap_or(false);
|
||||
Ok(DatasetStats {
|
||||
path: self.path.clone(),
|
||||
rows,
|
||||
disk_bytes,
|
||||
has_vector_index,
|
||||
has_doc_id_index,
|
||||
})
|
||||
}
|
||||
|
||||
/// Migrate a vectord-format Parquet file into a Lance dataset.
|
||||
///
|
||||
/// Input schema (vectord's binary-blob format):
|
||||
/// - source : Utf8
|
||||
/// - doc_id : Utf8
|
||||
/// - chunk_idx : Int32
|
||||
/// - chunk_text : Utf8
|
||||
/// - vector : Binary (raw f32 little-endian bytes)
|
||||
///
|
||||
/// Output schema (Lance-friendly):
|
||||
/// - source : Utf8
|
||||
/// - doc_id : Utf8
|
||||
/// - chunk_idx : Int32
|
||||
/// - chunk_text : Utf8
|
||||
/// - vector : FixedSizeList<Float32, dims>
|
||||
///
|
||||
/// Idempotent at the file level — if the target exists, it's
|
||||
/// overwritten. Caller must manage destination paths.
|
||||
pub async fn migrate_from_parquet_bytes(
|
||||
&self,
|
||||
parquet_bytes: &[u8],
|
||||
) -> Result<MigrationStats, String> {
|
||||
let t0 = Instant::now();
|
||||
let (schema, batches, rows) = read_parquet(parquet_bytes)?;
|
||||
let dims = detect_vector_dims(&batches)?;
|
||||
let (new_schema, new_batches) = convert_to_fixed_size_list(&schema, batches, dims)?;
|
||||
|
||||
// Overwrite any prior dataset at this path.
|
||||
let _ = std::fs::remove_dir_all(&strip_file_uri(&self.path));
|
||||
|
||||
write_dataset(&self.path, new_schema, new_batches).await?;
|
||||
let disk_bytes = dir_size_bytes(&strip_file_uri(&self.path));
|
||||
|
||||
Ok(MigrationStats {
|
||||
rows_written: rows,
|
||||
dimensions: dims,
|
||||
disk_bytes,
|
||||
duration_secs: t0.elapsed().as_secs_f32(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Native Lance append — does NOT rewrite existing files. New rows
|
||||
/// land as a separate fragment; readers union across fragments at
|
||||
/// query time. Contrast: our Parquet path requires rewriting the
|
||||
/// entire vector file to add rows.
|
||||
pub async fn append(
|
||||
&self,
|
||||
source: Option<String>,
|
||||
doc_ids: Vec<String>,
|
||||
chunk_idxs: Vec<i32>,
|
||||
chunk_texts: Vec<String>,
|
||||
vectors: Vec<Vec<f32>>,
|
||||
) -> Result<AppendStats, String> {
|
||||
let n = doc_ids.len();
|
||||
if n == 0 {
|
||||
return Ok(AppendStats { rows_appended: 0, disk_bytes_added: 0, duration_secs: 0.0 });
|
||||
}
|
||||
if chunk_idxs.len() != n || chunk_texts.len() != n || vectors.len() != n {
|
||||
return Err(format!(
|
||||
"append: length mismatch (doc_ids={n}, chunk_idxs={}, chunk_texts={}, vectors={})",
|
||||
chunk_idxs.len(), chunk_texts.len(), vectors.len(),
|
||||
));
|
||||
}
|
||||
let dims = vectors[0].len();
|
||||
for (i, v) in vectors.iter().enumerate() {
|
||||
if v.len() != dims {
|
||||
return Err(format!("append: row {i} has {} dims, expected {}", v.len(), dims));
|
||||
}
|
||||
}
|
||||
|
||||
let t0 = Instant::now();
|
||||
let pre_bytes = dir_size_bytes(&strip_file_uri(&self.path));
|
||||
|
||||
let src_arr = StringArray::from(
|
||||
(0..n).map(|_| source.clone()).collect::<Vec<_>>()
|
||||
);
|
||||
let doc_id_arr = StringArray::from(doc_ids);
|
||||
let chunk_idx_arr = Int32Array::from(chunk_idxs);
|
||||
let chunk_text_arr = StringArray::from(chunk_texts);
|
||||
|
||||
let mut flat: Vec<f32> = Vec::with_capacity(n * dims);
|
||||
for v in vectors { flat.extend(v); }
|
||||
let values = Float32Array::from(flat);
|
||||
let item_field = Arc::new(Field::new("item", DataType::Float32, true));
|
||||
let vec_arr = FixedSizeListArray::try_new(
|
||||
item_field.clone(), dims as i32, Arc::new(values), None,
|
||||
).map_err(e)?;
|
||||
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("source", DataType::Utf8, true),
|
||||
Field::new("doc_id", DataType::Utf8, false),
|
||||
Field::new("chunk_idx", DataType::Int32, true),
|
||||
Field::new("chunk_text", DataType::Utf8, true),
|
||||
Field::new("vector", DataType::FixedSizeList(item_field, dims as i32), false),
|
||||
]));
|
||||
|
||||
let arrays: Vec<ArrayRef> = vec![
|
||||
Arc::new(src_arr), Arc::new(doc_id_arr), Arc::new(chunk_idx_arr),
|
||||
Arc::new(chunk_text_arr), Arc::new(vec_arr),
|
||||
];
|
||||
let batch = RecordBatch::try_new(schema.clone(), arrays).map_err(e)?;
|
||||
let reader = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema);
|
||||
|
||||
use lance::dataset::{Dataset, WriteMode, WriteParams};
|
||||
let params = WriteParams { mode: WriteMode::Append, ..Default::default() };
|
||||
Dataset::write(reader, &self.path, Some(params)).await.map_err(e)?;
|
||||
|
||||
Ok(AppendStats {
|
||||
rows_appended: n,
|
||||
disk_bytes_added: dir_size_bytes(&strip_file_uri(&self.path)).saturating_sub(pre_bytes),
|
||||
duration_secs: t0.elapsed().as_secs_f32(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Build an IVF_PQ vector index. Replaces any prior index with the
|
||||
/// same name. Callers pass explicit params — sensible defaults for
|
||||
/// ~100K × 768d: num_partitions=316 (≈√N), num_bits=8, num_sub_vectors=48.
|
||||
pub async fn build_index(
|
||||
&self,
|
||||
num_partitions: u32,
|
||||
num_bits: u32,
|
||||
num_sub_vectors: u32,
|
||||
) -> Result<IndexStats, String> {
|
||||
use lance::dataset::Dataset;
|
||||
use lance::index::vector::VectorIndexParams;
|
||||
use lance_index::{DatasetIndexExt, IndexType};
|
||||
use lance_linalg::distance::MetricType;
|
||||
|
||||
let pre_bytes = dir_size_bytes(&strip_file_uri(&self.path));
|
||||
let t0 = Instant::now();
|
||||
|
||||
let mut dataset = Dataset::open(&self.path).await.map_err(e)?;
|
||||
let params = VectorIndexParams::ivf_pq(
|
||||
num_partitions as usize,
|
||||
num_bits as u8,
|
||||
num_sub_vectors as usize,
|
||||
MetricType::Cosine,
|
||||
50, // max_iterations — same as bench
|
||||
);
|
||||
dataset.create_index(
|
||||
&["vector"],
|
||||
IndexType::Vector,
|
||||
Some("vec_idx".into()),
|
||||
¶ms,
|
||||
true, // replace
|
||||
).await.map_err(e)?;
|
||||
|
||||
Ok(IndexStats {
|
||||
name: "vec_idx".into(),
|
||||
num_partitions,
|
||||
num_bits,
|
||||
num_sub_vectors,
|
||||
build_time_secs: t0.elapsed().as_secs_f32(),
|
||||
disk_bytes_added: dir_size_bytes(&strip_file_uri(&self.path)).saturating_sub(pre_bytes),
|
||||
})
|
||||
}
|
||||
|
||||
/// Build a scalar btree index on `doc_id`. Makes `get_by_doc_id`
|
||||
/// O(log N) instead of a filter-scan of every fragment. Small index
|
||||
/// footprint — a few hundred KB on 100K rows.
|
||||
pub async fn build_scalar_index(&self, column: &str) -> Result<ScalarIndexStats, String> {
|
||||
use lance::dataset::Dataset;
|
||||
use lance_index::{DatasetIndexExt, IndexType};
|
||||
use lance_index::scalar::ScalarIndexParams;
|
||||
|
||||
let pre_bytes = dir_size_bytes(&strip_file_uri(&self.path));
|
||||
let t0 = Instant::now();
|
||||
|
||||
let mut dataset = Dataset::open(&self.path).await.map_err(e)?;
|
||||
let idx_name = format!("{column}_btree");
|
||||
dataset.create_index(
|
||||
&[column],
|
||||
IndexType::Scalar,
|
||||
Some(idx_name.clone()),
|
||||
&ScalarIndexParams::default(),
|
||||
true,
|
||||
).await.map_err(e)?;
|
||||
|
||||
Ok(ScalarIndexStats {
|
||||
name: idx_name,
|
||||
column: column.to_string(),
|
||||
build_time_secs: t0.elapsed().as_secs_f32(),
|
||||
disk_bytes_added: dir_size_bytes(&strip_file_uri(&self.path)).saturating_sub(pre_bytes),
|
||||
})
|
||||
}
|
||||
|
||||
/// True if a scalar index exists on the named column.
|
||||
pub async fn has_scalar_index(&self, column: &str) -> Result<bool, String> {
|
||||
use lance_index::DatasetIndexExt;
|
||||
let dataset = match lance::dataset::Dataset::open(&self.path).await {
|
||||
Ok(d) => d,
|
||||
Err(_) => return Ok(false),
|
||||
};
|
||||
let indexes = dataset.load_indices().await.map_err(e)?;
|
||||
Ok(indexes.iter().any(|ix| {
|
||||
ix.fields.iter().any(|fid| {
|
||||
dataset.schema().field_by_id(*fid)
|
||||
.map(|f| f.name == column)
|
||||
.unwrap_or(false)
|
||||
})
|
||||
}))
|
||||
}
|
||||
|
||||
/// Search for top_k nearest neighbors of `query`. Uses the IVF_PQ
|
||||
/// index if one exists; otherwise does a full scan (slow but
|
||||
/// correct — useful during development before index build).
|
||||
pub async fn search(&self, query: &[f32], top_k: usize) -> Result<Vec<Hit>, String> {
|
||||
use lance::dataset::Dataset;
|
||||
|
||||
let dataset = Dataset::open(&self.path).await.map_err(e)?;
|
||||
let qarr = Float32Array::from(query.to_vec());
|
||||
|
||||
let mut scanner = dataset.scan();
|
||||
scanner.nearest("vector", &qarr, top_k as usize).map_err(e)?;
|
||||
scanner.project(&["doc_id", "chunk_text"]).map_err(e)?;
|
||||
|
||||
let mut stream = scanner.try_into_stream().await.map_err(e)?;
|
||||
let mut hits: Vec<Hit> = Vec::with_capacity(top_k);
|
||||
|
||||
while let Some(batch) = stream.next().await {
|
||||
let batch = batch.map_err(e)?;
|
||||
let doc_ids = batch.column_by_name("doc_id")
|
||||
.ok_or_else(|| "no doc_id column in search result".to_string())?
|
||||
.as_any().downcast_ref::<StringArray>()
|
||||
.ok_or_else(|| "doc_id is not StringArray".to_string())?;
|
||||
let chunk_texts = batch.column_by_name("chunk_text")
|
||||
.and_then(|c| c.as_any().downcast_ref::<StringArray>());
|
||||
// Lance tacks on a `_distance` column for nearest() queries.
|
||||
let distances = batch.column_by_name("_distance")
|
||||
.and_then(|c| c.as_any().downcast_ref::<Float32Array>());
|
||||
for row in 0..batch.num_rows() {
|
||||
let d = distances.map(|a| a.value(row));
|
||||
hits.push(Hit {
|
||||
doc_id: doc_ids.value(row).to_string(),
|
||||
chunk_text: chunk_texts.map(|a| a.value(row).to_string()).unwrap_or_default(),
|
||||
score: d.map(|d| 1.0 - d).unwrap_or(0.0), // cosine distance → similarity
|
||||
distance: d,
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(hits)
|
||||
}
|
||||
|
||||
/// Fetch one row by doc_id. Implementation: Lance filter-pushdown
|
||||
/// scan — O(1) with partition pruning on a proper btree index,
|
||||
/// O(N) on vector-only datasets (still far faster than reading
|
||||
/// the whole Parquet file). We don't build a scalar index on
|
||||
/// doc_id yet; that's a future optimization.
|
||||
pub async fn get_by_doc_id(&self, doc_id: &str) -> Result<Option<Row>, String> {
|
||||
use lance::dataset::Dataset;
|
||||
|
||||
let dataset = Dataset::open(&self.path).await.map_err(e)?;
|
||||
let filter = format!("doc_id = '{}'", doc_id.replace('\'', "''"));
|
||||
let mut scanner = dataset.scan();
|
||||
scanner.filter(&filter).map_err(e)?;
|
||||
scanner.limit(Some(1), None).map_err(e)?;
|
||||
let mut stream = scanner.try_into_stream().await.map_err(e)?;
|
||||
|
||||
while let Some(batch) = stream.next().await {
|
||||
let batch = batch.map_err(e)?;
|
||||
if batch.num_rows() == 0 { continue; }
|
||||
return Ok(Some(row_from_batch(&batch, 0)?));
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
// ================= Internal helpers =================
|
||||
|
||||
fn e<T: std::fmt::Display>(err: T) -> String { err.to_string() }
|
||||
|
||||
/// `file:///abs/path` → `/abs/path`. Leave other URI schemes as-is for
|
||||
/// helpers that only work on local paths (dir_size_bytes, remove_dir_all).
|
||||
fn strip_file_uri(uri: &str) -> String {
|
||||
uri.strip_prefix("file://").unwrap_or(uri).to_string()
|
||||
}
|
||||
|
||||
async fn open_or_err(path: &str) -> Result<lance::dataset::Dataset, String> {
|
||||
lance::dataset::Dataset::open(path).await.map_err(e)
|
||||
}
|
||||
|
||||
fn dir_size_bytes(path: &str) -> u64 {
|
||||
fn recurse(p: &std::path::Path) -> u64 {
|
||||
let Ok(meta) = std::fs::metadata(p) else { return 0; };
|
||||
if meta.is_file() { return meta.len(); }
|
||||
let Ok(entries) = std::fs::read_dir(p) else { return 0; };
|
||||
entries.filter_map(|e| e.ok()).map(|e| recurse(&e.path())).sum()
|
||||
}
|
||||
recurse(std::path::Path::new(path))
|
||||
}
|
||||
|
||||
fn read_parquet(bytes: &[u8]) -> Result<(Arc<Schema>, Vec<RecordBatch>, usize), String> {
|
||||
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
|
||||
let builder = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::copy_from_slice(bytes))
|
||||
.map_err(e)?;
|
||||
let schema = builder.schema().clone();
|
||||
let reader = builder.build().map_err(e)?;
|
||||
let batches: Vec<RecordBatch> = reader.collect::<Result<_, _>>().map_err(e)?;
|
||||
let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
|
||||
Ok((schema, batches, rows))
|
||||
}
|
||||
|
||||
fn detect_vector_dims(batches: &[RecordBatch]) -> Result<usize, String> {
|
||||
for batch in batches {
|
||||
let idx = batch.schema().index_of("vector")
|
||||
.map_err(|_| "no 'vector' column".to_string())?;
|
||||
let col = batch.column(idx);
|
||||
if let Some(binary) = col.as_any().downcast_ref::<BinaryArray>() {
|
||||
for i in 0..binary.len() {
|
||||
if !binary.is_null(i) {
|
||||
return Ok(binary.value(i).len() / 4);
|
||||
}
|
||||
}
|
||||
} else if let Some(fsl) = col.as_any().downcast_ref::<FixedSizeListArray>() {
|
||||
return Ok(fsl.value_length() as usize);
|
||||
}
|
||||
}
|
||||
Err("could not determine vector dimensions".into())
|
||||
}
|
||||
|
||||
fn convert_to_fixed_size_list(
|
||||
schema: &Arc<Schema>,
|
||||
batches: Vec<RecordBatch>,
|
||||
dims: usize,
|
||||
) -> Result<(Arc<Schema>, Vec<RecordBatch>), String> {
|
||||
let new_fields: Vec<Arc<Field>> = schema
|
||||
.fields()
|
||||
.iter()
|
||||
.map(|f| {
|
||||
if f.name() == "vector" {
|
||||
Arc::new(Field::new(
|
||||
"vector",
|
||||
DataType::FixedSizeList(
|
||||
Arc::new(Field::new("item", DataType::Float32, true)),
|
||||
dims as i32,
|
||||
),
|
||||
false,
|
||||
))
|
||||
} else {
|
||||
f.clone()
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let new_schema = Arc::new(Schema::new(new_fields));
|
||||
|
||||
let mut out = Vec::with_capacity(batches.len());
|
||||
for batch in batches {
|
||||
let vec_idx = batch.schema().index_of("vector").map_err(e)?;
|
||||
let mut new_cols: Vec<ArrayRef> = Vec::with_capacity(batch.num_columns());
|
||||
for (i, col) in batch.columns().iter().enumerate() {
|
||||
if i == vec_idx {
|
||||
if let Some(bin) = col.as_any().downcast_ref::<BinaryArray>() {
|
||||
new_cols.push(Arc::new(binary_to_fsl(bin, dims)?));
|
||||
} else if col.as_any().is::<FixedSizeListArray>() {
|
||||
// Already in the right shape — just clone.
|
||||
new_cols.push(col.clone());
|
||||
} else {
|
||||
return Err("vector column is neither Binary nor FixedSizeList".into());
|
||||
}
|
||||
} else {
|
||||
new_cols.push(col.clone());
|
||||
}
|
||||
}
|
||||
out.push(RecordBatch::try_new(new_schema.clone(), new_cols).map_err(e)?);
|
||||
}
|
||||
Ok((new_schema, out))
|
||||
}
|
||||
|
||||
fn binary_to_fsl(bin: &BinaryArray, dims: usize) -> Result<FixedSizeListArray, String> {
|
||||
let n = bin.len();
|
||||
let mut flat: Vec<f32> = Vec::with_capacity(n * dims);
|
||||
for i in 0..n {
|
||||
if bin.is_null(i) {
|
||||
flat.extend(std::iter::repeat(0.0).take(dims));
|
||||
continue;
|
||||
}
|
||||
let b = bin.value(i);
|
||||
if b.len() != dims * 4 {
|
||||
return Err(format!(
|
||||
"row {i}: {} bytes vs expected {} ({} × f32)",
|
||||
b.len(), dims * 4, dims,
|
||||
));
|
||||
}
|
||||
for c in b.chunks_exact(4) {
|
||||
flat.push(f32::from_le_bytes([c[0], c[1], c[2], c[3]]));
|
||||
}
|
||||
}
|
||||
let values = Float32Array::from(flat);
|
||||
let field = Arc::new(Field::new("item", DataType::Float32, true));
|
||||
FixedSizeListArray::try_new(field, dims as i32, Arc::new(values), None).map_err(e)
|
||||
}
|
||||
|
||||
async fn write_dataset(
|
||||
path: &str,
|
||||
schema: Arc<Schema>,
|
||||
batches: Vec<RecordBatch>,
|
||||
) -> Result<(), String> {
|
||||
use lance::dataset::{Dataset, WriteParams};
|
||||
let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema);
|
||||
Dataset::write(reader, path, Some(WriteParams::default()))
|
||||
.await.map_err(e)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn row_from_batch(batch: &RecordBatch, row: usize) -> Result<Row, String> {
|
||||
let doc_id = batch.column_by_name("doc_id")
|
||||
.and_then(|c| c.as_any().downcast_ref::<StringArray>())
|
||||
.map(|a| a.value(row).to_string())
|
||||
.ok_or_else(|| "missing doc_id".to_string())?;
|
||||
let chunk_text = batch.column_by_name("chunk_text")
|
||||
.and_then(|c| c.as_any().downcast_ref::<StringArray>())
|
||||
.map(|a| a.value(row).to_string())
|
||||
.unwrap_or_default();
|
||||
let source = batch.column_by_name("source")
|
||||
.and_then(|c| c.as_any().downcast_ref::<StringArray>())
|
||||
.and_then(|a| if a.is_null(row) { None } else { Some(a.value(row).to_string()) });
|
||||
let chunk_idx = batch.column_by_name("chunk_idx")
|
||||
.and_then(|c| c.as_any().downcast_ref::<Int32Array>())
|
||||
.and_then(|a| if a.is_null(row) { None } else { Some(a.value(row)) });
|
||||
|
||||
let vec_col = batch.column_by_name("vector")
|
||||
.ok_or_else(|| "no vector column in row fetch".to_string())?;
|
||||
let fsl = vec_col.as_any().downcast_ref::<FixedSizeListArray>()
|
||||
.ok_or_else(|| "vector column is not FixedSizeList".to_string())?;
|
||||
let inner = fsl.value(row);
|
||||
let floats = inner.as_any().downcast_ref::<Float32Array>()
|
||||
.ok_or_else(|| "vector inner not Float32".to_string())?;
|
||||
let mut v = Vec::with_capacity(floats.len());
|
||||
for i in 0..floats.len() { v.push(floats.value(i)); }
|
||||
|
||||
Ok(Row { doc_id, chunk_text, vector: v, source, chunk_idx })
|
||||
}
|
||||
@@ -7,6 +7,10 @@ edition = "2024"
|
||||
shared = { path = "../shared" }
|
||||
storaged = { path = "../storaged" }
|
||||
aibridge = { path = "../aibridge" }
|
||||
catalogd = { path = "../catalogd" }
|
||||
# ADR-019 firewall — vectord-lance owns its own Arrow 57 / Lance 4 deps.
|
||||
# Public API uses only std types so no version conflict propagates here.
|
||||
vectord-lance = { path = "../vectord-lance" }
|
||||
tokio = { workspace = true }
|
||||
axum = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
@@ -18,3 +22,4 @@ parquet = { workspace = true }
|
||||
arrow = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
instant-distance = { workspace = true }
|
||||
uuid = { workspace = true }
|
||||
|
||||
@@ -0,0 +1,711 @@
|
||||
//! Phase 16.2 + 16.5 — The autotune agent.
|
||||
//!
|
||||
//! A long-running tokio task that watches the trial journal and
|
||||
//! autonomously proposes + runs new HNSW configs. Distinct from
|
||||
//! `autotune::run_autotune` which is synchronous (one HTTP call, grid
|
||||
//! of trials, done). The agent is the continuous version: it sleeps,
|
||||
//! wakes on triggers, proposes configs based on prior trial history,
|
||||
//! runs them one at a time, and auto-promotes when it finds an
|
||||
//! improvement.
|
||||
//!
|
||||
//! Design invariants:
|
||||
//! - Trials are data (ADR-018). Every proposal reads the journal; every
|
||||
//! attempt appends to it. The journal is the agent's memory.
|
||||
//! - One trial at a time. Bounded Ollama load — the agent never fires
|
||||
//! multiple parallel embeddings and respects `cooldown_between_trials_secs`.
|
||||
//! - Rate-limited. `max_trials_per_hour` is a hard ceiling so a
|
||||
//! misbehaving proposal function can't saturate the system.
|
||||
//! - Never promotes below `min_recall`. Same safety gate as
|
||||
//! `run_autotune` — we will not make the index worse.
|
||||
//! - Triggered OR periodic. Ingest enqueues a `DatasetAppended` event
|
||||
//! when a new batch lands; the agent also wakes periodically to keep
|
||||
//! exploring even when nothing changed externally.
|
||||
//! - Graceful shutdown via the `stop_tx` signal — the handle's Drop
|
||||
//! doesn't force-kill, but `stop()` requests a clean exit after the
|
||||
//! current trial.
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use object_store::ObjectStore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::{Mutex, RwLock, mpsc, oneshot};
|
||||
|
||||
use aibridge::client::AiClient;
|
||||
use catalogd::registry::Registry as CatalogRegistry;
|
||||
|
||||
use crate::embedding_cache::EmbeddingCache;
|
||||
use crate::harness;
|
||||
use crate::hnsw::HnswStore;
|
||||
use crate::index_registry::IndexRegistry;
|
||||
use crate::promotion::{PromotionEntry, PromotionRegistry};
|
||||
use crate::trial::{HnswConfig, Trial, TrialJournal, TrialMetrics};
|
||||
|
||||
// -------- Public-facing types --------
|
||||
|
||||
/// Runtime configuration for the agent. Mirrored in shared::config under
|
||||
/// `[agent]`. Defaults are conservative — designed to tune slowly in the
|
||||
/// background without fighting real workloads for GPU time.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AgentConfig {
|
||||
/// Master switch. When false, `spawn` returns a handle but the loop
|
||||
/// doesn't run.
|
||||
pub enabled: bool,
|
||||
/// Periodic wake-up — even if the trigger queue is empty, every N
|
||||
/// seconds the agent picks an index with trials and proposes one
|
||||
/// more config. Keeps exploration alive on idle indexes.
|
||||
pub cycle_interval_secs: u64,
|
||||
/// Minimum gap between two trials on the SAME index. Prevents the
|
||||
/// agent from hammering Ollama when a hot index has many pending
|
||||
/// triggers in a row.
|
||||
pub cooldown_between_trials_secs: u64,
|
||||
/// Below this recall, a proposal is never promoted — even if it
|
||||
/// beats the champion on latency.
|
||||
pub min_recall: f32,
|
||||
/// Budget cap: hard ceiling on trials per hour across all indexes.
|
||||
/// When hit, the agent idles until the hour window rolls.
|
||||
pub max_trials_per_hour: u32,
|
||||
}
|
||||
|
||||
impl Default for AgentConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: false, // opt-in — don't auto-start until J turns it on
|
||||
cycle_interval_secs: 60,
|
||||
cooldown_between_trials_secs: 30,
|
||||
min_recall: 0.9,
|
||||
max_trials_per_hour: 30,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What caused the agent to look at a particular index. Recorded in the
|
||||
/// trial's note field so we can tell "new data arrived" trials from
|
||||
/// "periodic exploration" trials in the journal.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum TriggerReason {
|
||||
/// Ingest just appended to a dataset that has attached HNSW indexes.
|
||||
DatasetAppended { dataset: String },
|
||||
/// A human or another agent hit `/agent/enqueue/{index}`.
|
||||
Manual,
|
||||
/// Periodic wake — no external event, just keep exploring.
|
||||
Periodic,
|
||||
}
|
||||
|
||||
/// One unit of work for the agent.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TriggerEvent {
|
||||
pub index_name: String,
|
||||
pub reason: TriggerReason,
|
||||
pub enqueued_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
impl TriggerEvent {
|
||||
pub fn manual(index_name: impl Into<String>) -> Self {
|
||||
Self { index_name: index_name.into(), reason: TriggerReason::Manual, enqueued_at: Utc::now() }
|
||||
}
|
||||
pub fn dataset_appended(index_name: impl Into<String>, dataset: impl Into<String>) -> Self {
|
||||
Self {
|
||||
index_name: index_name.into(),
|
||||
reason: TriggerReason::DatasetAppended { dataset: dataset.into() },
|
||||
enqueued_at: Utc::now(),
|
||||
}
|
||||
}
|
||||
pub fn periodic(index_name: impl Into<String>) -> Self {
|
||||
Self { index_name: index_name.into(), reason: TriggerReason::Periodic, enqueued_at: Utc::now() }
|
||||
}
|
||||
}
|
||||
|
||||
/// Observable snapshot of the agent's state — what `/agent/status` returns.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct AgentStatus {
|
||||
pub running: bool,
|
||||
pub config: AgentConfig,
|
||||
pub queue_depth: usize,
|
||||
pub trials_run: u64,
|
||||
pub promotions: u64,
|
||||
pub trials_in_last_hour: u32,
|
||||
pub last_event: Option<AgentEvent>,
|
||||
pub started_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
/// Last thing that happened — useful for "why didn't it do anything?" debugging.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct AgentEvent {
|
||||
pub at: DateTime<Utc>,
|
||||
pub kind: String, // "trial_completed" | "promoted" | "skipped_rate_limit" | etc
|
||||
pub index_name: Option<String>,
|
||||
pub detail: String,
|
||||
}
|
||||
|
||||
/// Handle returned by `spawn`. Holds the trigger sender + shared status +
|
||||
/// stop signal.
|
||||
#[derive(Clone)]
|
||||
pub struct AgentHandle {
|
||||
trigger_tx: mpsc::Sender<TriggerEvent>,
|
||||
inner: Arc<AgentInner>,
|
||||
}
|
||||
|
||||
struct AgentInner {
|
||||
status: RwLock<AgentStatus>,
|
||||
stop_tx: Mutex<Option<oneshot::Sender<()>>>,
|
||||
queue_len: Mutex<usize>, // mirror of the channel capacity — for status reporting
|
||||
recent_trials: Mutex<VecDeque<DateTime<Utc>>>, // ring of recent trial timestamps for rate limit
|
||||
}
|
||||
|
||||
impl AgentHandle {
|
||||
/// Enqueue a trigger. Returns Err if the agent isn't running or the
|
||||
/// queue is full (backpressure — dropping events is correct here
|
||||
/// since periodic exploration will pick up the slack).
|
||||
pub async fn enqueue(&self, event: TriggerEvent) -> Result<(), String> {
|
||||
self.trigger_tx.try_send(event).map_err(|e| format!("enqueue: {e}"))?;
|
||||
let mut guard = self.inner.queue_len.lock().await;
|
||||
*guard = guard.saturating_add(1);
|
||||
// Update queue_depth in status for observability.
|
||||
let mut s = self.inner.status.write().await;
|
||||
s.queue_depth = *guard;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn status(&self) -> AgentStatus {
|
||||
let mut s = self.inner.status.read().await.clone();
|
||||
// Refresh rate-limit window from ring buffer.
|
||||
let cutoff = Utc::now() - chrono::Duration::hours(1);
|
||||
let ring = self.inner.recent_trials.lock().await;
|
||||
s.trials_in_last_hour = ring.iter().filter(|t| **t >= cutoff).count() as u32;
|
||||
s
|
||||
}
|
||||
|
||||
/// Request a graceful stop. Returns immediately — the loop exits
|
||||
/// after its current trial completes.
|
||||
pub async fn stop(&self) -> bool {
|
||||
let mut guard = self.inner.stop_tx.lock().await;
|
||||
if let Some(tx) = guard.take() {
|
||||
let _ = tx.send(());
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// -------- Agent state holder --------
|
||||
|
||||
/// Everything the agent needs to run a trial. Mirrors the fields of
|
||||
/// `VectorState` the agent actually uses. Kept separate so the service
|
||||
/// layer builds it explicitly — no clone of unneeded state.
|
||||
#[derive(Clone)]
|
||||
pub struct AgentDeps {
|
||||
pub store: Arc<dyn ObjectStore>,
|
||||
pub ai_client: AiClient,
|
||||
pub catalog: CatalogRegistry,
|
||||
pub index_registry: IndexRegistry,
|
||||
pub hnsw_store: HnswStore,
|
||||
pub embedding_cache: EmbeddingCache,
|
||||
pub trial_journal: TrialJournal,
|
||||
pub promotion_registry: PromotionRegistry,
|
||||
}
|
||||
|
||||
// -------- Spawn --------
|
||||
|
||||
/// Start the agent loop in a background tokio task. Returns a handle
|
||||
/// the caller uses to enqueue events and read status.
|
||||
pub fn spawn(config: AgentConfig, deps: AgentDeps) -> AgentHandle {
|
||||
let (trigger_tx, trigger_rx) = mpsc::channel::<TriggerEvent>(256);
|
||||
let (stop_tx, stop_rx) = oneshot::channel::<()>();
|
||||
|
||||
let status = AgentStatus {
|
||||
running: config.enabled,
|
||||
config: config.clone(),
|
||||
queue_depth: 0,
|
||||
trials_run: 0,
|
||||
promotions: 0,
|
||||
trials_in_last_hour: 0,
|
||||
last_event: None,
|
||||
started_at: if config.enabled { Some(Utc::now()) } else { None },
|
||||
};
|
||||
|
||||
let inner = Arc::new(AgentInner {
|
||||
status: RwLock::new(status),
|
||||
stop_tx: Mutex::new(Some(stop_tx)),
|
||||
queue_len: Mutex::new(0),
|
||||
recent_trials: Mutex::new(VecDeque::with_capacity(64)),
|
||||
});
|
||||
|
||||
if config.enabled {
|
||||
tracing::info!(
|
||||
"autotune agent started (cycle={}s, cooldown={}s, cap={}/hr, min_recall={})",
|
||||
config.cycle_interval_secs, config.cooldown_between_trials_secs,
|
||||
config.max_trials_per_hour, config.min_recall,
|
||||
);
|
||||
let loop_inner = inner.clone();
|
||||
let loop_deps = deps.clone();
|
||||
let loop_config = config.clone();
|
||||
tokio::spawn(async move {
|
||||
run_loop(loop_config, loop_deps, trigger_rx, stop_rx, loop_inner).await;
|
||||
});
|
||||
} else {
|
||||
// Agent disabled — still drain the channel so sends don't back up.
|
||||
tokio::spawn(async move {
|
||||
let mut rx = trigger_rx;
|
||||
while rx.recv().await.is_some() {}
|
||||
});
|
||||
tracing::info!("autotune agent configured but disabled (set [agent].enabled=true)");
|
||||
}
|
||||
|
||||
AgentHandle { trigger_tx, inner }
|
||||
}
|
||||
|
||||
// -------- Main loop --------
|
||||
|
||||
async fn run_loop(
|
||||
config: AgentConfig,
|
||||
deps: AgentDeps,
|
||||
mut trigger_rx: mpsc::Receiver<TriggerEvent>,
|
||||
mut stop_rx: oneshot::Receiver<()>,
|
||||
inner: Arc<AgentInner>,
|
||||
) {
|
||||
let mut periodic = tokio::time::interval(std::time::Duration::from_secs(config.cycle_interval_secs));
|
||||
// First tick fires immediately — skip it so we don't double-fire on startup.
|
||||
periodic.tick().await;
|
||||
|
||||
loop {
|
||||
let event = tokio::select! {
|
||||
_ = &mut stop_rx => {
|
||||
tracing::info!("autotune agent: stop signal received");
|
||||
let mut s = inner.status.write().await;
|
||||
s.running = false;
|
||||
return;
|
||||
}
|
||||
maybe = trigger_rx.recv() => match maybe {
|
||||
Some(ev) => {
|
||||
let mut guard = inner.queue_len.lock().await;
|
||||
*guard = guard.saturating_sub(1);
|
||||
let mut s = inner.status.write().await;
|
||||
s.queue_depth = *guard;
|
||||
ev
|
||||
}
|
||||
None => {
|
||||
tracing::info!("autotune agent: trigger channel closed");
|
||||
return;
|
||||
}
|
||||
},
|
||||
_ = periodic.tick() => {
|
||||
// Periodic wake — pick an index with existing trials.
|
||||
// If nothing's been tuned yet there's nothing to propose.
|
||||
match pick_periodic_target(&deps).await {
|
||||
Some(idx) => TriggerEvent::periodic(idx),
|
||||
None => continue,
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Rate limit: count recent trials, skip if over budget.
|
||||
if over_rate_limit(&inner, config.max_trials_per_hour).await {
|
||||
record_event(&inner, "skipped_rate_limit", Some(&event.index_name),
|
||||
format!("hit cap of {}/hour", config.max_trials_per_hour)).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Per-index cooldown.
|
||||
if cooling_down(&inner, &event.index_name, config.cooldown_between_trials_secs).await {
|
||||
record_event(&inner, "skipped_cooldown", Some(&event.index_name),
|
||||
format!("last trial too recent (<{}s)", config.cooldown_between_trials_secs)).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Run one trial.
|
||||
match run_one_cycle(&event, &deps, config.min_recall).await {
|
||||
Ok(outcome) => {
|
||||
mark_recent_trial(&inner).await;
|
||||
{
|
||||
let mut s = inner.status.write().await;
|
||||
s.trials_run += 1;
|
||||
if outcome.promoted { s.promotions += 1; }
|
||||
}
|
||||
record_event(&inner, if outcome.promoted { "promoted" } else { "trial_completed" },
|
||||
Some(&event.index_name),
|
||||
format!("config=ec{}/es{} recall={:.3} p50={:.0}us {}",
|
||||
outcome.trial.config.ef_construction,
|
||||
outcome.trial.config.ef_search,
|
||||
outcome.trial.metrics.recall_at_k,
|
||||
outcome.trial.metrics.search_latency_p50_us,
|
||||
if outcome.promoted { "★ PROMOTED" } else { "" })).await;
|
||||
}
|
||||
Err(e) => {
|
||||
record_event(&inner, "trial_error", Some(&event.index_name), e).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of one cycle — ran a trial, maybe promoted it.
|
||||
struct CycleOutcome {
|
||||
trial: Trial,
|
||||
promoted: bool,
|
||||
}
|
||||
|
||||
/// Core cycle: propose → build → bench → record → maybe promote.
|
||||
async fn run_one_cycle(
|
||||
event: &TriggerEvent,
|
||||
deps: &AgentDeps,
|
||||
min_recall: f32,
|
||||
) -> Result<CycleOutcome, String> {
|
||||
// Read history.
|
||||
let history = deps.trial_journal.list(&event.index_name).await
|
||||
.map_err(|e| format!("read journal: {e}"))?;
|
||||
if history.is_empty() {
|
||||
return Err(format!(
|
||||
"no trials yet for '{}' — seed with at least one POST /hnsw/trial first",
|
||||
event.index_name,
|
||||
));
|
||||
}
|
||||
|
||||
// Current champion (if any) is the promoted config.
|
||||
let champion = deps.promotion_registry.get_current(&event.index_name).await;
|
||||
let champion_trial = champion.as_ref().and_then(|p| {
|
||||
history.iter().find(|t| t.id == p.trial_id).cloned()
|
||||
});
|
||||
|
||||
// Propose the next config.
|
||||
let Some(next_config) = propose_next_config(&history, champion_trial.as_ref()) else {
|
||||
return Err("proposer returned None — search space exhausted".into());
|
||||
};
|
||||
|
||||
// Validate bounds defensively.
|
||||
if !(10..=400).contains(&next_config.ef_construction) {
|
||||
return Err(format!("proposed ef_construction={} out of bounds", next_config.ef_construction));
|
||||
}
|
||||
if !(10..=200).contains(&next_config.ef_search) {
|
||||
return Err(format!("proposed ef_search={} out of bounds", next_config.ef_search));
|
||||
}
|
||||
|
||||
// Need a harness to measure. Use the most recent one in history.
|
||||
// (A future refinement: remember per-index "canonical harness" on
|
||||
// the index metadata. For now: latest wins.)
|
||||
let harness_name = history.last().unwrap().eval_set.clone();
|
||||
let mut harness_set = harness::EvalSet::load(&deps.store, &harness_name).await
|
||||
.map_err(|e| format!("load harness '{harness_name}': {e}"))?;
|
||||
|
||||
let embeddings = deps.embedding_cache.get_or_load(&event.index_name).await
|
||||
.map_err(|e| format!("embeddings: {e}"))?;
|
||||
|
||||
if !harness_set.ground_truth_built {
|
||||
harness::compute_ground_truth(&mut harness_set, &embeddings, &deps.ai_client).await
|
||||
.map_err(|e| format!("ground truth: {e}"))?;
|
||||
harness_set.save(&deps.store).await.ok();
|
||||
}
|
||||
|
||||
// Build + bench.
|
||||
let trial_id = Trial::new_id();
|
||||
let slot = format!("{}__{}", event.index_name, trial_id);
|
||||
let build = deps.hnsw_store
|
||||
.build_index_with_config(&slot, (*embeddings).clone(), &next_config)
|
||||
.await?;
|
||||
|
||||
let query_vectors: Vec<Vec<f32>> = harness_set.queries
|
||||
.iter().filter_map(|q| q.query_embedding.clone()).collect();
|
||||
let bench = deps.hnsw_store.bench_search(&slot, &query_vectors, harness_set.k).await?;
|
||||
|
||||
let mut recalls = Vec::with_capacity(harness_set.queries.len());
|
||||
for (q, hits) in harness_set.queries.iter().zip(bench.retrieved.iter()) {
|
||||
if let Some(gt) = &q.ground_truth {
|
||||
recalls.push(harness::recall_at_k(hits, gt, harness_set.k));
|
||||
}
|
||||
}
|
||||
let mean_recall = if recalls.is_empty() { 0.0 } else {
|
||||
recalls.iter().sum::<f32>() / recalls.len() as f32
|
||||
};
|
||||
|
||||
let mut lats = bench.latencies_us.clone();
|
||||
lats.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
let p = |pct: f32| -> f32 {
|
||||
if lats.is_empty() { return 0.0; }
|
||||
let idx = ((lats.len() as f32 - 1.0) * pct).round() as usize;
|
||||
lats[idx.min(lats.len() - 1)]
|
||||
};
|
||||
|
||||
let dims = embeddings.first().map(|e| e.vector.len()).unwrap_or(0);
|
||||
let memory_bytes = (embeddings.len() * dims * std::mem::size_of::<f32>() + embeddings.len() * 128) as u64;
|
||||
|
||||
let note = match &event.reason {
|
||||
TriggerReason::DatasetAppended { dataset } => format!("agent: dataset_appended({dataset})"),
|
||||
TriggerReason::Manual => "agent: manual".to_string(),
|
||||
TriggerReason::Periodic => "agent: periodic".to_string(),
|
||||
};
|
||||
|
||||
let trial = Trial {
|
||||
id: trial_id,
|
||||
index_name: event.index_name.clone(),
|
||||
eval_set: harness_set.name.clone(),
|
||||
config: next_config.clone(),
|
||||
metrics: TrialMetrics {
|
||||
build_time_secs: build.build_time_secs,
|
||||
search_latency_p50_us: p(0.50),
|
||||
search_latency_p95_us: p(0.95),
|
||||
search_latency_p99_us: p(0.99),
|
||||
recall_at_k: mean_recall,
|
||||
memory_bytes,
|
||||
vectors: build.vectors,
|
||||
eval_queries: harness_set.queries.len(),
|
||||
brute_force_latency_us: 0.0,
|
||||
},
|
||||
created_at: Utc::now(),
|
||||
note: Some(note),
|
||||
};
|
||||
deps.trial_journal.append(&trial).await.ok();
|
||||
deps.hnsw_store.drop(&slot).await;
|
||||
|
||||
// Promotion decision: the new trial must meet recall gate AND beat
|
||||
// the current champion (higher recall OR same recall + lower p50).
|
||||
let promoted = if trial.metrics.recall_at_k < min_recall {
|
||||
false
|
||||
} else {
|
||||
let beats = match &champion_trial {
|
||||
None => true, // no champion yet — anything passing the gate wins
|
||||
Some(c) => beats_champion(&trial, c),
|
||||
};
|
||||
if beats {
|
||||
let entry = PromotionEntry {
|
||||
config: trial.config.clone(),
|
||||
trial_id: trial.id.clone(),
|
||||
promoted_at: Utc::now(),
|
||||
promoted_by: "agent".to_string(),
|
||||
note: Some(format!(
|
||||
"auto-promote: recall={:.3} p50={:.0}us (was {:.3}/{:.0}us)",
|
||||
trial.metrics.recall_at_k, trial.metrics.search_latency_p50_us,
|
||||
champion_trial.as_ref().map(|t| t.metrics.recall_at_k).unwrap_or(0.0),
|
||||
champion_trial.as_ref().map(|t| t.metrics.search_latency_p50_us).unwrap_or(0.0),
|
||||
)),
|
||||
};
|
||||
deps.promotion_registry.promote(&event.index_name, entry).await.is_ok()
|
||||
} else {
|
||||
false
|
||||
}
|
||||
};
|
||||
|
||||
Ok(CycleOutcome { trial, promoted })
|
||||
}
|
||||
|
||||
/// Champion-beat test: strictly higher recall, OR equal recall with
|
||||
/// lower p50. Same rule as autotune::pick_winner — kept consistent so the
|
||||
/// agent and the synchronous autotune agree on what "better" means.
|
||||
fn beats_champion(candidate: &Trial, champion: &Trial) -> bool {
|
||||
if candidate.metrics.recall_at_k > champion.metrics.recall_at_k {
|
||||
return true;
|
||||
}
|
||||
if (candidate.metrics.recall_at_k - champion.metrics.recall_at_k).abs() < 1e-4
|
||||
&& candidate.metrics.search_latency_p50_us < champion.metrics.search_latency_p50_us {
|
||||
return true;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// Propose the next HnswConfig given trial history and the current
|
||||
/// champion.
|
||||
///
|
||||
/// ============================================================
|
||||
/// J: THIS IS YOURS TO IMPLEMENT
|
||||
/// ============================================================
|
||||
///
|
||||
/// Inputs:
|
||||
/// - `history`: every trial ever run on this index, oldest first
|
||||
/// - `champion`: the currently-promoted trial, if any
|
||||
///
|
||||
/// Output:
|
||||
/// - `Some(HnswConfig)` with the config to try next
|
||||
/// - `None` if you decide there's nothing worth trying (agent logs
|
||||
/// "search space exhausted" and moves on)
|
||||
///
|
||||
/// Hard bounds the caller enforces: ef_construction ∈ [10, 400],
|
||||
/// ef_search ∈ [10, 200]. Stay inside those — configs outside get
|
||||
/// rejected and count as a wasted cycle.
|
||||
///
|
||||
/// Design options (pick one, or mix):
|
||||
///
|
||||
/// 1. LOCAL REFINEMENT (exploit-heavy):
|
||||
/// Sample near champion ± small delta. Converges fast, risks local
|
||||
/// minima. Good for "we know roughly where the sweet spot is."
|
||||
///
|
||||
/// 2. ε-GREEDY (mixed):
|
||||
/// With prob ε, random sample from full bounds (explore). Otherwise
|
||||
/// refinement around champion (exploit). ε=0.2 is a reasonable start.
|
||||
/// Good for long-running tune with no prior knowledge.
|
||||
///
|
||||
/// 3. COARSE→FINE (annealed):
|
||||
/// First N trials: wide random. Then shrink the neighborhood around
|
||||
/// champion as more trials accumulate. Mimics simulated annealing.
|
||||
///
|
||||
/// 4. DEDUP-AWARE:
|
||||
/// Whatever strategy, skip configs already in history. Prevents the
|
||||
/// agent from re-running the same (ec, es) pair twice.
|
||||
///
|
||||
/// A simple starter implementation is provided below (local refinement
|
||||
/// + dedup). Replace with your preferred strategy.
|
||||
pub fn propose_next_config(history: &[Trial], champion: Option<&Trial>) -> Option<HnswConfig> {
|
||||
// ε-greedy around the champion, dedup-aware.
|
||||
//
|
||||
// - With probability ε (≈0.25), sample a random config from the full
|
||||
// bounds. Keeps exploration alive so we don't get stuck hill-climbing
|
||||
// one axis.
|
||||
// - Otherwise: perturb the champion symmetrically on BOTH axes (not
|
||||
// just +20 / +40 like the starter did). Prefers small moves first
|
||||
// so recall stays near the current level.
|
||||
// - Always skip configs already in history — no point re-running.
|
||||
// - Deterministic per-history: RNG is seeded from history length so
|
||||
// the same journal state always proposes the same next config.
|
||||
// Makes tests + offline replay reproducible.
|
||||
|
||||
let base = champion
|
||||
.map(|t| t.config.clone())
|
||||
.or_else(|| history.last().map(|t| t.config.clone()))
|
||||
.unwrap_or_default();
|
||||
|
||||
let tried = |ec: usize, es: usize| -> bool {
|
||||
history.iter().any(|t|
|
||||
t.config.ef_construction == ec && t.config.ef_search == es
|
||||
)
|
||||
};
|
||||
let clamp = |ec: i32, es: i32| -> (usize, usize) {
|
||||
(ec.clamp(10, 400) as usize, es.clamp(10, 200) as usize)
|
||||
};
|
||||
|
||||
// Tiny xorshift — no rand crate dep. Seeded from history length so the
|
||||
// proposer is deterministic for a given journal state.
|
||||
let mut rng = (history.len() as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
|
||||
let mut next_u = || {
|
||||
rng ^= rng << 13; rng ^= rng >> 7; rng ^= rng << 17;
|
||||
rng
|
||||
};
|
||||
|
||||
for _attempt in 0..32 {
|
||||
let explore = (next_u() % 100) < 25; // ε = 0.25
|
||||
let (ec, es) = if explore {
|
||||
clamp(10 + (next_u() % 391) as i32, 10 + (next_u() % 191) as i32)
|
||||
} else {
|
||||
// Symmetric perturbation — signed steps on both axes.
|
||||
let dec = [-40, -20, -10, 10, 20, 40][(next_u() % 6) as usize];
|
||||
let des = [-20, -10, -5, 5, 10, 20][(next_u() % 6) as usize];
|
||||
clamp(base.ef_construction as i32 + dec, base.ef_search as i32 + des)
|
||||
};
|
||||
if !tried(ec, es) {
|
||||
return Some(HnswConfig { ef_construction: ec, ef_search: es, seed: Some(42) });
|
||||
}
|
||||
}
|
||||
None // 32 attempts all landed on duplicates — likely saturated
|
||||
}
|
||||
|
||||
// -------- Helpers --------
|
||||
|
||||
/// Find an index to poke on a periodic wake. Strategy: the one with the
|
||||
/// most recent trial. If nothing's been trialed yet, return None.
|
||||
async fn pick_periodic_target(deps: &AgentDeps) -> Option<String> {
|
||||
// `/agent` runs against any index that has a trial journal. We don't
|
||||
// have a "list all journals" helper, so we derive candidates from the
|
||||
// promotion registry (indexes with a human ever promoting are live).
|
||||
let promos = deps.promotion_registry.list_all().await.ok()?;
|
||||
// Prefer the one most recently promoted — it's the one a human cares
|
||||
// about right now.
|
||||
promos.into_iter()
|
||||
.filter_map(|f| f.current.map(|c| (f.index_name, c.promoted_at)))
|
||||
.max_by_key(|(_, at)| *at)
|
||||
.map(|(name, _)| name)
|
||||
}
|
||||
|
||||
async fn over_rate_limit(inner: &Arc<AgentInner>, cap: u32) -> bool {
|
||||
let cutoff = Utc::now() - chrono::Duration::hours(1);
|
||||
let ring = inner.recent_trials.lock().await;
|
||||
ring.iter().filter(|t| **t >= cutoff).count() as u32 >= cap
|
||||
}
|
||||
|
||||
async fn cooling_down(inner: &Arc<AgentInner>, _index: &str, cooldown_secs: u64) -> bool {
|
||||
// Minimal impl: gate on global most-recent trial rather than per-index.
|
||||
// Per-index cooldown would be easy to add — keep a HashMap<String, DateTime>
|
||||
// — but for Phase 16.2 MVP, global is fine. Ollama is the shared resource.
|
||||
let ring = inner.recent_trials.lock().await;
|
||||
if let Some(last) = ring.back() {
|
||||
let since = Utc::now().signed_duration_since(*last);
|
||||
return since < chrono::Duration::seconds(cooldown_secs as i64);
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
async fn mark_recent_trial(inner: &Arc<AgentInner>) {
|
||||
let mut ring = inner.recent_trials.lock().await;
|
||||
ring.push_back(Utc::now());
|
||||
// Keep bounded.
|
||||
while ring.len() > 256 {
|
||||
ring.pop_front();
|
||||
}
|
||||
}
|
||||
|
||||
async fn record_event(
|
||||
inner: &Arc<AgentInner>,
|
||||
kind: &str,
|
||||
index: Option<&str>,
|
||||
detail: String,
|
||||
) {
|
||||
tracing::info!("agent: {} {}{}", kind,
|
||||
index.map(|i| format!("[{i}] ")).unwrap_or_default(), detail);
|
||||
let mut s = inner.status.write().await;
|
||||
s.last_event = Some(AgentEvent {
|
||||
at: Utc::now(),
|
||||
kind: kind.to_string(),
|
||||
index_name: index.map(String::from),
|
||||
detail,
|
||||
});
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn mk_trial(ec: usize, es: usize, recall: f32, p50: f32) -> Trial {
|
||||
Trial {
|
||||
id: format!("t-{ec}-{es}"),
|
||||
index_name: "test".into(),
|
||||
eval_set: "eval".into(),
|
||||
config: HnswConfig { ef_construction: ec, ef_search: es, seed: Some(42) },
|
||||
metrics: TrialMetrics {
|
||||
build_time_secs: 1.0,
|
||||
search_latency_p50_us: p50,
|
||||
search_latency_p95_us: p50 * 1.5,
|
||||
search_latency_p99_us: p50 * 2.0,
|
||||
recall_at_k: recall,
|
||||
memory_bytes: 0, vectors: 1000, eval_queries: 10,
|
||||
brute_force_latency_us: 100.0,
|
||||
},
|
||||
created_at: Utc::now(),
|
||||
note: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn propose_skips_duplicates() {
|
||||
let hist = vec![
|
||||
mk_trial(80, 30, 1.0, 500.0),
|
||||
mk_trial(100, 30, 1.0, 520.0), // ec+20
|
||||
];
|
||||
let next = propose_next_config(&hist, Some(&hist[0])).unwrap();
|
||||
// ec+20 is taken, so the proposer should skip it.
|
||||
assert!(next.ef_construction != 100 || next.ef_search != 30);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn beats_champion_strict_recall() {
|
||||
let champ = mk_trial(80, 30, 0.95, 500.0);
|
||||
let better_recall = mk_trial(80, 30, 0.99, 600.0);
|
||||
let worse_recall = mk_trial(80, 30, 0.90, 100.0);
|
||||
assert!(beats_champion(&better_recall, &champ));
|
||||
assert!(!beats_champion(&worse_recall, &champ));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn beats_champion_same_recall_lower_latency() {
|
||||
let champ = mk_trial(80, 30, 1.0, 500.0);
|
||||
let faster = mk_trial(60, 30, 1.0, 400.0);
|
||||
let slower = mk_trial(60, 30, 1.0, 600.0);
|
||||
assert!(beats_champion(&faster, &champ));
|
||||
assert!(!beats_champion(&slower, &champ));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,328 @@
|
||||
//! Phase 16: Autotune — the agent loop for HNSW tuning.
|
||||
//!
|
||||
//! Given an index and a harness, run a grid of (ef_construction,
|
||||
//! ef_search) trials, pick the Pareto winner, and promote it.
|
||||
//!
|
||||
//! This is deliberately NOT a sophisticated optimizer — random-grid
|
||||
//! search is good enough when the config space is small (a dozen
|
||||
//! combinations) and the work per trial is seconds, not hours. Think
|
||||
//! of it as a "smart operator" rather than Bayesian optimization.
|
||||
//!
|
||||
//! Score: we want high recall and low p50 latency. The function below
|
||||
//! picks the Pareto-dominating trial, breaking ties by prefer-lower-
|
||||
//! build-time (cheaper to rebuild on data changes).
|
||||
//!
|
||||
//! Safeguards (PRD risk mitigation):
|
||||
//! - Never promote a trial with recall < `min_recall` (default 0.9)
|
||||
//! - Hard config bounds: ef_construction ∈ [10, 400], ef_search ∈ [10, 200]
|
||||
//! - Concurrency cap: one autotune per index at a time (tracked by
|
||||
//! JobTracker — caller is responsible for not firing in parallel
|
||||
//! for the same index)
|
||||
|
||||
use chrono::Utc;
|
||||
use object_store::ObjectStore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
|
||||
use aibridge::client::AiClient;
|
||||
use catalogd::registry::Registry as CatalogRegistry;
|
||||
|
||||
use crate::embedding_cache::EmbeddingCache;
|
||||
use crate::harness;
|
||||
use crate::hnsw::HnswStore;
|
||||
use crate::index_registry::IndexRegistry;
|
||||
use crate::jobs::JobTracker;
|
||||
use crate::promotion::{PromotionEntry, PromotionRegistry};
|
||||
use crate::trial::{HnswConfig, Trial, TrialJournal, TrialMetrics};
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct AutotuneRequest {
|
||||
pub index_name: String,
|
||||
pub harness: String,
|
||||
/// Minimum recall a trial must achieve to be eligible for promotion.
|
||||
/// Default 0.9 — below that, the tuning is not producing useful
|
||||
/// answers. Never promote a bad index.
|
||||
#[serde(default = "default_min_recall")]
|
||||
pub min_recall: f32,
|
||||
/// Override the default config grid (rarely needed). When omitted,
|
||||
/// a small curated grid runs.
|
||||
#[serde(default)]
|
||||
pub grid: Option<Vec<HnswConfig>>,
|
||||
}
|
||||
|
||||
fn default_min_recall() -> f32 { 0.9 }
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct AutotuneResult {
|
||||
pub index_name: String,
|
||||
pub trials_run: usize,
|
||||
pub trials_eligible: usize,
|
||||
pub winner: Option<WinnerReport>,
|
||||
pub promoted: bool,
|
||||
pub duration_secs: f32,
|
||||
pub skipped: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct WinnerReport {
|
||||
pub trial_id: String,
|
||||
pub config: HnswConfig,
|
||||
pub recall_at_k: f32,
|
||||
pub search_latency_p50_us: f32,
|
||||
pub build_time_secs: f32,
|
||||
}
|
||||
|
||||
/// Default exploration grid. Covers the practical range of settings
|
||||
/// from "fast build, OK recall" to "thorough build, great recall".
|
||||
fn default_grid() -> Vec<HnswConfig> {
|
||||
vec![
|
||||
HnswConfig { ef_construction: 20, ef_search: 30, seed: Some(42) },
|
||||
HnswConfig { ef_construction: 40, ef_search: 30, seed: Some(42) },
|
||||
HnswConfig { ef_construction: 80, ef_search: 30, seed: Some(42) },
|
||||
HnswConfig { ef_construction: 80, ef_search: 60, seed: Some(42) },
|
||||
HnswConfig { ef_construction: 160, ef_search: 30, seed: Some(42) },
|
||||
]
|
||||
}
|
||||
|
||||
/// Config-space bounds. Rejects anything outside these before running a
|
||||
/// trial — defends against an agent getting creative with absurd values.
|
||||
fn sanitize_grid(grid: Vec<HnswConfig>) -> (Vec<HnswConfig>, Vec<String>) {
|
||||
let mut ok = Vec::new();
|
||||
let mut rejected = Vec::new();
|
||||
for cfg in grid {
|
||||
if !(10..=400).contains(&cfg.ef_construction) {
|
||||
rejected.push(format!(
|
||||
"ef_construction={} out of bounds [10,400]", cfg.ef_construction,
|
||||
));
|
||||
continue;
|
||||
}
|
||||
if !(10..=200).contains(&cfg.ef_search) {
|
||||
rejected.push(format!(
|
||||
"ef_search={} out of bounds [10,200]", cfg.ef_search,
|
||||
));
|
||||
continue;
|
||||
}
|
||||
ok.push(cfg);
|
||||
}
|
||||
(ok, rejected)
|
||||
}
|
||||
|
||||
/// Pareto winner: highest recall, then lowest p50 latency among
|
||||
/// eligible trials (recall >= min_recall).
|
||||
fn pick_winner(trials: &[Trial], min_recall: f32) -> Option<&Trial> {
|
||||
trials
|
||||
.iter()
|
||||
.filter(|t| t.metrics.recall_at_k >= min_recall)
|
||||
.min_by(|a, b| {
|
||||
// Compare by (−recall, p50_latency) so minimum is the
|
||||
// best trial: higher recall wins, then lower latency.
|
||||
let a_key = (-a.metrics.recall_at_k, a.metrics.search_latency_p50_us);
|
||||
let b_key = (-b.metrics.recall_at_k, b.metrics.search_latency_p50_us);
|
||||
a_key.partial_cmp(&b_key).unwrap_or(std::cmp::Ordering::Equal)
|
||||
})
|
||||
}
|
||||
|
||||
/// Run the autotune loop. Synchronous — caller can wrap in tokio::spawn
|
||||
/// to make it fire-and-forget. Uses the existing trial flow so every
|
||||
/// autotune step lands in the trial journal.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn run_autotune(
|
||||
req: AutotuneRequest,
|
||||
store: &Arc<dyn ObjectStore>,
|
||||
catalog: &CatalogRegistry,
|
||||
ai_client: &AiClient,
|
||||
embedding_cache: &EmbeddingCache,
|
||||
hnsw_store: &HnswStore,
|
||||
index_registry: &IndexRegistry,
|
||||
trial_journal: &TrialJournal,
|
||||
promotion_registry: &PromotionRegistry,
|
||||
_job_tracker: &JobTracker,
|
||||
) -> Result<AutotuneResult, String> {
|
||||
let t0 = std::time::Instant::now();
|
||||
|
||||
// Sanity: the index has to exist.
|
||||
if index_registry.get(&req.index_name).await.is_none() {
|
||||
return Err(format!("index not found: {}", req.index_name));
|
||||
}
|
||||
let _ = catalog; // reserved for future audit emission
|
||||
|
||||
// Load the harness once, compute ground truth once.
|
||||
let mut harness_set = harness::EvalSet::load(store, &req.harness)
|
||||
.await
|
||||
.map_err(|e| format!("load harness: {e}"))?;
|
||||
|
||||
let embeddings = embedding_cache
|
||||
.get_or_load(&req.index_name)
|
||||
.await
|
||||
.map_err(|e| format!("embeddings: {e}"))?;
|
||||
|
||||
if !harness_set.ground_truth_built {
|
||||
harness::compute_ground_truth(&mut harness_set, &embeddings, ai_client)
|
||||
.await
|
||||
.map_err(|e| format!("ground truth: {e}"))?;
|
||||
harness_set.save(store).await.ok();
|
||||
}
|
||||
|
||||
let (grid, rejected) = sanitize_grid(req.grid.clone().unwrap_or_else(default_grid));
|
||||
if grid.is_empty() {
|
||||
return Err("all configs rejected by bounds check".into());
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
"autotune '{}': {} trial configs, min_recall={}",
|
||||
req.index_name, grid.len(), req.min_recall,
|
||||
);
|
||||
|
||||
// Run each trial through the existing pipeline so every run lands
|
||||
// in the trial journal and contributes to the long-term tuning
|
||||
// memory. This is the llms3.com-style "trials are data" pattern.
|
||||
let mut run_trials: Vec<Trial> = Vec::with_capacity(grid.len());
|
||||
for config in &grid {
|
||||
let trial = run_single_trial(
|
||||
&req.index_name,
|
||||
&harness_set,
|
||||
&embeddings,
|
||||
config,
|
||||
hnsw_store,
|
||||
trial_journal,
|
||||
ai_client,
|
||||
).await;
|
||||
match trial {
|
||||
Ok(t) => run_trials.push(t),
|
||||
Err(e) => tracing::warn!("autotune: trial ec={} es={} failed: {e}",
|
||||
config.ef_construction, config.ef_search),
|
||||
}
|
||||
}
|
||||
|
||||
let trials_eligible = run_trials
|
||||
.iter()
|
||||
.filter(|t| t.metrics.recall_at_k >= req.min_recall)
|
||||
.count();
|
||||
|
||||
let winner = pick_winner(&run_trials, req.min_recall);
|
||||
let mut promoted = false;
|
||||
|
||||
let winner_report = if let Some(w) = winner {
|
||||
let entry = PromotionEntry {
|
||||
config: w.config.clone(),
|
||||
trial_id: w.id.clone(),
|
||||
promoted_at: Utc::now(),
|
||||
promoted_by: "autotune".to_string(),
|
||||
note: Some(format!(
|
||||
"autotune winner: recall={:.3} p50={:.0}us",
|
||||
w.metrics.recall_at_k, w.metrics.search_latency_p50_us,
|
||||
)),
|
||||
};
|
||||
if let Err(e) = promotion_registry.promote(&req.index_name, entry).await {
|
||||
tracing::warn!("autotune: promote failed: {e}");
|
||||
} else {
|
||||
promoted = true;
|
||||
}
|
||||
Some(WinnerReport {
|
||||
trial_id: w.id.clone(),
|
||||
config: w.config.clone(),
|
||||
recall_at_k: w.metrics.recall_at_k,
|
||||
search_latency_p50_us: w.metrics.search_latency_p50_us,
|
||||
build_time_secs: w.metrics.build_time_secs,
|
||||
})
|
||||
} else {
|
||||
tracing::warn!(
|
||||
"autotune '{}': no trial met min_recall={}; not promoting",
|
||||
req.index_name, req.min_recall,
|
||||
);
|
||||
None
|
||||
};
|
||||
|
||||
Ok(AutotuneResult {
|
||||
index_name: req.index_name,
|
||||
trials_run: run_trials.len(),
|
||||
trials_eligible,
|
||||
winner: winner_report,
|
||||
promoted,
|
||||
duration_secs: t0.elapsed().as_secs_f32(),
|
||||
skipped: rejected,
|
||||
})
|
||||
}
|
||||
|
||||
/// Run one trial and journal it. Returns the recorded Trial.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn run_single_trial(
|
||||
index_name: &str,
|
||||
harness_set: &harness::EvalSet,
|
||||
embeddings: &Arc<Vec<crate::store::StoredEmbedding>>,
|
||||
config: &HnswConfig,
|
||||
hnsw_store: &HnswStore,
|
||||
trial_journal: &TrialJournal,
|
||||
_ai_client: &AiClient,
|
||||
) -> Result<Trial, String> {
|
||||
let trial_id = Trial::new_id();
|
||||
let slot = format!("{}__{}", index_name, trial_id);
|
||||
|
||||
let build = hnsw_store
|
||||
.build_index_with_config(&slot, (**embeddings).clone(), config)
|
||||
.await?;
|
||||
|
||||
let query_vectors: Vec<Vec<f32>> = harness_set
|
||||
.queries
|
||||
.iter()
|
||||
.filter_map(|q| q.query_embedding.clone())
|
||||
.collect();
|
||||
let bench = hnsw_store
|
||||
.bench_search(&slot, &query_vectors, harness_set.k)
|
||||
.await?;
|
||||
|
||||
let mut recalls = Vec::with_capacity(harness_set.queries.len());
|
||||
for (q, hits) in harness_set.queries.iter().zip(bench.retrieved.iter()) {
|
||||
if let Some(gt) = &q.ground_truth {
|
||||
recalls.push(harness::recall_at_k(hits, gt, harness_set.k));
|
||||
}
|
||||
}
|
||||
let mean_recall = if recalls.is_empty() {
|
||||
0.0
|
||||
} else {
|
||||
recalls.iter().sum::<f32>() / recalls.len() as f32
|
||||
};
|
||||
|
||||
let mut lats = bench.latencies_us.clone();
|
||||
lats.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
let p = |pct: f32| -> f32 {
|
||||
if lats.is_empty() { return 0.0; }
|
||||
let idx = ((lats.len() as f32 - 1.0) * pct).round() as usize;
|
||||
lats[idx.min(lats.len() - 1)]
|
||||
};
|
||||
|
||||
let brute_us = if let Some(qv) = query_vectors.first() {
|
||||
let t = std::time::Instant::now();
|
||||
let _ = harness::brute_force_top_k(qv, embeddings, harness_set.k);
|
||||
t.elapsed().as_micros() as f32
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
let dims = embeddings.first().map(|e| e.vector.len()).unwrap_or(0);
|
||||
let memory_bytes =
|
||||
(embeddings.len() * dims * std::mem::size_of::<f32>() + embeddings.len() * 128) as u64;
|
||||
|
||||
let trial = Trial {
|
||||
id: trial_id.clone(),
|
||||
index_name: index_name.to_string(),
|
||||
eval_set: harness_set.name.clone(),
|
||||
config: config.clone(),
|
||||
metrics: TrialMetrics {
|
||||
build_time_secs: build.build_time_secs,
|
||||
search_latency_p50_us: p(0.50),
|
||||
search_latency_p95_us: p(0.95),
|
||||
search_latency_p99_us: p(0.99),
|
||||
recall_at_k: mean_recall,
|
||||
memory_bytes,
|
||||
vectors: build.vectors,
|
||||
eval_queries: harness_set.queries.len(),
|
||||
brute_force_latency_us: brute_us,
|
||||
},
|
||||
created_at: Utc::now(),
|
||||
note: Some("autotune".to_string()),
|
||||
};
|
||||
let _ = trial_journal.append(&trial).await;
|
||||
let _ = hnsw_store.drop(&slot).await;
|
||||
Ok(trial)
|
||||
}
|
||||
+106
-17
@@ -15,9 +15,28 @@ pub struct TextChunk {
|
||||
pub text: String,
|
||||
}
|
||||
|
||||
/// Round `idx` UP to the nearest UTF-8 char boundary in `s`.
|
||||
/// Needed because `&str[a..b]` panics if `a` or `b` lands mid-character —
|
||||
/// and any text imported from Postgres/user data will contain multi-byte
|
||||
/// chars like `\u{202f}` (narrow no-break space), em-dashes, etc.
|
||||
fn ceil_char_boundary(s: &str, mut idx: usize) -> usize {
|
||||
if idx >= s.len() { return s.len(); }
|
||||
while idx <= s.len() && !s.is_char_boundary(idx) {
|
||||
idx += 1;
|
||||
}
|
||||
idx
|
||||
}
|
||||
|
||||
/// Split text into overlapping chunks.
|
||||
/// - `chunk_size`: target characters per chunk (not tokens — chars are a good proxy)
|
||||
/// - `overlap`: characters of overlap between consecutive chunks
|
||||
/// - `chunk_size`: target bytes per chunk (UTF-8 boundary-respecting)
|
||||
/// - `overlap`: bytes of overlap between consecutive chunks
|
||||
///
|
||||
/// Safety contract:
|
||||
/// - Never slices at a byte that isn't a UTF-8 char boundary
|
||||
/// - Always advances forward (monotonic `start`) — no infinite loops on
|
||||
/// pathological inputs
|
||||
/// - Bounded iteration — caps total chunks at `text.len()` to prevent
|
||||
/// runaway allocation on degenerate text/parameter combinations
|
||||
pub fn chunk_text(
|
||||
text: &str,
|
||||
source: &str,
|
||||
@@ -44,28 +63,44 @@ pub fn chunk_text(
|
||||
let mut start = 0;
|
||||
let mut idx = 0u32;
|
||||
|
||||
while start < text.len() {
|
||||
let end = (start + chunk_size).min(text.len());
|
||||
// Safety cap: a well-formed chunk_size advances `start` by at least
|
||||
// `chunk_size - overlap` each iteration. Even in degenerate cases we
|
||||
// should never emit more chunks than there are bytes of text.
|
||||
let max_chunks = text.len();
|
||||
let mut iterations = 0;
|
||||
|
||||
// Try to break at a sentence or paragraph boundary
|
||||
let chunk_text = &text[start..end];
|
||||
while start < text.len() {
|
||||
iterations += 1;
|
||||
if iterations > max_chunks {
|
||||
tracing::warn!(
|
||||
"chunker: exceeded safety cap ({}) on doc '{}' of length {} — stopping",
|
||||
max_chunks, doc_id, text.len(),
|
||||
);
|
||||
break;
|
||||
}
|
||||
|
||||
// UTF-8 safe upper bound: move forward to a char boundary.
|
||||
let end = ceil_char_boundary(text, (start + chunk_size).min(text.len()));
|
||||
|
||||
let chunk_text_slice = &text[start..end];
|
||||
let actual_end = if end < text.len() {
|
||||
// Look for last sentence boundary in the chunk
|
||||
if let Some(pos) = chunk_text.rfind(". ") {
|
||||
// Look for last sentence boundary inside the slice.
|
||||
if let Some(pos) = chunk_text_slice.rfind(". ") {
|
||||
start + pos + 2
|
||||
} else if let Some(pos) = chunk_text.rfind('\n') {
|
||||
} else if let Some(pos) = chunk_text_slice.rfind('\n') {
|
||||
start + pos + 1
|
||||
} else if let Some(pos) = chunk_text_slice.rfind(' ') {
|
||||
start + pos + 1
|
||||
} else {
|
||||
// Fall back to word boundary
|
||||
if let Some(pos) = chunk_text.rfind(' ') {
|
||||
start + pos + 1
|
||||
} else {
|
||||
end
|
||||
}
|
||||
end
|
||||
}
|
||||
} else {
|
||||
end
|
||||
};
|
||||
// All the rfind'd positions are at char boundaries (they match
|
||||
// ASCII byte patterns), so actual_end should already be safe.
|
||||
// Defensive belt-and-suspenders: round up anyway.
|
||||
let actual_end = ceil_char_boundary(text, actual_end);
|
||||
|
||||
let chunk = text[start..actual_end].trim();
|
||||
if !chunk.is_empty() {
|
||||
@@ -78,11 +113,24 @@ pub fn chunk_text(
|
||||
idx += 1;
|
||||
}
|
||||
|
||||
// Advance with overlap
|
||||
// Advance with overlap — must strictly progress to avoid infinite
|
||||
// loops on degenerate inputs (e.g. chunk_size ≤ overlap, or
|
||||
// boundary-finding returning a position at `start`).
|
||||
if actual_end >= text.len() {
|
||||
break;
|
||||
}
|
||||
start = if actual_end > overlap { actual_end - overlap } else { actual_end };
|
||||
let tentative = if actual_end > overlap { actual_end - overlap } else { actual_end };
|
||||
let next_start = ceil_char_boundary(text, tentative);
|
||||
if next_start <= start {
|
||||
// No progress — force forward one char to guarantee termination.
|
||||
let mut forced = start + 1;
|
||||
while forced < text.len() && !text.is_char_boundary(forced) {
|
||||
forced += 1;
|
||||
}
|
||||
start = forced;
|
||||
} else {
|
||||
start = next_start;
|
||||
}
|
||||
}
|
||||
|
||||
chunks
|
||||
@@ -132,4 +180,45 @@ mod tests {
|
||||
assert_eq!(chunks[0].source, "candidates");
|
||||
assert_eq!(chunks[0].doc_id, "CAND-001");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handles_multibyte_utf8_at_chunk_boundary() {
|
||||
// Regression: multi-byte unicode (narrow no-break space \u{202f},
|
||||
// 3 bytes) landing inside a chunk-size window used to panic with
|
||||
// "byte index N is not a char boundary". Found on pg-imported
|
||||
// data that contained editorial unicode punctuation.
|
||||
let text = "**Claim 1**\u{202f}\u{2013} *Each large language model has its own personality \
|
||||
and way of interacting.*\u{202f}\u{2013}\u{202f}Verdict:\u{202f}VERIFIED. \
|
||||
Reasoning: widely accepted observation that different models exhibit distinct response \
|
||||
styles stemming from variations in training data.".repeat(5);
|
||||
let chunks = chunk_text(&text, "test", "unicode-1", 500, 50);
|
||||
// Must not panic, must produce at least one chunk.
|
||||
assert!(!chunks.is_empty());
|
||||
// All chunks must be valid UTF-8 substrings of the original.
|
||||
for chunk in &chunks {
|
||||
assert!(!chunk.text.is_empty());
|
||||
assert!(text.contains(&chunk.text));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_infinite_loop_on_no_spaces() {
|
||||
// Long text with no spaces, sentences, or newlines — chunker must
|
||||
// still advance via the hard `end` fallback.
|
||||
let text = "a".repeat(5000);
|
||||
let chunks = chunk_text(&text, "test", "noSpace", 500, 50);
|
||||
assert!(!chunks.is_empty());
|
||||
// Expected ~5000/450 chunks (with overlap=50 → step=450)
|
||||
assert!(chunks.len() < 50, "chunk count {} seems excessive", chunks.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_infinite_loop_on_degenerate_params() {
|
||||
// chunk_size equals overlap — previously would have made zero
|
||||
// forward progress. Our safety forced-advance must catch this.
|
||||
let text = "a b c d e f g h i j k l m n o p q r s t u v w x y z".repeat(20);
|
||||
let chunks = chunk_text(&text, "test", "degen", 100, 100);
|
||||
assert!(!chunks.is_empty());
|
||||
assert!(chunks.len() <= text.len(), "runaway chunk count {}", chunks.len());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
/// In-memory cache for StoredEmbedding vectors.
|
||||
///
|
||||
/// Rationale: loading 100K embeddings from Parquet takes ~2-5s. When an AI agent
|
||||
/// iterates on HNSW parameters, each trial would repeat that cost. The cache
|
||||
/// pins embeddings in memory so trials reuse them.
|
||||
///
|
||||
/// This is a pure performance layer — the Parquet file is still the source of
|
||||
/// truth (ADR-008). Eviction is safe; worst case is one slow reload.
|
||||
|
||||
use object_store::ObjectStore;
|
||||
use serde::Serialize;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::store::{self, StoredEmbedding};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct EmbeddingCache {
|
||||
store: Arc<dyn ObjectStore>,
|
||||
cache: Arc<RwLock<HashMap<String, Arc<Vec<StoredEmbedding>>>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct CacheEntry {
|
||||
pub index_name: String,
|
||||
pub vectors: usize,
|
||||
pub dimensions: usize,
|
||||
pub memory_bytes: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct CacheStats {
|
||||
pub entries: Vec<CacheEntry>,
|
||||
pub total_memory_bytes: u64,
|
||||
}
|
||||
|
||||
impl EmbeddingCache {
|
||||
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||
Self {
|
||||
store,
|
||||
cache: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Return cached embeddings, loading from object storage on first request.
|
||||
pub async fn get_or_load(
|
||||
&self,
|
||||
index_name: &str,
|
||||
) -> Result<Arc<Vec<StoredEmbedding>>, String> {
|
||||
if let Some(cached) = self.cache.read().await.get(index_name) {
|
||||
return Ok(cached.clone());
|
||||
}
|
||||
|
||||
// Load under a write lock so concurrent callers only hit disk once.
|
||||
let mut guard = self.cache.write().await;
|
||||
if let Some(cached) = guard.get(index_name) {
|
||||
return Ok(cached.clone());
|
||||
}
|
||||
tracing::info!("embedding_cache: loading '{index_name}' from object storage");
|
||||
let t0 = std::time::Instant::now();
|
||||
let loaded = store::load_embeddings(&self.store, index_name).await?;
|
||||
let n = loaded.len();
|
||||
let arc = Arc::new(loaded);
|
||||
guard.insert(index_name.to_string(), arc.clone());
|
||||
tracing::info!(
|
||||
"embedding_cache: loaded '{index_name}' — {n} vectors in {:.2}s",
|
||||
t0.elapsed().as_secs_f32()
|
||||
);
|
||||
Ok(arc)
|
||||
}
|
||||
|
||||
pub async fn evict(&self, index_name: &str) -> bool {
|
||||
self.cache.write().await.remove(index_name).is_some()
|
||||
}
|
||||
|
||||
pub async fn stats(&self) -> CacheStats {
|
||||
let cache = self.cache.read().await;
|
||||
let mut entries = Vec::with_capacity(cache.len());
|
||||
let mut total: u64 = 0;
|
||||
for (name, embs) in cache.iter() {
|
||||
let dims = embs.first().map(|e| e.vector.len()).unwrap_or(0);
|
||||
// Rough estimate: vector data + chunk_text + metadata overhead.
|
||||
let vec_bytes = (embs.len() * dims * std::mem::size_of::<f32>()) as u64;
|
||||
let text_bytes: u64 = embs.iter().map(|e| e.chunk_text.len() as u64).sum();
|
||||
let overhead = (embs.len() * 128) as u64; // strings + struct overhead
|
||||
let mem = vec_bytes + text_bytes + overhead;
|
||||
total += mem;
|
||||
entries.push(CacheEntry {
|
||||
index_name: name.clone(),
|
||||
vectors: embs.len(),
|
||||
dimensions: dims,
|
||||
memory_bytes: mem,
|
||||
});
|
||||
}
|
||||
CacheStats { entries, total_memory_bytes: total }
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,216 @@
|
||||
/// Eval harness for HNSW tuning.
|
||||
///
|
||||
/// An EvalSet is a named collection of queries against a specific vector index.
|
||||
/// Each query has a query_text and (optionally pre-computed) ground_truth
|
||||
/// doc_ids — the "correct" top-k results that brute-force cosine returns.
|
||||
/// HNSW trials are scored by how well they recreate that top-k (recall@k).
|
||||
///
|
||||
/// Storage: `_hnsw_evals/{name}.json` as a single JSON document. Small.
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use object_store::ObjectStore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
|
||||
use aibridge::client::{AiClient, EmbedRequest};
|
||||
use storaged::ops;
|
||||
|
||||
use crate::store::StoredEmbedding;
|
||||
|
||||
/// A single eval query with optional pre-computed ground truth.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct EvalQuery {
|
||||
pub id: String,
|
||||
pub query_text: String,
|
||||
/// Ordered list of doc_ids that brute-force returns as top-k.
|
||||
/// `None` means the ground truth hasn't been computed yet; `compute_ground_truth`
|
||||
/// fills it in against the current embedding set.
|
||||
#[serde(default)]
|
||||
pub ground_truth: Option<Vec<String>>,
|
||||
/// Optional pre-computed query embedding. Filling this avoids re-embedding
|
||||
/// the same query on every trial.
|
||||
#[serde(default)]
|
||||
pub query_embedding: Option<Vec<f32>>,
|
||||
}
|
||||
|
||||
/// A named eval set tied to a specific vector index.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct EvalSet {
|
||||
pub name: String,
|
||||
pub index_name: String,
|
||||
pub k: usize, // top-k used for recall calculation
|
||||
pub queries: Vec<EvalQuery>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub ground_truth_built: bool,
|
||||
}
|
||||
|
||||
impl EvalSet {
|
||||
pub fn new(name: &str, index_name: &str, k: usize) -> Self {
|
||||
Self {
|
||||
name: name.to_string(),
|
||||
index_name: index_name.to_string(),
|
||||
k,
|
||||
queries: Vec::new(),
|
||||
created_at: Utc::now(),
|
||||
ground_truth_built: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn storage_key(&self) -> String {
|
||||
format!("_hnsw_evals/{}.json", self.name)
|
||||
}
|
||||
|
||||
pub async fn save(&self, store: &Arc<dyn ObjectStore>) -> Result<(), String> {
|
||||
let json = serde_json::to_vec_pretty(self).map_err(|e| e.to_string())?;
|
||||
ops::put(store, &self.storage_key(), json.into()).await
|
||||
}
|
||||
|
||||
pub async fn load(store: &Arc<dyn ObjectStore>, name: &str) -> Result<Self, String> {
|
||||
let key = format!("_hnsw_evals/{}.json", name);
|
||||
let data = ops::get(store, &key).await?;
|
||||
serde_json::from_slice(&data).map_err(|e| format!("parse eval set: {e}"))
|
||||
}
|
||||
|
||||
pub async fn list(store: &Arc<dyn ObjectStore>) -> Result<Vec<String>, String> {
|
||||
let keys = ops::list(store, Some("_hnsw_evals/")).await?;
|
||||
Ok(keys
|
||||
.into_iter()
|
||||
.filter(|k| k.ends_with(".json"))
|
||||
.map(|k| {
|
||||
k.trim_start_matches("_hnsw_evals/")
|
||||
.trim_end_matches(".json")
|
||||
.to_string()
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
/// Cosine similarity for two same-length f32 slices.
|
||||
fn cosine(a: &[f32], b: &[f32]) -> f32 {
|
||||
let mut dot = 0.0f32;
|
||||
let mut na = 0.0f32;
|
||||
let mut nb = 0.0f32;
|
||||
for i in 0..a.len() {
|
||||
dot += a[i] * b[i];
|
||||
na += a[i] * a[i];
|
||||
nb += b[i] * b[i];
|
||||
}
|
||||
if na == 0.0 || nb == 0.0 {
|
||||
return 0.0;
|
||||
}
|
||||
dot / (na.sqrt() * nb.sqrt())
|
||||
}
|
||||
|
||||
/// Brute-force top-k search for a single query against all embeddings.
|
||||
/// This is the ground-truth oracle that HNSW trials must approximate.
|
||||
pub fn brute_force_top_k(
|
||||
query: &[f32],
|
||||
embeddings: &[StoredEmbedding],
|
||||
k: usize,
|
||||
) -> Vec<String> {
|
||||
let mut scored: Vec<(f32, usize)> = embeddings
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, e)| (cosine(query, &e.vector), i))
|
||||
.collect();
|
||||
// Partial sort — we only need top-k.
|
||||
scored.sort_unstable_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
|
||||
scored
|
||||
.into_iter()
|
||||
.take(k)
|
||||
.map(|(_, i)| embeddings[i].doc_id.clone())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Recall@k — fraction of ground-truth doc_ids that appear in predicted.
|
||||
pub fn recall_at_k(predicted: &[String], ground_truth: &[String], k: usize) -> f32 {
|
||||
if ground_truth.is_empty() || k == 0 {
|
||||
return 0.0;
|
||||
}
|
||||
let gt_set: std::collections::HashSet<&String> =
|
||||
ground_truth.iter().take(k).collect();
|
||||
let hits = predicted
|
||||
.iter()
|
||||
.take(k)
|
||||
.filter(|d| gt_set.contains(d))
|
||||
.count();
|
||||
hits as f32 / gt_set.len() as f32
|
||||
}
|
||||
|
||||
/// Populate query_embedding and ground_truth for every query that lacks them.
|
||||
pub async fn compute_ground_truth(
|
||||
eval: &mut EvalSet,
|
||||
embeddings: &[StoredEmbedding],
|
||||
ai_client: &AiClient,
|
||||
) -> Result<(), String> {
|
||||
let need_embed: Vec<(usize, String)> = eval
|
||||
.queries
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, q)| q.query_embedding.is_none())
|
||||
.map(|(i, q)| (i, q.query_text.clone()))
|
||||
.collect();
|
||||
|
||||
if !need_embed.is_empty() {
|
||||
// Embed in one batch to keep things simple; for very large eval sets
|
||||
// we'd batch this in chunks of 32.
|
||||
let texts: Vec<String> = need_embed.iter().map(|(_, t)| t.clone()).collect();
|
||||
let resp = ai_client
|
||||
.embed(EmbedRequest { texts, model: None })
|
||||
.await
|
||||
.map_err(|e| format!("embed queries: {e}"))?;
|
||||
for ((idx, _), vec) in need_embed.iter().zip(resp.embeddings.iter()) {
|
||||
let v: Vec<f32> = vec.iter().map(|&x| x as f32).collect();
|
||||
eval.queries[*idx].query_embedding = Some(v);
|
||||
}
|
||||
}
|
||||
|
||||
for q in eval.queries.iter_mut() {
|
||||
if q.ground_truth.is_some() {
|
||||
continue;
|
||||
}
|
||||
let emb = q.query_embedding.as_ref().ok_or("missing embedding")?;
|
||||
q.ground_truth = Some(brute_force_top_k(emb, embeddings, eval.k));
|
||||
}
|
||||
eval.ground_truth_built = true;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Auto-generate a synthetic eval set by sampling every Nth chunk's text as
|
||||
/// its own query. Useful for a quick-start eval when the user doesn't have
|
||||
/// real natural-language queries yet.
|
||||
pub fn synthetic_from_chunks(
|
||||
eval_name: &str,
|
||||
index_name: &str,
|
||||
embeddings: &[StoredEmbedding],
|
||||
sample_count: usize,
|
||||
k: usize,
|
||||
) -> EvalSet {
|
||||
let n = embeddings.len();
|
||||
let sample_count = sample_count.min(n);
|
||||
let stride = (n / sample_count.max(1)).max(1);
|
||||
|
||||
let mut queries = Vec::with_capacity(sample_count);
|
||||
for i in 0..sample_count {
|
||||
let idx = (i * stride).min(n - 1);
|
||||
let chunk = &embeddings[idx];
|
||||
// Use the first ~200 chars of the chunk as the "query" — it should find
|
||||
// itself and nearby chunks as top results.
|
||||
let query_text: String = chunk.chunk_text.chars().take(200).collect();
|
||||
queries.push(EvalQuery {
|
||||
id: format!("syn-{}", i),
|
||||
query_text,
|
||||
ground_truth: None,
|
||||
query_embedding: None,
|
||||
});
|
||||
}
|
||||
|
||||
EvalSet {
|
||||
name: eval_name.to_string(),
|
||||
index_name: index_name.to_string(),
|
||||
k,
|
||||
queries,
|
||||
created_at: Utc::now(),
|
||||
ground_truth_built: false,
|
||||
}
|
||||
}
|
||||
@@ -7,6 +7,7 @@ use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::store::StoredEmbedding;
|
||||
use crate::trial::HnswConfig;
|
||||
|
||||
/// A vector point for HNSW — wraps f32 slice with cosine distance.
|
||||
#[derive(Clone)]
|
||||
@@ -62,18 +63,32 @@ impl HnswStore {
|
||||
}
|
||||
}
|
||||
|
||||
/// Build an HNSW index from stored embeddings.
|
||||
/// Build an HNSW index from stored embeddings with default config.
|
||||
pub async fn build_index(
|
||||
&self,
|
||||
index_name: &str,
|
||||
embeddings: Vec<StoredEmbedding>,
|
||||
) -> Result<BuildStats, String> {
|
||||
self.build_index_with_config(index_name, embeddings, &HnswConfig::default()).await
|
||||
}
|
||||
|
||||
/// Build an HNSW index from stored embeddings with explicit config.
|
||||
/// Used by the trial system — each trial calls this with different params.
|
||||
pub async fn build_index_with_config(
|
||||
&self,
|
||||
index_name: &str,
|
||||
embeddings: Vec<StoredEmbedding>,
|
||||
config: &HnswConfig,
|
||||
) -> Result<BuildStats, String> {
|
||||
let n = embeddings.len();
|
||||
if n == 0 {
|
||||
return Err("no embeddings to index".into());
|
||||
}
|
||||
|
||||
tracing::info!("building HNSW index '{}' from {} vectors", index_name, n);
|
||||
tracing::info!(
|
||||
"building HNSW '{}': {} vectors, ef_construction={} ef_search={} seed={:?}",
|
||||
index_name, n, config.ef_construction, config.ef_search, config.seed,
|
||||
);
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
// Separate points and metadata
|
||||
@@ -92,14 +107,17 @@ impl HnswStore {
|
||||
values.push(i);
|
||||
}
|
||||
|
||||
// Build HNSW — this is the expensive part
|
||||
let map = Builder::default()
|
||||
.ef_construction(40) // balanced for 100K scale
|
||||
.ef_search(50) // fast search with good recall
|
||||
.build(points, values);
|
||||
// Build HNSW — the expensive part
|
||||
let mut builder = Builder::default()
|
||||
.ef_construction(config.ef_construction)
|
||||
.ef_search(config.ef_search);
|
||||
if let Some(seed) = config.seed {
|
||||
builder = builder.seed(seed);
|
||||
}
|
||||
let map = builder.build(points, values);
|
||||
|
||||
let build_time = start.elapsed().as_secs_f32();
|
||||
tracing::info!("HNSW index '{}' built: {} vectors in {:.1}s", index_name, n, build_time);
|
||||
tracing::info!("HNSW '{}' built: {} vectors in {:.1}s", index_name, n, build_time);
|
||||
|
||||
let index = Arc::new(HnswIndex { map, metadata });
|
||||
self.indexes.write().await.insert(index_name.to_string(), index);
|
||||
@@ -111,6 +129,43 @@ impl HnswStore {
|
||||
})
|
||||
}
|
||||
|
||||
/// Run a batch of search queries and return raw per-query latencies in microseconds.
|
||||
/// Also returns the retrieved doc_ids per query (for recall calculation).
|
||||
pub async fn bench_search(
|
||||
&self,
|
||||
index_name: &str,
|
||||
query_vectors: &[Vec<f32>],
|
||||
top_k: usize,
|
||||
) -> Result<BenchResult, String> {
|
||||
let indexes = self.indexes.read().await;
|
||||
let index = indexes
|
||||
.get(index_name)
|
||||
.ok_or_else(|| format!("HNSW index not found: {index_name}"))?
|
||||
.clone();
|
||||
drop(indexes);
|
||||
|
||||
let mut latencies_us = Vec::with_capacity(query_vectors.len());
|
||||
let mut retrieved: Vec<Vec<String>> = Vec::with_capacity(query_vectors.len());
|
||||
|
||||
for qv in query_vectors {
|
||||
let query_point = VecPoint(qv.clone());
|
||||
let t0 = std::time::Instant::now();
|
||||
let mut search = Search::default();
|
||||
let results = index.map.search(&query_point, &mut search);
|
||||
let ids: Vec<String> = results
|
||||
.take(top_k)
|
||||
.map(|item| {
|
||||
let meta_idx = *item.value;
|
||||
index.metadata[meta_idx].doc_id.clone()
|
||||
})
|
||||
.collect();
|
||||
latencies_us.push(t0.elapsed().as_micros() as f32);
|
||||
retrieved.push(ids);
|
||||
}
|
||||
|
||||
Ok(BenchResult { latencies_us, retrieved })
|
||||
}
|
||||
|
||||
/// Search an HNSW index. Returns approximate nearest neighbors.
|
||||
pub async fn search(
|
||||
&self,
|
||||
@@ -166,3 +221,9 @@ pub struct BuildStats {
|
||||
pub vectors: usize,
|
||||
pub build_time_secs: f32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BenchResult {
|
||||
pub latencies_us: Vec<f32>,
|
||||
pub retrieved: Vec<Vec<String>>,
|
||||
}
|
||||
|
||||
@@ -29,8 +29,21 @@ pub struct IndexMeta {
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub build_time_secs: f32,
|
||||
pub chunks_per_sec: f32,
|
||||
/// Federation layer 2: which bucket holds this index's artifacts
|
||||
/// (trial journal + promotion file). Defaults to "primary" for
|
||||
/// pre-federation indexes — the serde default keeps old metadata
|
||||
/// files readable without migration.
|
||||
#[serde(default = "default_bucket")]
|
||||
pub bucket: String,
|
||||
/// ADR-019: which physical backend stores this index. `Parquet`
|
||||
/// means storage_key points at our binary-blob Parquet file;
|
||||
/// `Lance` means it points at a Lance dataset directory.
|
||||
#[serde(default)]
|
||||
pub vector_backend: shared::types::VectorBackend,
|
||||
}
|
||||
|
||||
fn default_bucket() -> String { "primary".to_string() }
|
||||
|
||||
/// Registry of all vector indexes.
|
||||
#[derive(Clone)]
|
||||
pub struct IndexRegistry {
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
//! ADR-019 hybrid: routing layer for the Lance vector backend.
|
||||
//!
|
||||
//! Holds a `LanceRegistry` which maps an index name to its
|
||||
//! `LanceVectorStore` instance, lazy-creating on first touch. Path
|
||||
//! resolution: the index's bucket gives us a bucket root; we append
|
||||
//! `lance/{index_name}` and use that as the dataset URI. For local
|
||||
//! buckets that means a directory under the bucket's root.
|
||||
//!
|
||||
//! S3 buckets: in principle Lance accepts s3://... URIs, but the
|
||||
//! firewall crate (`vectord-lance`) doesn't pull the S3 feature in by
|
||||
//! default. When we promote profile buckets onto S3 in a future phase,
|
||||
//! enable that feature and update `lance_uri_for` accordingly.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use storaged::registry::BucketRegistry;
|
||||
use vectord_lance::LanceVectorStore;
|
||||
|
||||
use crate::index_registry::IndexRegistry;
|
||||
|
||||
/// Convert a bucket+index pair into the URI Lance should use as the
|
||||
/// dataset path. Supports both local (filesystem) and S3 buckets.
|
||||
///
|
||||
/// **Local buckets:** path resolution mirrors lakehouse.toml's convention.
|
||||
/// Returns an absolute filesystem path.
|
||||
///
|
||||
/// **S3 buckets:** returns `s3://{s3_bucket}/lance/{index_name}`. Lance's
|
||||
/// internal object_store crate reads `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY`
|
||||
/// / `AWS_ENDPOINT` from environment (or the S3 feature's default chain).
|
||||
/// For MinIO: set `AWS_ENDPOINT=http://localhost:9000` and
|
||||
/// `AWS_ALLOW_HTTP=true` before starting the gateway.
|
||||
///
|
||||
/// Refuses unknown buckets so a typo doesn't silently land Lance data
|
||||
/// in a directory / prefix the rest of the system can't see.
|
||||
pub fn lance_uri_for(
|
||||
buckets: &BucketRegistry,
|
||||
bucket: &str,
|
||||
index_name: &str,
|
||||
) -> Result<String, String> {
|
||||
if !buckets.contains(bucket) {
|
||||
return Err(format!("bucket '{bucket}' not registered"));
|
||||
}
|
||||
// Check if this bucket is S3-backed by looking for a bucket config
|
||||
// with backend="s3". BucketRegistry exposes backend type through
|
||||
// the list() info, but that's async. The simpler signal: if the
|
||||
// bucket name matches one we know is S3 (configured via lakehouse.toml
|
||||
// with backend="s3"), use the s3:// URI scheme.
|
||||
//
|
||||
// For the synchronous path, we check a naming convention: buckets
|
||||
// whose name starts with "s3:" are treated as S3 targets. The rest
|
||||
// of the name is the S3 bucket name. Convention-based, explicit,
|
||||
// no async needed.
|
||||
//
|
||||
// Additionally, any bucket registered with backend="s3" in the
|
||||
// config will have its BucketConfig.bucket field set — that's the
|
||||
// actual S3 bucket name. We can't access BucketConfig synchronously
|
||||
// from the current registry API, so for now the naming convention
|
||||
// is the primary signal.
|
||||
if bucket.starts_with("s3:") {
|
||||
let s3_bucket = &bucket["s3:".len()..];
|
||||
return Ok(format!("s3://{s3_bucket}/lance/{index_name}"));
|
||||
}
|
||||
|
||||
// Local path resolution.
|
||||
let root: PathBuf = match bucket {
|
||||
"primary" => PathBuf::from("./data"),
|
||||
"rescue" => PathBuf::from("./data/_rescue"),
|
||||
"testing" => PathBuf::from("./data/_testing"),
|
||||
b if b.starts_with("profile:") => {
|
||||
let safe = b.replace(':', "_");
|
||||
PathBuf::from(format!("./data/_profiles/{safe}"))
|
||||
}
|
||||
b => PathBuf::from(format!("./data/_buckets/{}", b.replace(':', "_"))),
|
||||
};
|
||||
let _ = std::fs::create_dir_all(root.join("lance"));
|
||||
let abs = match std::fs::canonicalize(&root) {
|
||||
Ok(p) => p.join("lance").join(index_name),
|
||||
Err(_) => root.join("lance").join(index_name),
|
||||
};
|
||||
Ok(abs.to_string_lossy().to_string())
|
||||
}
|
||||
|
||||
/// Lookup-by-index registry of `LanceVectorStore` handles. One handle
|
||||
/// per index, lazy-created on first call. Cheap to keep alive — it's
|
||||
/// just a path string + handle to Lance's metadata cache.
|
||||
#[derive(Clone)]
|
||||
pub struct LanceRegistry {
|
||||
buckets: Arc<BucketRegistry>,
|
||||
indexes: IndexRegistry,
|
||||
stores: Arc<RwLock<HashMap<String, LanceVectorStore>>>,
|
||||
}
|
||||
|
||||
impl LanceRegistry {
|
||||
pub fn new(buckets: Arc<BucketRegistry>, indexes: IndexRegistry) -> Self {
|
||||
Self {
|
||||
buckets,
|
||||
indexes,
|
||||
stores: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get (or lazy-create) the LanceVectorStore for an index.
|
||||
/// Resolves bucket → URI via the registry; backend stays whatever
|
||||
/// IndexMeta says (no enforcement here — caller decides whether to
|
||||
/// route to Lance based on profile/index backend choice).
|
||||
pub async fn store_for(&self, index_name: &str) -> Result<LanceVectorStore, String> {
|
||||
if let Some(s) = self.stores.read().await.get(index_name) {
|
||||
return Ok(s.clone());
|
||||
}
|
||||
let bucket = self.indexes
|
||||
.get(index_name).await
|
||||
.map(|m| m.bucket)
|
||||
.unwrap_or_else(|| "primary".to_string());
|
||||
let uri = lance_uri_for(&self.buckets, &bucket, index_name)?;
|
||||
let store = LanceVectorStore::new(uri);
|
||||
self.stores.write().await.insert(index_name.to_string(), store.clone());
|
||||
Ok(store)
|
||||
}
|
||||
|
||||
/// For freshly-created indexes that don't have IndexMeta yet — caller
|
||||
/// supplies the bucket explicitly.
|
||||
pub async fn store_for_new(&self, index_name: &str, bucket: &str) -> Result<LanceVectorStore, String> {
|
||||
let uri = lance_uri_for(&self.buckets, bucket, index_name)?;
|
||||
let store = LanceVectorStore::new(uri);
|
||||
self.stores.write().await.insert(index_name.to_string(), store.clone());
|
||||
Ok(store)
|
||||
}
|
||||
}
|
||||
@@ -1,9 +1,17 @@
|
||||
pub mod agent;
|
||||
pub mod autotune;
|
||||
pub mod lance_backend;
|
||||
pub mod chunker;
|
||||
pub mod embedding_cache;
|
||||
pub mod harness;
|
||||
pub mod hnsw;
|
||||
pub mod index_registry;
|
||||
pub mod jobs;
|
||||
pub mod promotion;
|
||||
pub mod refresh;
|
||||
pub mod store;
|
||||
pub mod search;
|
||||
pub mod rag;
|
||||
pub mod supervisor;
|
||||
pub mod service;
|
||||
pub mod trial;
|
||||
|
||||
@@ -0,0 +1,220 @@
|
||||
//! Phase 16: Promoted HNSW configs — the "active generation" pointer.
|
||||
//!
|
||||
//! An index's HNSW config used at build time normally defaults to the
|
||||
//! system-wide default (`HnswConfig::default()`). An operator or the
|
||||
//! autotune agent can *promote* a specific trial's config — subsequent
|
||||
//! HNSW builds against that index use the promoted config instead.
|
||||
//!
|
||||
//! Every promotion is history-tracked so `rollback` can revert. The
|
||||
//! history file lives at `primary://_hnsw_promotions/{index_name}.json`
|
||||
//! and is small (< few KB) so we rewrite it on every promotion rather
|
||||
//! than append-log.
|
||||
//!
|
||||
//! Not included here:
|
||||
//! - Atomic graph rebuild on promote — promotion only updates the sticky
|
||||
//! default. Next activation (or search that triggers lazy build) picks
|
||||
//! up the new config. That's "zero-downtime swap after build" which is
|
||||
//! what ADR-019 actually claimed; an instant-swap requires a
|
||||
//! pre-built graph pool which we don't have yet.
|
||||
//! - Agent loop — lives in `vectord::autotune`.
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use storaged::ops;
|
||||
use storaged::registry::BucketRegistry;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::index_registry::IndexRegistry;
|
||||
use crate::trial::HnswConfig;
|
||||
|
||||
const PROMOTION_PREFIX: &str = "_hnsw_promotions";
|
||||
|
||||
/// One promotion record. The `trial_id` is the origin of the config —
|
||||
/// lets operators trace back "why was this config picked?" to the exact
|
||||
/// trial in the trial journal.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PromotionEntry {
|
||||
pub config: HnswConfig,
|
||||
pub trial_id: String,
|
||||
pub promoted_at: DateTime<Utc>,
|
||||
#[serde(default)]
|
||||
pub promoted_by: String,
|
||||
#[serde(default)]
|
||||
pub note: Option<String>,
|
||||
}
|
||||
|
||||
/// Serialized form of an index's promotion history.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct PromotionFile {
|
||||
pub index_name: String,
|
||||
pub current: Option<PromotionEntry>,
|
||||
#[serde(default)]
|
||||
pub history: Vec<PromotionEntry>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct PromotionRegistry {
|
||||
buckets: Arc<BucketRegistry>,
|
||||
index_registry: IndexRegistry,
|
||||
cache: Arc<RwLock<HashMap<String, PromotionFile>>>,
|
||||
}
|
||||
|
||||
impl PromotionRegistry {
|
||||
pub fn new(buckets: Arc<BucketRegistry>, index_registry: IndexRegistry) -> Self {
|
||||
Self {
|
||||
buckets,
|
||||
index_registry,
|
||||
cache: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
fn key(index_name: &str) -> String {
|
||||
// Sanitize for object-store safety.
|
||||
let safe: String = index_name
|
||||
.chars()
|
||||
.map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' { c } else { '_' })
|
||||
.collect();
|
||||
format!("{PROMOTION_PREFIX}/{safe}.json")
|
||||
}
|
||||
|
||||
/// Resolve which bucket's store holds this index's promotion file.
|
||||
/// Same rules as TrialJournal::bucket_for — follows IndexMeta.bucket,
|
||||
/// defaults to primary when metadata is missing.
|
||||
async fn store_for(&self, index_name: &str) -> Result<Arc<dyn object_store::ObjectStore>, String> {
|
||||
let bucket = self.index_registry
|
||||
.get(index_name)
|
||||
.await
|
||||
.map(|m| m.bucket)
|
||||
.unwrap_or_else(|| "primary".to_string());
|
||||
self.buckets.get(&bucket)
|
||||
}
|
||||
|
||||
/// Load (and cache) the promotion file for an index.
|
||||
pub async fn load(&self, index_name: &str) -> Result<PromotionFile, String> {
|
||||
if let Some(cached) = self.cache.read().await.get(index_name) {
|
||||
return Ok(cached.clone());
|
||||
}
|
||||
let store = self.store_for(index_name).await?;
|
||||
let key = Self::key(index_name);
|
||||
let file = match ops::get(&store, &key).await {
|
||||
Ok(bytes) => serde_json::from_slice::<PromotionFile>(&bytes)
|
||||
.map_err(|e| format!("parse promotion file: {e}"))?,
|
||||
Err(_) => PromotionFile {
|
||||
index_name: index_name.to_string(),
|
||||
current: None,
|
||||
history: Vec::new(),
|
||||
},
|
||||
};
|
||||
self.cache.write().await.insert(index_name.to_string(), file.clone());
|
||||
Ok(file)
|
||||
}
|
||||
|
||||
/// Promote a config to the active slot. Pushes the current promotion
|
||||
/// (if any) onto the history stack. Persists before returning — the
|
||||
/// config is durable by the time this call completes.
|
||||
pub async fn promote(
|
||||
&self,
|
||||
index_name: &str,
|
||||
entry: PromotionEntry,
|
||||
) -> Result<PromotionFile, String> {
|
||||
let mut file = self.load(index_name).await?;
|
||||
if let Some(prior) = file.current.take() {
|
||||
file.history.push(prior);
|
||||
// Cap history to something sensible so this file doesn't grow
|
||||
// unbounded. 50 entries = 50 promotions — way more than any
|
||||
// sane workflow needs.
|
||||
const HISTORY_CAP: usize = 50;
|
||||
if file.history.len() > HISTORY_CAP {
|
||||
let drop = file.history.len() - HISTORY_CAP;
|
||||
file.history.drain(0..drop);
|
||||
}
|
||||
}
|
||||
file.current = Some(entry);
|
||||
file.index_name = index_name.to_string();
|
||||
|
||||
let store = self.store_for(index_name).await?;
|
||||
let key = Self::key(index_name);
|
||||
let json = serde_json::to_vec_pretty(&file).map_err(|e| e.to_string())?;
|
||||
ops::put(&store, &key, json.into()).await?;
|
||||
|
||||
self.cache.write().await.insert(index_name.to_string(), file.clone());
|
||||
tracing::info!(
|
||||
"promoted '{}' to config {:?} (trial={})",
|
||||
index_name, file.current.as_ref().unwrap().config, file.current.as_ref().unwrap().trial_id,
|
||||
);
|
||||
Ok(file)
|
||||
}
|
||||
|
||||
/// Pop the latest promotion back onto the current slot (if any
|
||||
/// history exists). If current is set but history is empty, the
|
||||
/// current promotion is cleared — the index falls back to defaults.
|
||||
pub async fn rollback(&self, index_name: &str) -> Result<PromotionFile, String> {
|
||||
let mut file = self.load(index_name).await?;
|
||||
match file.history.pop() {
|
||||
Some(prev) => {
|
||||
file.current = Some(prev);
|
||||
}
|
||||
None => {
|
||||
if file.current.is_none() {
|
||||
return Err(format!("no promotion to rollback for '{index_name}'"));
|
||||
}
|
||||
file.current = None;
|
||||
}
|
||||
}
|
||||
let store = self.store_for(index_name).await?;
|
||||
let key = Self::key(index_name);
|
||||
let json = serde_json::to_vec_pretty(&file).map_err(|e| e.to_string())?;
|
||||
ops::put(&store, &key, json.into()).await?;
|
||||
self.cache.write().await.insert(index_name.to_string(), file.clone());
|
||||
tracing::info!("rolled back promotion for '{}'", index_name);
|
||||
Ok(file)
|
||||
}
|
||||
|
||||
/// Get the currently-promoted config (if any). Callers use this to
|
||||
/// pick the right HnswConfig at build time.
|
||||
pub async fn get_current(&self, index_name: &str) -> Option<PromotionEntry> {
|
||||
self.load(index_name).await.ok().and_then(|f| f.current)
|
||||
}
|
||||
|
||||
/// Convenience: return the promoted config or the provided default.
|
||||
pub async fn config_or(&self, index_name: &str, default: HnswConfig) -> HnswConfig {
|
||||
match self.get_current(index_name).await {
|
||||
Some(entry) => entry.config,
|
||||
None => default,
|
||||
}
|
||||
}
|
||||
|
||||
/// List every index that has a promotion recorded (for operator UI).
|
||||
///
|
||||
/// Federation: scans EVERY registered bucket for promotion files.
|
||||
/// Per-profile buckets each have their own `_hnsw_promotions/` so we
|
||||
/// aggregate across them. Dedups by index_name — if the same index
|
||||
/// somehow has promotion files in multiple buckets, the one from the
|
||||
/// bucket recorded in IndexMeta wins.
|
||||
pub async fn list_all(&self) -> Result<Vec<PromotionFile>, String> {
|
||||
let bucket_infos = self.buckets.list().await;
|
||||
let mut by_name: HashMap<String, PromotionFile> = HashMap::new();
|
||||
|
||||
for b in &bucket_infos {
|
||||
let store = match self.buckets.get(&b.name) {
|
||||
Ok(s) => s,
|
||||
Err(_) => continue,
|
||||
};
|
||||
let keys = ops::list(&store, Some(&format!("{PROMOTION_PREFIX}/")))
|
||||
.await.unwrap_or_default();
|
||||
for key in keys {
|
||||
if !key.ends_with(".json") { continue; }
|
||||
let bytes = match ops::get(&store, &key).await {
|
||||
Ok(b) => b,
|
||||
Err(_) => continue,
|
||||
};
|
||||
if let Ok(f) = serde_json::from_slice::<PromotionFile>(&bytes) {
|
||||
by_name.insert(f.index_name.clone(), f);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(by_name.into_values().collect())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,312 @@
|
||||
//! Phase C: Decoupled embedding refresh.
|
||||
//!
|
||||
//! When a dataset's row-level data changes, the vector index it feeds is
|
||||
//! stale. Historically we coupled ingest and embedding — writing new rows
|
||||
//! also re-embedded every row, which is fine at 500 rows but blows up at
|
||||
//! 100K+. The llms3.com architecture calls out "asynchronous vector
|
||||
//! refresh cycles independent from transactional mutations" as the right
|
||||
//! pattern.
|
||||
//!
|
||||
//! This module implements the refresh side. The ingest path marks
|
||||
//! embeddings stale (see catalogd::registry::mark_embeddings_stale); this
|
||||
//! code clears that staleness by embedding only rows whose `doc_id` isn't
|
||||
//! already in the existing vector index.
|
||||
//!
|
||||
//! Scope — MVP:
|
||||
//! - Reads the dataset's Parquet, extracts (doc_id, text) pairs from named
|
||||
//! columns
|
||||
//! - Loads existing embeddings via EmbeddingCache
|
||||
//! - Filters to rows whose doc_id is NOT in the existing set
|
||||
//! - Chunks, embeds via Ollama, appends to the index parquet
|
||||
//! - Clears stale flag on success
|
||||
//!
|
||||
//! Not in MVP:
|
||||
//! - UPDATE semantics (same doc_id, new content) — would need content-hash
|
||||
//! comparison per row
|
||||
//! - Large-scale resilience (batching with checkpoints like the
|
||||
//! supervisor) — MVP does it inline
|
||||
//! - Lance backend — ADR-019 makes this straightforward later; MVP stays
|
||||
//! on Parquet sidecar indexes
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
|
||||
use aibridge::client::{AiClient, EmbedRequest};
|
||||
use arrow::array::{Array, Int32Array, Int64Array, StringArray};
|
||||
use catalogd::registry::Registry;
|
||||
use object_store::ObjectStore;
|
||||
|
||||
use crate::chunker::{self, TextChunk};
|
||||
use crate::embedding_cache::EmbeddingCache;
|
||||
use crate::index_registry::IndexRegistry;
|
||||
use crate::store::{self, StoredEmbedding};
|
||||
|
||||
#[derive(Debug, Clone, serde::Deserialize)]
|
||||
pub struct RefreshRequest {
|
||||
pub index_name: String,
|
||||
/// Column with the document id (row identity).
|
||||
pub id_column: String,
|
||||
/// Column with the text to embed.
|
||||
pub text_column: String,
|
||||
/// Chunk size (chars). Defaults to 500.
|
||||
#[serde(default)]
|
||||
pub chunk_size: Option<usize>,
|
||||
/// Overlap (chars). Defaults to 50.
|
||||
#[serde(default)]
|
||||
pub overlap: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct RefreshResult {
|
||||
pub index_name: String,
|
||||
pub dataset_name: String,
|
||||
pub pre_existing_docs: usize,
|
||||
pub dataset_docs: usize,
|
||||
pub new_docs_embedded: usize,
|
||||
pub new_chunks_embedded: usize,
|
||||
pub total_embeddings_after: usize,
|
||||
pub duration_secs: f32,
|
||||
pub stale_cleared: bool,
|
||||
}
|
||||
|
||||
/// Full refresh pipeline. Takes dataset_name as URL param, body has the
|
||||
/// column selectors.
|
||||
pub async fn refresh_index(
|
||||
dataset_name: &str,
|
||||
req: &RefreshRequest,
|
||||
store_: &Arc<dyn ObjectStore>,
|
||||
registry: &Registry,
|
||||
ai_client: &AiClient,
|
||||
embedding_cache: &EmbeddingCache,
|
||||
index_registry: &IndexRegistry,
|
||||
) -> Result<RefreshResult, String> {
|
||||
let t0 = std::time::Instant::now();
|
||||
|
||||
// 1. Find the dataset manifest, pull its object storage key
|
||||
let manifest = registry
|
||||
.get_by_name(dataset_name)
|
||||
.await
|
||||
.ok_or_else(|| format!("dataset not found: {dataset_name}"))?;
|
||||
|
||||
if manifest.objects.is_empty() {
|
||||
return Err(format!("dataset '{dataset_name}' has no object references"));
|
||||
}
|
||||
|
||||
// 2. Read dataset rows — extract (doc_id, text) pairs from the
|
||||
// specified columns. Multi-object datasets: concat across all.
|
||||
let mut doc_id_to_text: Vec<(String, String)> = Vec::new();
|
||||
for obj in &manifest.objects {
|
||||
let data = storaged::ops::get(store_, &obj.key).await
|
||||
.map_err(|e| format!("read {}: {e}", obj.key))?;
|
||||
let (schema, batches) = shared::arrow_helpers::parquet_to_record_batches(&data)
|
||||
.map_err(|e| format!("parse {}: {e}", obj.key))?;
|
||||
|
||||
let id_idx = schema.index_of(&req.id_column)
|
||||
.map_err(|_| format!("id column '{}' not in dataset schema", req.id_column))?;
|
||||
let text_idx = schema.index_of(&req.text_column)
|
||||
.map_err(|_| format!("text column '{}' not in dataset schema", req.text_column))?;
|
||||
|
||||
for batch in &batches {
|
||||
let text_col = batch
|
||||
.column(text_idx)
|
||||
.as_any()
|
||||
.downcast_ref::<StringArray>()
|
||||
.ok_or_else(|| format!("text column '{}' is not Utf8", req.text_column))?;
|
||||
|
||||
// Accept Utf8, Int32, or Int64 as id — that covers CSV-derived
|
||||
// and Postgres-imported schemas without forcing upstream casts.
|
||||
let id_reader: Box<dyn Fn(usize) -> Option<String>> = {
|
||||
let col = batch.column(id_idx);
|
||||
if let Some(s) = col.as_any().downcast_ref::<StringArray>() {
|
||||
let s = s.clone();
|
||||
Box::new(move |row| if s.is_null(row) { None } else { Some(s.value(row).to_string()) })
|
||||
} else if let Some(a) = col.as_any().downcast_ref::<Int32Array>() {
|
||||
let a = a.clone();
|
||||
Box::new(move |row| if a.is_null(row) { None } else { Some(a.value(row).to_string()) })
|
||||
} else if let Some(a) = col.as_any().downcast_ref::<Int64Array>() {
|
||||
let a = a.clone();
|
||||
Box::new(move |row| if a.is_null(row) { None } else { Some(a.value(row).to_string()) })
|
||||
} else {
|
||||
return Err(format!(
|
||||
"id column '{}' must be Utf8, Int32, or Int64 — got {}",
|
||||
req.id_column,
|
||||
col.data_type(),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
for row in 0..batch.num_rows() {
|
||||
if text_col.is_null(row) { continue; }
|
||||
let Some(id) = id_reader(row) else { continue; };
|
||||
let text = text_col.value(row).to_string();
|
||||
if text.trim().is_empty() { continue; }
|
||||
doc_id_to_text.push((id, text));
|
||||
}
|
||||
}
|
||||
}
|
||||
let dataset_docs = doc_id_to_text.len();
|
||||
tracing::info!("refresh '{}': dataset has {dataset_docs} rows", dataset_name);
|
||||
|
||||
// 3. Load existing embeddings (empty if no index yet)
|
||||
let existing: Vec<StoredEmbedding> = match embedding_cache.get_or_load(&req.index_name).await {
|
||||
Ok(arc) => arc.as_ref().clone(),
|
||||
Err(_) => Vec::new(), // first-time index build
|
||||
};
|
||||
let pre_existing_docs: HashSet<String> = existing
|
||||
.iter()
|
||||
.map(|e| e.doc_id.clone())
|
||||
.collect();
|
||||
let pre_existing_count = pre_existing_docs.len();
|
||||
|
||||
// 4. Delta — rows whose doc_id isn't already indexed
|
||||
let new_rows: Vec<(String, String)> = doc_id_to_text
|
||||
.into_iter()
|
||||
.filter(|(id, _)| !pre_existing_docs.contains(id))
|
||||
.collect();
|
||||
let new_docs = new_rows.len();
|
||||
|
||||
if new_docs == 0 {
|
||||
tracing::info!("refresh '{}': no new docs to embed", dataset_name);
|
||||
// Even on no-op, make sure index metadata is registered so
|
||||
// downstream discovery (profile activation, A/B search) sees
|
||||
// indexes built by earlier refresh calls that predated the
|
||||
// auto-register behavior.
|
||||
let _ = try_update_index_meta(index_registry, &req.index_name, existing.len()).await;
|
||||
registry.clear_embeddings_stale(dataset_name).await?;
|
||||
return Ok(RefreshResult {
|
||||
index_name: req.index_name.clone(),
|
||||
dataset_name: dataset_name.to_string(),
|
||||
pre_existing_docs: pre_existing_count,
|
||||
dataset_docs,
|
||||
new_docs_embedded: 0,
|
||||
new_chunks_embedded: 0,
|
||||
total_embeddings_after: existing.len(),
|
||||
duration_secs: t0.elapsed().as_secs_f32(),
|
||||
stale_cleared: true,
|
||||
});
|
||||
}
|
||||
|
||||
// 5. Chunk the new rows
|
||||
let chunk_size = req.chunk_size.unwrap_or(500);
|
||||
let overlap = req.overlap.unwrap_or(50);
|
||||
let doc_ids: Vec<String> = new_rows.iter().map(|(id, _)| id.clone()).collect();
|
||||
let texts: Vec<String> = new_rows.iter().map(|(_, t)| t.clone()).collect();
|
||||
let chunks: Vec<TextChunk> = chunker::chunk_column(
|
||||
dataset_name, &doc_ids, &texts, chunk_size, overlap,
|
||||
);
|
||||
let new_chunks = chunks.len();
|
||||
tracing::info!("refresh '{}': {} new docs -> {} chunks", dataset_name, new_docs, new_chunks);
|
||||
|
||||
// 6. Embed via Ollama (batched)
|
||||
let batch_size = 32;
|
||||
let mut all_vectors: Vec<Vec<f64>> = Vec::with_capacity(new_chunks);
|
||||
for batch in chunks.chunks(batch_size) {
|
||||
let batch_texts: Vec<String> = batch.iter().map(|c| c.text.clone()).collect();
|
||||
let resp = ai_client
|
||||
.embed(EmbedRequest { texts: batch_texts, model: None })
|
||||
.await
|
||||
.map_err(|e| format!("embed: {e}"))?;
|
||||
all_vectors.extend(resp.embeddings);
|
||||
}
|
||||
|
||||
// 7. Combine existing + new and write back as a single parquet
|
||||
// (MVP — append-as-rewrite. ADR-019 points to Lance for true native
|
||||
// append; for the Parquet sidecar we rewrite at refresh time.)
|
||||
let mut new_stored: Vec<StoredEmbedding> = existing.clone();
|
||||
for (chunk, vector) in chunks.into_iter().zip(all_vectors.iter()) {
|
||||
new_stored.push(StoredEmbedding {
|
||||
source: chunk.source,
|
||||
doc_id: chunk.doc_id,
|
||||
chunk_idx: chunk.chunk_idx,
|
||||
chunk_text: chunk.text,
|
||||
vector: vector.iter().map(|&x| x as f32).collect(),
|
||||
});
|
||||
}
|
||||
|
||||
// We have to reconstruct chunks + vectors for store_embeddings — but
|
||||
// store_embeddings takes &[TextChunk] and vectors separately. Convert
|
||||
// the combined StoredEmbedding back to that shape.
|
||||
let combined_chunks: Vec<TextChunk> = new_stored
|
||||
.iter()
|
||||
.map(|e| TextChunk {
|
||||
source: e.source.clone(),
|
||||
doc_id: e.doc_id.clone(),
|
||||
chunk_idx: e.chunk_idx,
|
||||
text: e.chunk_text.clone(),
|
||||
})
|
||||
.collect();
|
||||
let combined_vectors: Vec<Vec<f64>> = new_stored
|
||||
.iter()
|
||||
.map(|e| e.vector.iter().map(|&f| f as f64).collect())
|
||||
.collect();
|
||||
|
||||
let _key = store::store_embeddings(
|
||||
store_, &req.index_name, &combined_chunks, &combined_vectors,
|
||||
).await?;
|
||||
|
||||
// 8. Evict embedding cache — next read will pick up the new file
|
||||
let _ = embedding_cache.evict(&req.index_name).await;
|
||||
|
||||
// 9. Update index registry metadata (row/chunk counts; others unchanged
|
||||
// so we don't disturb the existing entry if present)
|
||||
let _ = try_update_index_meta(index_registry, &req.index_name, new_stored.len()).await;
|
||||
|
||||
// 10. Clear stale flag
|
||||
registry.clear_embeddings_stale(dataset_name).await?;
|
||||
|
||||
let total = new_stored.len();
|
||||
Ok(RefreshResult {
|
||||
index_name: req.index_name.clone(),
|
||||
dataset_name: dataset_name.to_string(),
|
||||
pre_existing_docs: pre_existing_count,
|
||||
dataset_docs,
|
||||
new_docs_embedded: new_docs,
|
||||
new_chunks_embedded: new_chunks,
|
||||
total_embeddings_after: total,
|
||||
duration_secs: t0.elapsed().as_secs_f32(),
|
||||
stale_cleared: true,
|
||||
})
|
||||
}
|
||||
|
||||
/// Best-effort refresh of index registry metadata.
|
||||
/// - If the index already exists, bump the chunk_count.
|
||||
/// - If it's brand new (first-time build via refresh), REGISTER it so
|
||||
/// downstream discovery (profile activation, A/B search) can find it.
|
||||
async fn try_update_index_meta(
|
||||
index_registry: &IndexRegistry,
|
||||
index_name: &str,
|
||||
chunk_count: usize,
|
||||
) -> Result<(), String> {
|
||||
if let Some(mut meta) = index_registry.get(index_name).await {
|
||||
meta.chunk_count = chunk_count;
|
||||
return index_registry.register(meta).await;
|
||||
}
|
||||
// First-time registration — infer reasonable defaults. Convention:
|
||||
// index name `{source}_v1` or `{source}_vN` implies a source dataset
|
||||
// named by stripping the `_vN` suffix. Otherwise use the index name
|
||||
// as the source and let the caller patch later.
|
||||
let source = match index_name.rsplit_once('_') {
|
||||
Some((base, suffix)) if suffix.starts_with('v') && suffix[1..].chars().all(|c| c.is_ascii_digit()) => {
|
||||
base.to_string()
|
||||
}
|
||||
_ => index_name.to_string(),
|
||||
};
|
||||
let meta = crate::index_registry::IndexMeta {
|
||||
index_name: index_name.to_string(),
|
||||
source,
|
||||
model_name: "nomic-embed-text".to_string(),
|
||||
model_version: "latest".to_string(),
|
||||
dimensions: 768,
|
||||
chunk_count,
|
||||
doc_count: chunk_count,
|
||||
chunk_size: 500,
|
||||
overlap: 50,
|
||||
storage_key: format!("vectors/{index_name}.parquet"),
|
||||
created_at: chrono::Utc::now(),
|
||||
build_time_secs: 0.0,
|
||||
chunks_per_sec: 0.0,
|
||||
bucket: "primary".to_string(),
|
||||
vector_backend: shared::types::VectorBackend::Parquet,
|
||||
};
|
||||
index_registry.register(meta).await
|
||||
}
|
||||
+1266
-10
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,240 @@
|
||||
/// Trial journal for HNSW parameter tuning.
|
||||
///
|
||||
/// Every HNSW build+eval is recorded as a Trial. The journal is append-only
|
||||
/// and stored under `_hnsw_trials/{index_name}/` as batched JSONL files —
|
||||
/// an AI agent iterating on configs reads prior trials to decide what to
|
||||
/// try next, and writes a new trial on each attempt.
|
||||
///
|
||||
/// Storage uses the shared `storaged::append_log::AppendLog` so appends are
|
||||
/// write-once (new file per batch) rather than rewriting a single growing
|
||||
/// JSONL on every event. See `append_log.rs` for the full rationale.
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use object_store::ObjectStore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use storaged::append_log::AppendLog;
|
||||
use storaged::registry::BucketRegistry;
|
||||
|
||||
use crate::index_registry::IndexRegistry;
|
||||
|
||||
/// HNSW build/search parameters the agent can tune.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct HnswConfig {
|
||||
pub ef_construction: usize,
|
||||
pub ef_search: usize,
|
||||
#[serde(default)]
|
||||
pub seed: Option<u64>,
|
||||
}
|
||||
|
||||
impl Default for HnswConfig {
|
||||
/// Production default, locked in 2026-04-16 based on trial grid against
|
||||
/// resumes_100k_v2 (100K vectors, 20 queries, recall@10):
|
||||
/// ec=20 es=30 → recall 0.960, p50 509us, build 8s
|
||||
/// ec=80 es=30 → recall 1.000, p50 873us, build 230s ← sweet spot
|
||||
/// ec=200 es=30 → recall 1.000, p50 874us, build 106s (no recall gain)
|
||||
///
|
||||
/// `ec=80` is the smallest value that reaches 100% recall. Higher values
|
||||
/// waste build time. `es=30` gives faster search than `es=100` with no
|
||||
/// recall penalty at this scale.
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
ef_construction: 80,
|
||||
ef_search: 30,
|
||||
seed: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Metrics collected on every trial. All latencies in microseconds.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TrialMetrics {
|
||||
pub build_time_secs: f32,
|
||||
pub search_latency_p50_us: f32,
|
||||
pub search_latency_p95_us: f32,
|
||||
pub search_latency_p99_us: f32,
|
||||
pub recall_at_k: f32,
|
||||
pub memory_bytes: u64,
|
||||
pub vectors: usize,
|
||||
pub eval_queries: usize,
|
||||
/// Brute-force latency for comparison — how much speedup did HNSW buy us?
|
||||
pub brute_force_latency_us: f32,
|
||||
}
|
||||
|
||||
/// A single tuning attempt.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Trial {
|
||||
pub id: String,
|
||||
pub index_name: String,
|
||||
pub eval_set: String,
|
||||
pub config: HnswConfig,
|
||||
pub metrics: TrialMetrics,
|
||||
pub created_at: DateTime<Utc>,
|
||||
/// Free-form note — the agent can record why it tried this config.
|
||||
#[serde(default)]
|
||||
pub note: Option<String>,
|
||||
}
|
||||
|
||||
impl Trial {
|
||||
pub fn new_id() -> String {
|
||||
format!(
|
||||
"trial-{}-{}",
|
||||
Utc::now().timestamp_millis(),
|
||||
&uuid::Uuid::new_v4().to_string()[..8]
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-index append log, lazy-created on first write.
|
||||
///
|
||||
/// Federation layer 2: the journal resolves each index's bucket from the
|
||||
/// index registry and writes its JSONL batches to THAT bucket, not
|
||||
/// primary. Back-compat is preserved by `IndexMeta::bucket` defaulting
|
||||
/// to "primary" for pre-federation indexes. Indexes the registry has
|
||||
/// never heard of (edge case — trials run before first register) fall
|
||||
/// through to primary as well.
|
||||
#[derive(Clone)]
|
||||
pub struct TrialJournal {
|
||||
buckets: Arc<BucketRegistry>,
|
||||
index_registry: IndexRegistry,
|
||||
/// Cache per (bucket, index) AppendLog so the in-memory buffer persists
|
||||
/// across calls. Keyed by `(bucket, index_name)` so moving an index
|
||||
/// between buckets is clean — the old journal stays intact.
|
||||
logs: Arc<RwLock<HashMap<(String, String), Arc<AppendLog>>>>,
|
||||
}
|
||||
|
||||
impl TrialJournal {
|
||||
pub fn new(buckets: Arc<BucketRegistry>, index_registry: IndexRegistry) -> Self {
|
||||
Self {
|
||||
buckets,
|
||||
index_registry,
|
||||
logs: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
fn prefix(index_name: &str) -> String {
|
||||
format!("_hnsw_trials/{}", index_name)
|
||||
}
|
||||
|
||||
/// Resolve which bucket holds this index's trial artifacts.
|
||||
/// Falls back to primary for indexes without recorded metadata.
|
||||
async fn bucket_for(&self, index_name: &str) -> String {
|
||||
self.index_registry
|
||||
.get(index_name)
|
||||
.await
|
||||
.map(|m| m.bucket)
|
||||
.unwrap_or_else(|| "primary".to_string())
|
||||
}
|
||||
|
||||
async fn log_for(&self, index_name: &str) -> Result<Arc<AppendLog>, String> {
|
||||
let bucket = self.bucket_for(index_name).await;
|
||||
let key = (bucket.clone(), index_name.to_string());
|
||||
|
||||
if let Some(log) = self.logs.read().await.get(&key) {
|
||||
return Ok(log.clone());
|
||||
}
|
||||
let mut guard = self.logs.write().await;
|
||||
if let Some(log) = guard.get(&key) {
|
||||
return Ok(log.clone());
|
||||
}
|
||||
let store = self.buckets.get(&bucket)?;
|
||||
// Trials arrive one at a time during human/agent iteration — a low
|
||||
// threshold gives "hit /trials and see my latest attempt" immediacy
|
||||
// without creating one file per event.
|
||||
let log = Arc::new(
|
||||
AppendLog::new(store, Self::prefix(index_name))
|
||||
.with_flush_threshold(4),
|
||||
);
|
||||
guard.insert(key, log.clone());
|
||||
Ok(log)
|
||||
}
|
||||
|
||||
/// Append a trial record. In-memory buffered; persisted in batches.
|
||||
pub async fn append(&self, trial: &Trial) -> Result<(), String> {
|
||||
let line = serde_json::to_vec(trial).map_err(|e| e.to_string())?;
|
||||
let log = self.log_for(&trial.index_name).await?;
|
||||
log.append(line).await
|
||||
}
|
||||
|
||||
/// Read all trials for an index (flushed batches + unflushed buffer).
|
||||
pub async fn list(&self, index_name: &str) -> Result<Vec<Trial>, String> {
|
||||
let log = self.log_for(index_name).await?;
|
||||
let lines = log.read_all().await?;
|
||||
let mut trials = Vec::with_capacity(lines.len());
|
||||
for line in lines {
|
||||
match serde_json::from_slice::<Trial>(&line) {
|
||||
Ok(t) => trials.push(t),
|
||||
Err(e) => tracing::warn!("trial journal: skip malformed line: {e}"),
|
||||
}
|
||||
}
|
||||
Ok(trials)
|
||||
}
|
||||
|
||||
/// Explicit flush for callers that want write-through semantics
|
||||
/// (e.g. an agent that wants to commit a trial before querying stats).
|
||||
pub async fn flush(&self, index_name: &str) -> Result<(), String> {
|
||||
let log = self.log_for(index_name).await?;
|
||||
log.flush().await
|
||||
}
|
||||
|
||||
/// Compact all batch files for an index into one.
|
||||
pub async fn compact(&self, index_name: &str) -> Result<storaged::append_log::CompactStats, String> {
|
||||
let log = self.log_for(index_name).await?;
|
||||
log.compact().await
|
||||
}
|
||||
|
||||
/// Current champion for an index by the named metric.
|
||||
/// Valid metrics: `recall`, `latency`, `pareto`.
|
||||
///
|
||||
/// The `pareto` strategy is a placeholder — J should tune the scoring
|
||||
/// function to match what matters in production. Right now it's a simple
|
||||
/// weighted sum.
|
||||
pub async fn best(
|
||||
&self,
|
||||
index_name: &str,
|
||||
metric: &str,
|
||||
) -> Result<Option<Trial>, String> {
|
||||
let trials = self.list(index_name).await?;
|
||||
if trials.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let best = match metric {
|
||||
"recall" => trials
|
||||
.into_iter()
|
||||
.max_by(|a, b| {
|
||||
a.metrics
|
||||
.recall_at_k
|
||||
.partial_cmp(&b.metrics.recall_at_k)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
})
|
||||
.unwrap(),
|
||||
"latency" => trials
|
||||
.into_iter()
|
||||
.min_by(|a, b| {
|
||||
a.metrics
|
||||
.search_latency_p95_us
|
||||
.partial_cmp(&b.metrics.search_latency_p95_us)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
})
|
||||
.unwrap(),
|
||||
"pareto" | _ => trials
|
||||
.into_iter()
|
||||
.max_by(|a, b| pareto_score(a).partial_cmp(&pareto_score(b)).unwrap())
|
||||
.unwrap(),
|
||||
};
|
||||
Ok(Some(best))
|
||||
}
|
||||
}
|
||||
|
||||
/// Simple Pareto-style score: reward recall, penalize p95 latency.
|
||||
/// Tunable — J should swap this in production to match what matters.
|
||||
fn pareto_score(t: &Trial) -> f32 {
|
||||
// Recall is [0, 1]. Latency is us — assume 100us baseline.
|
||||
let recall = t.metrics.recall_at_k;
|
||||
let latency_penalty = (t.metrics.search_latency_p95_us / 1000.0).min(1.0); // cap at 1ms
|
||||
recall - 0.2 * latency_penalty
|
||||
}
|
||||
@@ -0,0 +1,322 @@
|
||||
# ADR-017: Federated Multi-Bucket Storage
|
||||
|
||||
**Status:** Accepted — 2026-04-16
|
||||
**Owner:** J
|
||||
**Implements:** Phase 15 horizon item "Federated multi-bucket query"
|
||||
**Depends on:** ADR-001 (object storage as source of truth), ADR-009 (delta files)
|
||||
|
||||
---
|
||||
|
||||
## Problem
|
||||
|
||||
Today every dataset lives in exactly one object storage backend (`/home/profit/lakehouse/data` on local disk). Three scenarios break that assumption:
|
||||
|
||||
1. **Multi-tenant hosting** — Client A's staffing data must stay in Client A's S3 bucket. Client B's in Client B's. The lakehouse is the compute plane, their bucket is the data plane.
|
||||
2. **Data residency** — European clients require their data never leaves an EU-region bucket. Same lakehouse instance, different physical storage.
|
||||
3. **Cross-dataset analytics** — A user asks "average placement margin across all clients we manage." The query must scan Parquet across N buckets as a single logical table.
|
||||
|
||||
None of this is possible with a single `object_store` instance. `DataFusion` *can* register multiple object stores under different URL schemes, but nothing in the current stack surfaces that capability.
|
||||
|
||||
## Non-Goals (for the MVP)
|
||||
|
||||
- **Not** building cross-bucket JOIN optimization. DataFusion pushes predicates down per-partition already; that's sufficient.
|
||||
- **Not** handling bucket-level auth policies (who-can-read-which-bucket). That's Phase 13's job.
|
||||
- **Not** supporting heterogeneous backends in a single logical dataset (e.g. half in S3, half local). One dataset = one bucket.
|
||||
- **Not** automatic replication across buckets. Each bucket stands alone.
|
||||
- **Not** scheduled sync/migration. Manual today.
|
||||
|
||||
---
|
||||
|
||||
## Design
|
||||
|
||||
### Core invariant
|
||||
|
||||
Every `ObjectRef` belongs to exactly one named bucket. The catalog is the single authority over which bucket holds which dataset. Queries transparently span buckets by virtue of DataFusion's multi-store capability — the catalog tells the query engine where each file lives.
|
||||
|
||||
### Naming
|
||||
|
||||
Three classes of bucket:
|
||||
|
||||
- **`primary`** — system-wide default. Shared catalog, reference datasets, and vectors that aren't tenant-specific. Always present, always the fallback. Corresponds to today's `./data` directory.
|
||||
- **`profile:{user_id}`** — per-user workspace bucket. Each user's personal data, their vector indexes, their workspace state. Provisioned on first use. Pre-loaded ("hot loaded") when the user activates their profile.
|
||||
- **`{named}`** — tenant/client buckets. Explicit, configured up front. E.g. `client_a`, `client_eu`.
|
||||
|
||||
Conventions:
|
||||
- **Bucket URL** = `{bucket_id}://{key}` internally (e.g. `primary://datasets/candidates.parquet`, `profile:daisy://vectors/notes.parquet`)
|
||||
- DataFusion sees each bucket as a separate registered object store; `ListingTable` paths include the bucket scheme
|
||||
- Profile buckets use the `profile:` prefix as both a namespace marker and a DataFusion URL scheme
|
||||
|
||||
### Config shape — `lakehouse.toml`
|
||||
|
||||
```toml
|
||||
[storage]
|
||||
# Backward compat — the existing [storage] block still creates a "primary"
|
||||
# bucket if no [[storage.buckets]] entries are defined.
|
||||
root = "./data"
|
||||
# Where profile buckets are rooted when not otherwise configured.
|
||||
# A profile:{user} bucket resolves to {profile_root}/{user}/ unless the user
|
||||
# config overrides.
|
||||
profile_root = "./data/_profiles"
|
||||
|
||||
[[storage.buckets]]
|
||||
name = "primary"
|
||||
backend = "local"
|
||||
root = "./data"
|
||||
|
||||
rescue_bucket = "rescue" # single shared fallback for failed reads (see §Failure mode)
|
||||
|
||||
[[storage.buckets]]
|
||||
name = "rescue"
|
||||
backend = "local"
|
||||
root = "./data/_rescue"
|
||||
|
||||
[[storage.buckets]]
|
||||
name = "client_a"
|
||||
backend = "s3"
|
||||
bucket = "client-a-lakehouse"
|
||||
region = "us-east-1"
|
||||
endpoint = "https://s3.amazonaws.com"
|
||||
secret_ref = "client_a_aws" # NOT the literal secret — a handle
|
||||
|
||||
[[storage.buckets]]
|
||||
name = "client_eu"
|
||||
backend = "s3"
|
||||
bucket = "client-eu-lakehouse"
|
||||
region = "eu-west-1"
|
||||
secret_ref = "client_eu_aws"
|
||||
```
|
||||
|
||||
Rules:
|
||||
- Credentials **never appear in this file**. `secret_ref` is a handle resolved through the secrets layer (see §Secrets).
|
||||
- If `[[storage.buckets]]` is absent, the existing `[storage]` block creates a single `"primary"` bucket. Zero-config upgrade.
|
||||
- Adding a bucket requires a restart (MVP); runtime addition via API is a polish item.
|
||||
|
||||
### Secrets
|
||||
|
||||
Credentials are never in `lakehouse.toml`. A pluggable `SecretsProvider` trait resolves `secret_ref` handles to credential structs at startup:
|
||||
|
||||
```rust
|
||||
pub trait SecretsProvider: Send + Sync {
|
||||
async fn resolve(&self, handle: &str) -> Result<BucketCredentials, String>;
|
||||
}
|
||||
```
|
||||
|
||||
MVP ships **one** implementation — `FileSecretsProvider` reading `/etc/lakehouse/secrets.toml` (or a path given by `LAKEHOUSE_SECRETS` env var):
|
||||
|
||||
```toml
|
||||
# /etc/lakehouse/secrets.toml — root:root, mode 0600, NEVER in git
|
||||
[client_a_aws]
|
||||
access_key = "AKIA..."
|
||||
secret_key = "wJalrXUt..."
|
||||
|
||||
[client_eu_aws]
|
||||
access_key = "AKIA..."
|
||||
secret_key = "..."
|
||||
```
|
||||
|
||||
Future providers plug into the same trait without touching core code:
|
||||
- `VaultSecretsProvider` — HashiCorp Vault
|
||||
- `SopsSecretsProvider` — age/gpg-encrypted files in git
|
||||
- `KeyringSecretsProvider` — OS-level keyring
|
||||
|
||||
Startup fails fast if a `secret_ref` can't resolve — no silent fallback to anonymous S3.
|
||||
|
||||
### Schema changes
|
||||
|
||||
**`ObjectRef`** already has a `bucket: String` field — currently it carries the S3 bucket name or `"local"` inconsistently. Repurpose it as the **catalog bucket name**:
|
||||
|
||||
```rust
|
||||
pub struct ObjectRef {
|
||||
pub bucket: String, // NOW: catalog bucket name, e.g. "primary" or "client_a"
|
||||
pub key: String, // object key within that bucket
|
||||
pub size_bytes: u64,
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
```
|
||||
|
||||
Migration: a `resync-missing`-style one-shot sets `bucket = "primary"` on every existing ObjectRef whose value is empty or ambiguous.
|
||||
|
||||
**`DatasetManifest`** — no new top-level field. `objects: Vec<ObjectRef>` is already where bucket info lives. Per the design invariant (one dataset, one bucket), all ObjectRefs inside a manifest share a bucket, so we can add a convenience accessor `manifest.bucket_id()` → `objects[0].bucket.clone()`.
|
||||
|
||||
### Storage layer changes
|
||||
|
||||
**`storaged::BucketRegistry`** — new struct, replaces the single `Arc<dyn ObjectStore>` currently threaded through gateway.
|
||||
|
||||
```rust
|
||||
pub struct BucketRegistry {
|
||||
buckets: HashMap<String, Arc<dyn ObjectStore>>,
|
||||
default: String, // "primary"
|
||||
}
|
||||
|
||||
impl BucketRegistry {
|
||||
pub fn from_config(cfg: &StorageConfig) -> Result<Self, String>;
|
||||
pub fn get(&self, bucket_id: &str) -> Result<&Arc<dyn ObjectStore>, String>;
|
||||
pub fn default_store(&self) -> &Arc<dyn ObjectStore>;
|
||||
pub fn list_buckets(&self) -> Vec<BucketInfo>;
|
||||
}
|
||||
```
|
||||
|
||||
**`storaged::ops`** gets a bucket parameter on every call:
|
||||
```rust
|
||||
pub async fn put(registry: &BucketRegistry, bucket: &str, key: &str, data: Bytes) -> Result<(), String>;
|
||||
pub async fn get(registry: &BucketRegistry, bucket: &str, key: &str) -> Result<Bytes, String>;
|
||||
// etc.
|
||||
```
|
||||
|
||||
Every existing call site passes `"primary"` initially; specific call sites (ingestd, vectord) learn to route.
|
||||
|
||||
### Query layer changes
|
||||
|
||||
**`queryd`** — on startup, register every bucket as a separate `object_store` URL with DataFusion:
|
||||
|
||||
```rust
|
||||
for (name, store) in bucket_registry.iter() {
|
||||
let url = format!("lakehouse-{}://", name).parse::<Url>()?;
|
||||
session_ctx.register_object_store(&url, store.clone());
|
||||
}
|
||||
```
|
||||
|
||||
**`ListingTable`** URLs reference the bucket:
|
||||
```rust
|
||||
let url = format!("lakehouse-{}://{}", obj.bucket, obj.key);
|
||||
```
|
||||
|
||||
DataFusion handles cross-bucket scans natively — each partition gets routed to the correct store.
|
||||
|
||||
### Gateway API additions
|
||||
|
||||
| Endpoint | Purpose |
|
||||
|---|---|
|
||||
| `GET /storage/buckets` | List all configured buckets + their backend + reachability status |
|
||||
| `GET /storage/errors` | Recent bucket operation failures (filtered by bucket/time) |
|
||||
| `GET /storage/health` | Summary: which buckets have errored in the last 5 minutes |
|
||||
| `POST /ingest/file` + `X-Lakehouse-Bucket: client_a` | Header selects target bucket; absent → `primary` |
|
||||
| `POST /profile/{user}/activate` | Pre-load user's profile bucket: embedding caches warm, HNSW indexes rebuilt, workspace state hydrated |
|
||||
| `POST /profile/{user}/deactivate` | Evict profile bucket's cached state |
|
||||
| `POST /catalog/datasets/by-name/{name}/relocate` | Move a dataset to another bucket (polish) |
|
||||
|
||||
**Routing rule for all existing endpoints:**
|
||||
1. `X-Lakehouse-Bucket: {name}` header → use that bucket
|
||||
2. No header → `primary`
|
||||
3. Unknown bucket name → 404
|
||||
|
||||
Catalog and query endpoints span all buckets by default — a query sees all buckets the user has access to unless explicitly scoped.
|
||||
|
||||
### Profile hot load
|
||||
|
||||
When `POST /profile/{user}/activate` fires:
|
||||
|
||||
1. Resolve `profile:{user}` bucket — create if missing (first-time activation)
|
||||
2. Scan the catalog for all datasets whose `owner == user` or which live in that bucket
|
||||
3. For each vector index belonging to the profile: pre-load embeddings into `EmbeddingCache`, pre-build HNSW with the locked-in default config (ec=80 es=30)
|
||||
4. Hydrate any saved workspace state (see Phase 8.5) into memory
|
||||
5. Return a manifest summary — what's hot, what wasn't found, total memory used
|
||||
|
||||
Design contract: `/profile/{user}/activate` should be idempotent. Calling it twice is cheap; the second call is a no-op because the cache is already warm.
|
||||
|
||||
### Failure mode — rescue bucket + loud errors
|
||||
|
||||
The design goal is **failures stay findable**. When something breaks, you should know within seconds which bucket, which key, which operation, and what happened — without grepping logs.
|
||||
|
||||
**Writes**: always go to the target bucket. If the target bucket is unreachable, the request fails hard with 503 and the failure is recorded in the error journal. No queueing, no silent fallback — writes that vanish silently are the worst possible failure mode.
|
||||
|
||||
**Reads**:
|
||||
- First attempt: target bucket.
|
||||
- On failure: fall through to the single configured `rescue_bucket`.
|
||||
- If the key exists in rescue, serve it with `X-Lakehouse-Rescue-Used: true` and `X-Lakehouse-Original-Bucket: {target}` response headers so the caller knows fallback happened.
|
||||
- If the key isn't in rescue either: 404 with both buckets listed in the error body.
|
||||
- Every rescue fallback is appended to the **error journal** (see below).
|
||||
|
||||
**Rescue bucket** is a single, well-known bucket configured at the storage level. Typical setup: a periodic sync from `primary` → `rescue` so that when `primary` has an outage, recent data is still reachable. Sync is outside this ADR's scope — operator's job.
|
||||
|
||||
**Rest of the system keeps working**: a failed bucket never takes down the gateway. Other buckets remain queryable. Cross-bucket queries scan what they can and surface per-partition failures (which bucket, which key) in the error body.
|
||||
|
||||
### Error journal — "find errors ez"
|
||||
|
||||
A simple append-only log of every bucket operation that failed or fell through. Stored at `primary://_errors/bucket_errors.jsonl` (yes, primary — because we need it findable even if other buckets are down).
|
||||
|
||||
```json
|
||||
{"ts":"2026-04-16T10:30:15Z","op":"read","target":"client_a","key":"datasets/candidates.parquet","error":"connection refused","rescued":true,"rescue_key_found":true}
|
||||
{"ts":"2026-04-16T10:30:18Z","op":"write","target":"client_a","key":"datasets/new.parquet","error":"connection refused","rescued":false}
|
||||
```
|
||||
|
||||
Exposed via:
|
||||
- `GET /storage/errors?limit=50` — recent errors
|
||||
- `GET /storage/errors?bucket=client_a` — filter by bucket
|
||||
- `GET /storage/errors?since=2026-04-16T10:00` — filter by time
|
||||
- `GET /storage/health` — summary: which buckets have errored in the last 5 minutes
|
||||
|
||||
The journal doesn't replace logs (tracing still gets everything), but it gives you a **single authoritative place** to answer the question "has anything been failing?" without tailing systemd journals. If the journal is empty, nothing is broken.
|
||||
|
||||
### Vectors are bucket-scoped
|
||||
|
||||
Vector indexes live in the same bucket as their source data:
|
||||
|
||||
- Data in `primary` → vectors in `primary://vectors/{index_name}.parquet`
|
||||
- Data in `profile:daisy` → vectors in `profile:daisy://vectors/{index_name}.parquet`
|
||||
- Data in `client_a` → vectors in `client_a://vectors/{index_name}.parquet`
|
||||
|
||||
The existing `_hnsw_trials/` and `_hnsw_evals/` prefixes also live per-bucket. The trial journal for a profile's index stays inside that profile's bucket — so tenants don't see each other's tuning history.
|
||||
|
||||
**Profile hot load consequence:** activating a profile pre-loads ALL vector indexes in that profile's bucket into `EmbeddingCache` + builds HNSW with the default config. First query after activation is warm.
|
||||
|
||||
---
|
||||
|
||||
## Implementation phases
|
||||
|
||||
### MVP (this is what I'd build next session)
|
||||
|
||||
1. **`shared::config`** — extend `StorageConfig` with `buckets: Vec<BucketConfig>` and `profile_root`, with backward-compat fallback from single `[storage]` block
|
||||
2. **`shared::secrets`** — `SecretsProvider` trait + `FileSecretsProvider` impl reading `/etc/lakehouse/secrets.toml`
|
||||
3. **`storaged::BucketRegistry`** — multi-backend registry; mirrors-of tracking; lazy profile bucket creation
|
||||
4. **`storaged::ops`** — add `bucket: &str` param to every call; read path falls back to mirror on unreachable; write path hard-fails
|
||||
5. **`gateway`** — `X-Lakehouse-Bucket` header middleware → routing decision per request
|
||||
6. **`gateway`** — `GET /storage/buckets` + `POST /profile/{user}/activate|deactivate`
|
||||
7. **`catalogd`** — one-shot migration: every ObjectRef without an explicit bucket gets `"primary"`
|
||||
8. **`queryd::session`** — register every bucket as a DataFusion ObjectStore under its own URL scheme
|
||||
9. **`vectord`** — every vector/trial/eval storage path becomes `{bucket}://vectors/...`; profile activate pre-loads HNSW
|
||||
10. **`storaged::error_journal`** — append-only JSONL writer at `primary://_errors/bucket_errors.jsonl`, hooked into every ops call-site
|
||||
11. **Test harness** — configure `primary` + `profile:testuser` + a mock `client_a` pointing at a 2nd local dir + `rescue` as a 3rd local dir; verify cross-bucket join, profile activate, rescue fallback, error journal visibility
|
||||
|
||||
**Success gate #1 — cross-bucket query:** `SELECT SUM(hours_regular) FROM primary.timesheets t1 JOIN "profile:testuser".timesheets t2 USING (placement_id)` returns a sensible number.
|
||||
|
||||
**Success gate #2 — profile hot load:** `POST /profile/daisy/activate` pre-builds HNSW over her personal vector index; next `/search` call against that index is <1ms cold.
|
||||
|
||||
**Success gate #3 — rescue fallback + error visibility:** rename `client_a`'s directory to simulate outage. GET for a key returns the copy from `rescue` with `X-Lakehouse-Rescue-Used: true` + `X-Lakehouse-Original-Bucket: client_a` headers. PUT returns 503. `GET /storage/errors` shows both events with full context. `GET /storage/health` flags `client_a` as unhealthy.
|
||||
|
||||
### Polish (follow-up)
|
||||
|
||||
8. **Bucket health check** — `GET /storage/buckets` shows reachable/unreachable
|
||||
9. **Bucket relocation** — `POST /catalog/datasets/{name}/relocate` (copy to new bucket, update manifest, delete from old)
|
||||
10. **Runtime bucket addition** — `POST /storage/buckets` adds a bucket without restart (stores in a separate `buckets.toml`)
|
||||
11. **S3 streaming optimization** — for very large Parquet files, stream the footer instead of loading the whole blob
|
||||
|
||||
---
|
||||
|
||||
## Decisions (2026-04-16)
|
||||
|
||||
1. **Credentials** — Pluggable `SecretsProvider` trait. MVP ships `FileSecretsProvider` at `/etc/lakehouse/secrets.toml` (root:root, 0600). Future providers (Vault, SOPS, OS keyring) plug in without touching core.
|
||||
2. **Primary default + profile bucket model** — `primary` always exists as system-wide fallback. Per-user `profile:{user}` buckets are provisioned on first activate. Named tenant buckets configured explicitly.
|
||||
3. **Ingest routing** — `X-Lakehouse-Bucket` header. Absent → `primary`. Unknown → 404.
|
||||
4. **Failure** — Writes: hard fail on target, logged to error journal. Reads: fall through to a single shared `rescue_bucket`; response headers `X-Lakehouse-Rescue-Used` + `X-Lakehouse-Original-Bucket` make the fallback visible. Every failure appended to an error journal at `primary://_errors/bucket_errors.jsonl`, queryable via `GET /storage/errors`. Failed buckets never take down unrelated paths.
|
||||
5. **Vectors in their bucket** — Vectors, trial journals, and eval sets all live in the same bucket as their source data. Profile activate pre-loads the full vector stack for that profile using the locked-in HNSW default (ec=80 es=30).
|
||||
|
||||
---
|
||||
|
||||
## Risks
|
||||
|
||||
| Risk | Severity | Mitigation |
|
||||
|---|---|---|
|
||||
| DataFusion multi-store predicate pushdown lands on cross-bucket JOINs | **Medium** | Benchmark early. If bad, mark cross-bucket JOINs as a known sharp edge. |
|
||||
| S3 credentials accidentally logged | **High** | Never stringify credential fields. Add a test that redacts `secret_key` in every code path. |
|
||||
| Existing ObjectRefs have inconsistent `bucket` values (some `"local"`, some empty) | **Low** | One-shot migration normalizes to `"primary"`. |
|
||||
| Bucket config drift between restarts | **Medium** | Config-driven means restart = canonical. A runtime API would need careful state reconciliation; defer. |
|
||||
|
||||
---
|
||||
|
||||
## Decision log
|
||||
|
||||
- **1 dataset = 1 bucket** (rather than datasets spanning buckets) — vastly simpler, matches the "tenant isolation" real use case. Revisit only if someone has an actual multi-bucket dataset requirement.
|
||||
- **Config-first, runtime-later** — YAML/TOML for the MVP; runtime bucket add/remove is polish. Deployment restart is cheap in this stack.
|
||||
- **`bucket: String` repurposed, no new field** — avoids a schema migration on every manifest.
|
||||
- **DataFusion's multi-store primitive handles cross-bucket queries** — don't reinvent. We're wiring it up, not building it.
|
||||
@@ -0,0 +1,105 @@
|
||||
# ADR-019: Vector Storage — Parquet+HNSW stays, Lance joins as second tier
|
||||
|
||||
**Status:** Accepted — 2026-04-16
|
||||
**Implements:** Phase 18 from PRD (Lance evaluation)
|
||||
**Supersedes:** nothing (augments ADR-008)
|
||||
**Owner:** J
|
||||
|
||||
---
|
||||
|
||||
## Context
|
||||
|
||||
Phase 18 of the PRD committed to settling "Parquet+sidecar vs Lance" with measurements, not vibes. This ADR records the benchmark outcome and the resulting architectural direction.
|
||||
|
||||
Input data: `data/vectors/resumes_100k_v2.parquet` — 100,000 × 768d embeddings, the same index we tuned HNSW against in Phase 15.
|
||||
|
||||
Benchmark harness: `crates/lance-bench/src/main.rs` — standalone binary, deliberately not integrated into the workspace's common deps to avoid forcing DataFusion/Arrow upgrades on the rest of the stack until we'd decided.
|
||||
|
||||
## The scorecard
|
||||
|
||||
All numbers measured on the same 128GB server, same 100K × 768d index, release build:
|
||||
|
||||
| Dimension | Parquet + HNSW (current) | Lance 4.0 IVF_PQ (candidate) | Winner |
|
||||
|---|---|---|---|
|
||||
| Cold load | 0.17s | 0.13s | Lance, 1.27× — *does not clear 2× decision threshold* |
|
||||
| Disk size (data only) | 330.3 MB | 330.4 MB | Tie |
|
||||
| Index on-disk footprint | 0 (HNSW is RAM-only) | 7.4 MB | Lance |
|
||||
| Index build time | 230s (ec=80 es=30) | 16s | **Lance, 14× faster** |
|
||||
| Search p50 | 873us (recall@10 = 1.00) | 2229us (recall unmeasured, likely 0.85-0.95) | **Parquet+HNSW, 2.55× faster** |
|
||||
| Search p95 | 1413us | 4998us | **Parquet+HNSW, 3.54× faster** |
|
||||
| Speedup vs brute force (p50) | 50.4× | 19.7× | Parquet+HNSW |
|
||||
| Random row access (fetch by id) | ~35ms (full-file scan) | 311us | **Lance, 112× faster** |
|
||||
| Append 10K rows | Full-file rewrite (~330MB + re-embed + re-index) | 0.08s, +31MB delta | **Lance, structurally different** |
|
||||
|
||||
## Applying the decision rules from EXECUTION_PLAN.md
|
||||
|
||||
Original rules:
|
||||
- *Lance wins cold-load by ≥2× AND matches search latency → migrate*
|
||||
- *Within 50% across board → stay Parquet, document ceiling*
|
||||
- *Lance loses → close the door*
|
||||
|
||||
Strict reading: cold-load is **1.27×, not ≥2×**. Search latency is **2.55× worse, not matching**. By the written rule, we stay.
|
||||
|
||||
But the written rule missed something. It assumed Lance's value would show up as raw-speed wins across the whole table. The actual benchmark reveals Lance's value is **in capabilities the current stack doesn't have**, not in the metrics we scoped:
|
||||
|
||||
1. **Random row access** is 112× faster. Our Parquet design can't do O(1) random access to a row — RAG text retrieval is a full-file scan today. Lance makes this native.
|
||||
2. **Append** is structurally different. Adding 10K rows is 0.08s on Lance; on our stack it's a full rewrite of the entire 330MB Parquet file plus re-embedding plus re-indexing.
|
||||
3. **Index build** is 14× faster. The HNSW `ec=80 es=30` production default takes 230s; Lance IVF_PQ takes 16s. Hot-swap generation (Phase 16) is much more feasible at 16s per build.
|
||||
|
||||
## The decision
|
||||
|
||||
**Hybrid architecture — neither replace nor reject.**
|
||||
|
||||
### What stays
|
||||
|
||||
- `vectord::store` with Parquet + binary-blob vectors → **primary vector backend**
|
||||
- `vectord::hnsw::HnswStore` → in-RAM HNSW for search at 100K-scale indexes
|
||||
- All Phase 15 trial infrastructure → keeps working, unchanged
|
||||
- Production default `ec=80 es=30` → still the right call for in-RAM use
|
||||
|
||||
### What gets added
|
||||
|
||||
- **`vectord::lance_store`** — second backend using Lance as the persistence layer
|
||||
- Scope: indexes where *any* of the following apply:
|
||||
- Corpus exceeds ~5M vectors (our in-RAM ceiling)
|
||||
- Workload is append-heavy (incremental ingest from streaming sources)
|
||||
- Text retrieval dominates (point lookups by doc_id for RAG)
|
||||
- Hot-swap generations are required (Phase 16)
|
||||
- Implemented as a standalone crate first (follow the pilot layout), promoted into vectord when the API stabilizes
|
||||
- **Profile-level configuration** — `ModelProfile.vector_backend: Parquet | Lance` so each profile picks the tier that matches its workload
|
||||
|
||||
### What we keep watching (but don't act on yet)
|
||||
|
||||
- **Lance search latency at scale.** 2229us at 100K is worse than HNSW. At 10M we expect Lance to pull ahead because HNSW doesn't fit in RAM. Re-benchmark when we have a 10M-vector corpus to test against.
|
||||
- **IVF_PQ recall.** We measured latency but not recall — I picked `num_partitions=316, nbits=8, num_sub_vectors=48` blindly. A proper recall sweep is part of Phase C when we integrate Lance into the trial system.
|
||||
- **Lance's own HNSW-on-disk variant** (`with_ivf_hnsw_pq_params`). Might close the in-RAM latency gap. Left for a future pilot.
|
||||
|
||||
## Why this isn't moving the goalposts
|
||||
|
||||
The EXECUTION_PLAN rule was "migrate or don't migrate." The evidence says neither is correct — one stack can't serve both the staffing SQL workload AND the LLM-brain append-heavy random-access workload at all scales. The honest answer is two backends, each doing what it's good at, selected per-profile.
|
||||
|
||||
This matches the dual-use framing in the 2026-04-16 PRD update: different workloads, shared substrate, per-profile specialization. We wrote that principle into the PRD; the benchmark data just made it concrete for the vector tier.
|
||||
|
||||
## Follow-up work (updates EXECUTION_PLAN.md)
|
||||
|
||||
- **Phase C (decoupled embedding refresh)** gets easier — Lance's native append removes the need to invent a "vectors delta" Parquet layer. When we build Phase C, use Lance as the embedding-layer backend.
|
||||
- **Phase 16 (hot-swap)** becomes feasible — 16s index builds mean online re-trials are cheap. When we build Phase 16, Lance is the storage for index generations.
|
||||
- **Phase 17 (model profiles)** gains a new field: `vector_backend: Parquet | Lance`. Default Parquet for backward compatibility. Agents can opt into Lance.
|
||||
|
||||
## Costs we accept
|
||||
|
||||
- **Second dependency tree.** Lance pulls in DataFusion 52 and Arrow 57, while our main stack runs DataFusion 47 and Arrow 55. Keeping lance-bench isolated works for a pilot; productionizing will need either workspace-wide upgrade or a firewall via a dedicated `vectord-lance` crate.
|
||||
- **Second API surface.** Lance's vector-index API is different from our HNSW code. Per-profile abstraction cost is real.
|
||||
- **Operational complexity.** Two vector storage implementations to debug and monitor.
|
||||
|
||||
Worth it because the alternative — forcing every workload through one backend — means either the staffing case or the LLM-brain case is served badly.
|
||||
|
||||
## Ceilings this updates in PRD
|
||||
|
||||
The PRD "Known ceilings" table had:
|
||||
|
||||
> Vector count per index | ~5M vectors on 128GB RAM | 10M+ (serious web crawl) | Phase 18 Lance migration OR mmap'd embeddings
|
||||
|
||||
Update to:
|
||||
|
||||
> Vector count per index | ~5M vectors on 128GB RAM (Parquet+HNSW in-RAM) | Past 5M | Switch that profile's `vector_backend` to Lance; IVF_PQ keeps working on disk-resident quantized codes
|
||||
@@ -79,3 +79,18 @@
|
||||
**Date:** 2026-03-27
|
||||
**Decision:** Each contract/search gets a named workspace with saved queries, shortlists, activity logs, and delta layers. Workspaces have daily/weekly/monthly tiers and support instant zero-copy handoff between agents.
|
||||
**Rationale:** Staffing workflows are inherently agent-centric — a recruiter works a contract, builds context, then may need to hand it off. The workspace captures that context in a structured, queryable, transferable format. Without it, handoff means "read the email thread and figure it out."
|
||||
|
||||
## ADR-017: Federated multi-bucket storage
|
||||
**Date:** 2026-04-16
|
||||
**Decision:** Every `ObjectRef` belongs to exactly one named bucket. Three bucket classes: `primary` (system default, always present), `profile:{user}` (per-user/per-model workspace bucket), and named tenant buckets (`client_a`, `client_eu`). A single shared `rescue_bucket` handles read fallback on target failure. Writes hard-fail on unreachable target; no silent fallback. Every bucket op failure lands in an error journal at `primary://_errors/bucket_errors/`, queryable via `/storage/errors`. Credentials never live in `lakehouse.toml` — a pluggable `SecretsProvider` trait resolves opaque `secret_ref` handles.
|
||||
**Rationale:** The single-backend assumption breaks when we want: tenant data isolation, data residency (EU bucket), per-profile workspaces for local models. DataFusion already supports multiple registered object stores; catalogd is the cross-bucket metadata authority. Rescue bucket + visible error journal makes operational failures diagnosable in one HTTP call. See `docs/ADR-017-federation.md` for full design + success gates.
|
||||
|
||||
## ADR-018: Write-once batched append pattern for journals
|
||||
**Date:** 2026-04-16
|
||||
**Decision:** All append-only journals (error journal, HNSW trial journal, future audit logs) use the `storaged::append_log::AppendLog` helper. Events accumulate in an in-memory buffer; on threshold or explicit `flush()`, the buffer is written as one new timestamped file (`batch_{epoch_us}.jsonl`). Existing files are never rewritten. `compact()` merges all batches into one with a fresh timestamp, preserving chronological sort order.
|
||||
**Rationale:** Object stores have no append primitive. Naive "read-modify-write the whole JSONL file on every event" is O(N²) cumulative work and creates the classic small-file / rewrite-amplification anti-pattern that llms3.com flags as the top lakehouse pitfall. Write-once batching is the LSM-tree idea applied to small JSONL events — bounded write amplification, append-only semantics, optional compaction for read efficiency. The in-memory ring buffer preserves O(1) recent-event reads for the `/storage/errors` and `/hnsw/trials` query endpoints.
|
||||
|
||||
## ADR-019: Vector storage — Parquet+HNSW primary, Lance secondary (hybrid)
|
||||
**Date:** 2026-04-16
|
||||
**Decision:** Keep Parquet + binary-blob vectors + in-RAM HNSW as the primary vector backend. Add Lance as a second backend available per-profile for workloads where Lance wins architecturally. Per-profile `vector_backend: Parquet | Lance` field becomes part of Phase 17 model profiles. Implementation kicks off via the standalone `crates/lance-bench` crate and is promoted into `vectord::lance_store` when the API stabilizes.
|
||||
**Rationale:** Head-to-head benchmark on the 100K × 768d `resumes_100k_v2` index (see `docs/ADR-019-vector-storage.md` for the full scorecard). Parquet+HNSW wins current-scale search latency by 2.55× (873us vs 2229us p50). Lance wins index build time by 14× (16s vs 230s), random row access by 112× (311us vs ~35ms full-file scan), and append speed structurally (0.08s vs full Parquet rewrite). Neither strictly dominates — the dual-use PRD framing (staffing + LLM brain) means both workloads exist in the same system. Keeps ADR-008's "Parquet is the format" principle intact for dataset tables; adds Lance as a purpose-built vector-tier option without discarding the tuned HNSW stack.
|
||||
|
||||
@@ -0,0 +1,233 @@
|
||||
# Execution Plan — Phases B through E
|
||||
|
||||
**Created:** 2026-04-16
|
||||
**Status:** Active planning document — update as phases complete or scope shifts
|
||||
**Owner:** J
|
||||
|
||||
This plan sequences the work J and Claw agreed on during the 2026-04-16 reframe session, after stress-testing the "dual-use substrate" vision and aligning it with llms3.com's architectural patterns.
|
||||
|
||||
---
|
||||
|
||||
## The four phases, at a glance
|
||||
|
||||
| Phase | Work | Prereq | Estimated cost | Risk |
|
||||
|---|---|---|---|---|
|
||||
| B | Lance pilot on one vector index | None | 1 focused session | Medium — new dep, unfamiliar surface |
|
||||
| C | Decoupled embedding refresh pipeline | Benefits from B's outcome | 1 focused session | Low — additive, doesn't break existing |
|
||||
| D | AI-safe views | Phase 13 (done) | 1 focused session | Low — builds on existing catalog + tool registry |
|
||||
| E | Soft deletes / tombstones | None | 1-2 focused sessions | Medium — touches query path, compaction |
|
||||
|
||||
Each "focused session" ≈ 3-4 hours of coding + verification + doc update.
|
||||
|
||||
---
|
||||
|
||||
## Phase B — Lance pilot (the storage format question)
|
||||
|
||||
### Why now
|
||||
|
||||
J's LLMS3 knowledge base explicitly positions Lance as `alternative_to` Parquet for vector workloads. We admitted in the 2026-04-16 stress test that our Parquet-portability argument for vectors is weaker than advertised — our vector Parquet blobs aren't readable by DuckDB/Polars anyway. Lance could unlock: random-row access, disk-resident indexes, time-travel, better compression. Or it could disappoint and we lock in Parquet with a written reason why.
|
||||
|
||||
**Commit to one answer backed by measurements.** No ambiguity after this phase.
|
||||
|
||||
### Scope
|
||||
|
||||
1. Add `lance` crate as a dep in `vectord`, behind a `feature = "lance"` flag initially so we can build without forcing users to install it
|
||||
2. New module `vectord::lance_store` — mirror of `vectord::store` but against Lance format
|
||||
3. New endpoint `POST /vectors/lance/index` — build a Lance index for a named source (parallel to the existing Parquet path)
|
||||
4. Benchmark script that runs against `resumes_100k_v2` (the existing 100K reference)
|
||||
|
||||
### Measured dimensions
|
||||
|
||||
All benchmarks against `resumes_100k_v2` (100K × 768d), cold start:
|
||||
|
||||
| Metric | Parquet baseline | Lance | Threshold to migrate |
|
||||
|---|---|---|---|
|
||||
| Cold load from disk | ~2.8s (measured) | ? | ≥2× faster |
|
||||
| Search latency p50 | 873us (ec=80 es=30) | ? | Within 50% |
|
||||
| Disk size | 330MB | ? | Comparable or better |
|
||||
| Single-row random access | N/A (full scan) | ? | <10ms |
|
||||
| Append 10K new rows | Full rewrite (~3s) | ? | Incremental <500ms |
|
||||
|
||||
### Decision rules
|
||||
|
||||
- **Lance wins cold-load by ≥2× AND matches search latency:** migrate vector storage to Lance. Dataset tables stay Parquet. Update ADR-008 → ADR-019.
|
||||
- **Lance within 50% across board:** stay Parquet. Document ceiling honestly in PRD (already done). Lance revisit when we have a problem Parquet can't solve.
|
||||
- **Lance loses:** close the door, don't revisit without new evidence. Write ADR-019 as "Lance evaluated, rejected, here's why with numbers."
|
||||
|
||||
### Success gate
|
||||
|
||||
Benchmark output table posted to `docs/ADR-019-vector-storage.md` with measured numbers in each cell of the table above. Decision rule applied mechanically. No "let's defer the call."
|
||||
|
||||
### Rollback
|
||||
|
||||
The `feature = "lance"` flag means if the pilot goes badly, `cargo build` without the flag is unchanged. No production migration happens until ADR-019 commits to the change. Safe experiment.
|
||||
|
||||
---
|
||||
|
||||
## Phase C — Decoupled embedding refresh
|
||||
|
||||
### Why now
|
||||
|
||||
llms3.com's lakehouse architecture explicitly separates "transactional data mutations" from "asynchronous vector refresh cycles." Today we couple them — an ingest writes rows AND embeddings in one flow. That means:
|
||||
|
||||
- Adding 1K rows to a 100K-row dataset forces re-embedding of ALL rows (or nothing)
|
||||
- No notion of "embeddings are stale, schedule a refresh tonight"
|
||||
- The embedding cost (Ollama-bound, the bottleneck) is synchronous with ingest
|
||||
|
||||
### Scope
|
||||
|
||||
1. Add fields to `DatasetManifest`:
|
||||
- `last_embedded_at: Option<DateTime>`
|
||||
- `embedding_stale_since: Option<DateTime>` (set when data written but embeddings not refreshed)
|
||||
- `embedding_refresh_policy: RefreshPolicy` — `Manual` | `OnAppend` | `Scheduled(cron)`
|
||||
2. Decouple ingest from embed: ingest writes data + marks embeddings stale; embedding runs separately
|
||||
3. New endpoint: `POST /vectors/refresh/{dataset}` — diffs existing vectors vs current rows, only embeds new/changed (keyed by `doc_id`)
|
||||
4. Background scheduler (or cron trigger) — for `Scheduled` policy, re-runs refresh per schedule
|
||||
5. `GET /vectors/stale` — lists datasets with stale embeddings and how stale
|
||||
|
||||
### Measured success
|
||||
|
||||
- Ingest a 1K-row append to `kb_team_runs` (currently 586 rows, Postgres-sourced).
|
||||
- Catalog shows `embedding_stale_since = now`.
|
||||
- `POST /vectors/refresh/kb_team_runs` embeds only the 1K new rows, not all 1586.
|
||||
- Result: new rows searchable, old embeddings unchanged, total Ollama time ~5s instead of ~30s.
|
||||
|
||||
### Dependencies on Phase B
|
||||
|
||||
If Lance wins Phase B, this is dramatically easier — Lance supports native append. If we stay Parquet, we need a "vectors delta" Parquet file that merges at read time (same pattern as Phase 8's data delta files). ~100 extra LOC if we stay Parquet.
|
||||
|
||||
### Rollback
|
||||
|
||||
The `refresh_policy` field defaults to `Manual` for all existing datasets, so no behavior change for anything already in the system. Opt-in per dataset.
|
||||
|
||||
---
|
||||
|
||||
## Phase D — AI-safe views
|
||||
|
||||
### Why now
|
||||
|
||||
llms3.com's framing: "AI-safe views enforcing row/column security + PII tokenization before model exposure." Phase 13 gave us role-based column masking at query time. That's per-query enforcement. "Views" means pre-materialized: create `candidates_safe` once, bind model X to that view, model X can never accidentally see raw `candidates`.
|
||||
|
||||
This is also the precondition for Phase 17 (model profiles) to be meaningfully safe. "Bind model to dataset" isn't enough — needs to be "bind model to a safe view of the dataset."
|
||||
|
||||
### Scope
|
||||
|
||||
1. New catalog entity: `AiView` with fields `name`, `base_dataset`, `columns[]` (whitelist), `row_filter` (optional SQL WHERE clause), `column_redactions[]` (PII tokenization rules)
|
||||
2. Persistence: `_catalog/views/{name}.json` alongside manifests
|
||||
3. Query-rewrite layer: when a query references `candidates_safe`, DataFusion sees an equivalent `SELECT (whitelisted cols) FROM candidates WHERE (row_filter)` — with redactions applied as expressions
|
||||
4. Endpoint: `POST /catalog/views` to create, `GET /catalog/views` to list, `GET /catalog/views/{name}/preview` to see what it looks like
|
||||
5. Tool registry integration: tools can bind to an AiView instead of a raw table; agent invocations go through the view
|
||||
|
||||
### Measured success
|
||||
|
||||
- Create view `candidates_safe` = `SELECT candidate_id, skills, city FROM candidates WHERE status != 'blocked'`.
|
||||
- Agent (tool registry) calls `search_candidates` bound to `candidates_safe`.
|
||||
- Agent cannot see `email`, `phone`, `ssn`, or `status='blocked'` rows, even if it writes raw SQL.
|
||||
- Audit log records agent accessed `candidates_safe`, not `candidates`.
|
||||
|
||||
### Dependencies
|
||||
|
||||
- Phase 13 already provides the sensitivity classification layer
|
||||
- Phase 12 tool registry already exists
|
||||
- This phase is the bridge between them for agent access
|
||||
|
||||
### Rollback
|
||||
|
||||
Views are additive. Dropping the feature = delete view definitions, tool registry falls back to direct table access. No data migration needed.
|
||||
|
||||
---
|
||||
|
||||
## Phase E — Soft deletes / tombstones
|
||||
|
||||
### Why now
|
||||
|
||||
GDPR/CCPA compliance for staffing data. Today, `ops::delete` physically deletes a parquet object — fine for whole datasets, useless for "delete one candidate's record." To delete one row we'd have to rewrite the whole `candidates.parquet` which at 100K rows is 10MB of churn per deletion.
|
||||
|
||||
llms3.com lists "deletion vectors" as a core lakehouse pattern (Iceberg/Delta/Hudi all implement it). This is the single biggest compliance gap in the current system.
|
||||
|
||||
### Scope
|
||||
|
||||
1. New sidecar per dataset: `{dataset}_tombstones.parquet` with columns `{row_key, deleted_at, actor, reason}`
|
||||
2. Delete API: `POST /catalog/datasets/{name}/tombstone` with `{row_keys[], reason, actor}`
|
||||
3. Query-time filter: `queryd` automatically LEFT JOINs tombstones and filters out deleted rows
|
||||
4. Compaction integration: Phase 8 compaction reads base + delta + tombstones, writes a clean base without tombstoned rows, clears the tombstone sidecar
|
||||
5. Event journal integration (Phase 9): every tombstone emits a journal event with full context
|
||||
|
||||
### Measured success
|
||||
|
||||
- `POST /catalog/datasets/candidates/tombstone` with `{row_keys: ["CAND-123"], reason: "GDPR request", actor: "legal@company"}`
|
||||
- `SELECT COUNT(*) FROM candidates` drops by 1 immediately
|
||||
- `SELECT * FROM candidates WHERE candidate_id = 'CAND-123'` returns empty
|
||||
- `GET /journal/history/CAND-123` shows the tombstone event
|
||||
- After scheduled compaction, the tombstone is materialized — `candidates.parquet` no longer contains CAND-123, tombstone sidecar is emptied for that row key
|
||||
|
||||
### Dependencies
|
||||
|
||||
- Phase 8 delta/merge-on-read pattern (done) — tombstones are a third layer at read time
|
||||
- Phase 9 event journal (done) — tombstones emit journal events
|
||||
|
||||
### Rollback
|
||||
|
||||
If query rewrite becomes too complex, fallback: tombstones stored but applied only during compaction (not at query time). Queries return deleted rows until compaction runs. Less useful but safer.
|
||||
|
||||
---
|
||||
|
||||
## Cross-phase concerns
|
||||
|
||||
### Phases that need federation layer 2 (task #5)
|
||||
|
||||
Every phase above assumes the federation foundation (shipped 2026-04-16) but NOT federation layer 2 (cross-bucket SQL, profile activation, `X-Lakehouse-Bucket` header).
|
||||
|
||||
**Implication:** Phases B-E can proceed on the `primary` bucket without blocking on federation layer 2. Federation layer 2 becomes valuable when we want multi-profile model scoping (Phase 17). Sequence:
|
||||
|
||||
```
|
||||
A (done) → B → C (+D in parallel) → federation layer 2 → Phase 16 → Phase 17 → E
|
||||
```
|
||||
|
||||
### Phases that need federation layer 2 FIRST
|
||||
|
||||
None of B/C/D/E strictly need it. Phase 16 (hot-swap) benefits from it. Phase 17 (model profiles) depends on it heavily.
|
||||
|
||||
### What NOT to build in B-E
|
||||
|
||||
- Distributed query — wait for a real scale problem
|
||||
- Replacement of DataFusion — working fine, stay put
|
||||
- Iceberg/Delta Lake migration — explicitly out of scope per ADR-009
|
||||
- Live streaming / CDC — explicit non-goal
|
||||
|
||||
### Definition of done for each phase
|
||||
|
||||
Each phase completes when:
|
||||
1. Code shipped and building clean
|
||||
2. Success gate measurably passed
|
||||
3. Relevant ADR added to `docs/DECISIONS.md` (or updated)
|
||||
4. `docs/PHASES.md` checkbox flipped with measurement data
|
||||
5. PRD invariants checked — if a new invariant emerged, add it
|
||||
6. One regression test in the crate or HTTP integration test
|
||||
|
||||
---
|
||||
|
||||
## Session plan (what to do in what order)
|
||||
|
||||
### Next session
|
||||
1. **Phase B — Lance pilot.** Single session. Answers the biggest open architectural question.
|
||||
2. Based on outcome, **write ADR-019** with the decision + data.
|
||||
3. **Update PHASES.md** with Phase 18 status (Lance evaluated).
|
||||
|
||||
### Session after
|
||||
4. **Phase C — Decoupled embedding refresh.** Implementation shaped by B's outcome (append is easy on Lance, requires delta logic on Parquet).
|
||||
|
||||
### Session after that
|
||||
5. **Federation layer 2** OR **Phase D (AI-safe views)** — J decides based on priority. Federation layer 2 unlocks model profiles (Phase 17); AI-safe views is standalone value.
|
||||
|
||||
### Final session for this track
|
||||
6. **Phase E — Soft deletes.** The compliance-driven phase. Fits cleanly after everything else because it touches the query path and wants to be built after query optimizations stabilize.
|
||||
|
||||
### Milestone checkpoint
|
||||
After Phase E, stop and reassess. We'll have:
|
||||
- Lance decision made and committed
|
||||
- Decoupled embedding pipeline
|
||||
- AI-safe view enforcement
|
||||
- Soft delete semantics
|
||||
|
||||
That's a substantial capability increase. Plausible "pause, write a retrospective, decide on Phase 16/17" moment.
|
||||
+122
-7
@@ -121,15 +121,130 @@
|
||||
- [x] AI migration prompt builder for complex cases
|
||||
- [x] 5 unit tests
|
||||
|
||||
## Phase 15+: Horizon ⬜
|
||||
- [ ] HNSW vector index (100K search: 4.5s → <50ms)
|
||||
- [ ] Federated multi-bucket query
|
||||
- [ ] Database connector ingest (Postgres/MySQL)
|
||||
- [ ] PDF OCR (Tesseract)
|
||||
- [ ] Scheduled ingest (cron)
|
||||
## Phase 15+: Horizon
|
||||
- [x] HNSW vector index with iteration-friendly trial system (2026-04-16)
|
||||
- `HnswStore.build_index_with_config` — parameterized ef_construction, ef_search, seed
|
||||
- `EmbeddingCache` — pins 100K vectors in memory, shared across trials
|
||||
- `harness::EvalSet` — named query sets with brute-force ground truth
|
||||
- `TrialJournal` — append-only JSONL at `_hnsw_trials/{index}.jsonl`
|
||||
- Endpoints: `/vectors/hnsw/trial`, `/hnsw/trials/{idx}`, `/hnsw/trials/{idx}/best?metric={recall|latency|pareto}`, `/hnsw/evals`, `/hnsw/evals/{name}/autogen`, `/hnsw/cache/stats`
|
||||
- Measured on 100K resumes: **brute-force 44-54ms → HNSW 509us-1830us**, recall 0.92-1.00 depending on `ef_construction`. Sweet spot: ec=80 es=30 → p50=873us recall=1.00 — locked in as `HnswConfig::default()`
|
||||
- [x] Catalog manifest repair — `POST /catalog/resync-missing` restores row_count and columns from parquet footers (2026-04-16). All 7 staffing tables recovered to PRD-matching 2.47M rows.
|
||||
- [~] Federated multi-bucket query — **foundation complete 2026-04-16**, see ADR-017
|
||||
- [x] `StorageConfig.buckets` + `rescue_bucket` + `profile_root` config shape
|
||||
- [x] `SecretsProvider` trait + `FileSecretsProvider` (reads /etc/lakehouse/secrets.toml, checks 0600 perms)
|
||||
- [x] `storaged::BucketRegistry` — multi-backend, rescue-aware, reachability probes
|
||||
- [x] `storaged::error_journal::ErrorJournal` — append-only JSONL at `primary://_errors/bucket_errors.jsonl`
|
||||
- [x] Endpoints: `GET /storage/buckets`, `GET /storage/errors`, `GET /storage/bucket-health`
|
||||
- [x] Bucket-aware I/O: `PUT/GET /storage/buckets/{bucket}/objects/{*key}` with rescue fallback + `X-Lakehouse-Rescue-Used` observability headers
|
||||
- [x] Backward compat: empty `[[storage.buckets]]` synthesizes a `primary` from legacy `root`
|
||||
- [x] Three-bucket test (primary + rescue + testing) verified: normal reads, rescue fallback with headers, hard-fail missing, write to unknown bucket 503, error journal + health summary
|
||||
- [x] `X-Lakehouse-Bucket` header middleware on ingest endpoints (2026-04-16)
|
||||
- [x] Catalog migration: `POST /catalog/migrate-buckets` stamps `bucket = "primary"` on legacy refs (12 renamed, 14 total now canonical)
|
||||
- [x] `queryd` registers every bucket with DataFusion for cross-bucket SQL — verified with people_test (testing) × animals (primary) CROSS JOIN
|
||||
- [x] Profile hot-load endpoints: bucket auto-provisioning on `POST /vectors/profile/{id}/activate` (2026-04-17)
|
||||
- [x] `vectord` bucket-scoped paths: TrialJournal + PromotionRegistry resolve per-index via IndexMeta.bucket (2026-04-17)
|
||||
- [x] Runtime bucket lifecycle: `POST /storage/buckets` (provision) + `DELETE /storage/buckets/{name}` (unregister, refuses primary/rescue) (2026-04-17)
|
||||
- [x] ModelProfile.bucket field — per-profile artifact isolation (2026-04-17)
|
||||
- [x] Database connector ingest (Postgres first) — 2026-04-16
|
||||
- `pg_stream::stream_table_to_parquet` — ORDER BY + LIMIT/OFFSET pagination, configurable batch_size
|
||||
- `parse_dsn` — postgresql:// and postgres:// URL scheme, user/password/host/port/db
|
||||
- `POST /ingest/db` endpoint: `{dsn, table, dataset_name?, batch_size?, order_by?, limit?}` → streams to Parquet, registers in catalog with PII detection + redacted-password lineage
|
||||
- Existing `POST /ingest/postgres/import` (structured config) preserved alongside
|
||||
- 4 DSN-parser unit tests + live end-to-end test against `knowledge_base.team_runs` (586 rows, 13 cols, 6 batches, 196ms)
|
||||
- [x] Phase B: Lance storage evaluation — 2026-04-16
|
||||
- `crates/lance-bench` standalone pilot (Lance 4.0) avoids DataFusion/Arrow version conflict with main stack
|
||||
- 8-dimension benchmark on resumes_100k_v2 — see docs/ADR-019-vector-storage.md for scorecard
|
||||
- Decision: hybrid architecture. Parquet+HNSW stays primary (2.55× faster search at 100K in-RAM). Lance added as per-profile second backend for random access (112× faster), append (0.08s vs full rewrite), hot-swap (14× faster index builds), and scale past 5M RAM ceiling.
|
||||
- [x] Phase E.2 — Compaction integrates tombstones (physical deletion) — 2026-04-16
|
||||
- `delta::compact` accepts `tombstones: &[Tombstone]` param, filters rows at merge time via arrow `filter_record_batch`
|
||||
- CompactResult gains `tombstones_applied` + `rows_dropped_by_tombstones`
|
||||
- Atomic write: ArrowWriter → single Parquet file (fixes latent bug where concatenated Parquet byte streams produced garbage — footer-only-first-segment visible), verify-parse before overwrite, temp_key staging, delete delta files AFTER base write succeeds
|
||||
- Snappy compression on output matches ingest defaults (avoids 3× size inflation on every compact)
|
||||
- `TombstoneStore::clear` drops all batch files for a dataset; called by queryd after successful compact
|
||||
- Query engine exposes `catalog()` accessor so service handler can reach the tombstone store
|
||||
- E2E verified on candidates (100K rows): tombstone 3 IDs → compact → 99,997 rows physically in parquet, tombstones empty, IDs gone from `__raw__candidates` too; file size 10.59 MB → 10.72 MB (proportional to data, not inflated)
|
||||
- [x] Phase 16: Hot-swap generations + autotune agent — 2026-04-16
|
||||
- `vectord::promotion::PromotionRegistry` — per-index current config + history at `_hnsw_promotions/{index}.json`, cap 50 history entries
|
||||
- Endpoints: `POST /vectors/hnsw/promote/{index}/{trial_id}`, `POST /vectors/hnsw/rollback/{index}`, `GET /vectors/hnsw/promoted/{index}`
|
||||
- `vectord::autotune::run_autotune` — grid of trials (configurable or default 5 configs), Pareto winner selection (max recall, then min p50), min_recall safety gate (default 0.9), config bounds (ec ∈ [10,400], es ∈ [10,200])
|
||||
- `POST /vectors/hnsw/autotune` — runs the full loop synchronously, journals every trial, auto-promotes winner
|
||||
- `activate_profile` uses `promotion_registry.config_or(..., profile_default)` so newly-promoted configs flow automatically into next activation
|
||||
- End-to-end: autogen harness for threat_intel_v1 (10 queries), autotune ran 5 trials (all recall=1.00, p50 64-68us), promoted ec=20 es=30 at recall=1.0 p50=64us as winner. Manual promote of ec=80 es=30 pushed autotune pick onto history. Rollback restored autotune winner. Second rollback cleared to None. Re-promote + restart verified persistence. Activation after promotion logged "building HNSW ef_construction=80 ef_search=30 seed=42" — config flowed through correctly.
|
||||
- [x] Phase 17: Model profiles + scoped search — 2026-04-16
|
||||
- `shared::types::ModelProfile` — { id, ollama_name, description, bound_datasets, hnsw_config, embed_model, created_at, created_by }
|
||||
- `shared::types::ProfileHnswConfig` — mirror of vectord's HnswConfig to avoid cross-crate dep cycle (defaults ec=80 es=30 matching Phase 15 winner)
|
||||
- `Registry::{put_profile, get_profile, list_profiles, delete_profile}` persisted at `_catalog/profiles/{id}.json`, validates bindings exist (raw dataset OR AiView)
|
||||
- Endpoints: `POST/GET /catalog/profiles`, `GET/DELETE /catalog/profiles/{id}`
|
||||
- `POST /vectors/profile/{id}/activate` — warms EmbeddingCache + builds HNSW with profile's config for every bound dataset's vector index; reports warmed indexes + failures + duration
|
||||
- `POST /vectors/profile/{id}/search` — rejects 403 if requested index's source isn't in profile.bound_datasets; falls through to HNSW if warm, brute-force otherwise
|
||||
- Fixed refresh to register new index metadata (was silently no-op for first-time indexes)
|
||||
- End-to-end: security-analyst profile bound to threat_intel → activate warms 54 vectors in 156ms → within-scope HNSW search works (0.625 score); out-of-scope search for candidates returns 403 with allowed bindings listed
|
||||
- [x] Phase E: Soft deletes (tombstones) — 2026-04-16
|
||||
- `shared::types::Tombstone` — { dataset, row_key_column, row_key_value, deleted_at, actor, reason }
|
||||
- `catalogd::tombstones::TombstoneStore` per-dataset append-log at `_catalog/tombstones/{dataset}/`, flush_threshold=1 + explicit flush so every tombstone is durable on return (compliance requirement)
|
||||
- All tombstones for a dataset must share the same `row_key_column` (validated at write — query filter is built as a single WHERE NOT IN clause)
|
||||
- `Registry::add_tombstone / list_tombstones`
|
||||
- Endpoint: `POST /catalog/datasets/by-name/{name}/tombstone` accepting `{row_key_column, row_key_values[], actor, reason}`; companion `GET` lists active tombstones
|
||||
- `queryd::context::build_context` wraps tombstoned tables: raw goes to `__raw__{name}`, public name becomes a DataFusion view with `WHERE CAST(col AS VARCHAR) NOT IN (...)` filter
|
||||
- End-to-end on candidates: tombstone 3 IDs, COUNT drops 100,000 → 99,997, specific WHERE returns empty, AiView candidates_safe transitively excludes them too, restart preserves all tombstones
|
||||
- Limits / not in MVP: physical compaction (Phase 8 doesn't yet read tombstones during merge); journal integration (tombstones don't yet emit Phase 9 mutation events — covered by audit fields on the tombstone itself)
|
||||
- [x] Phase D: AI-safe views — 2026-04-16
|
||||
- `shared::types::AiView` — name, base_dataset, columns whitelist, optional row_filter, column_redactions
|
||||
- `shared::types::Redaction` — Null | Hash | Mask { keep_prefix, keep_suffix }
|
||||
- `Registry::put_view / get_view / list_views / delete_view` persisted to `_catalog/views/{name}.json`
|
||||
- `queryd::context` registers each view as a DataFusion view with the safe projection + filter + redactions baked into the SELECT
|
||||
- Endpoints: `POST/GET /catalog/views`, `GET/DELETE /catalog/views/{name}`
|
||||
- End-to-end on candidates: `candidates_safe` view exposes 8 of 15 columns, masks `candidate_id` (CAN******01), filters out `status='blocked'`. `SELECT * FROM candidates_safe` returns whitelist only; `SELECT email FROM candidates_safe` fails. View survives restart.
|
||||
- Capability surface — raw `candidates` still accessible by name; Phase 13 access control is the layer that enforces who can query what
|
||||
- [x] Phase C: Decoupled embedding refresh — 2026-04-16
|
||||
- `DatasetManifest`: `last_embedded_at`, `embedding_stale_since`, `embedding_refresh_policy` (Manual | OnAppend | Scheduled)
|
||||
- `Registry::mark_embeddings_stale` / `clear_embeddings_stale` / `stale_datasets`
|
||||
- Ingest paths (CSV pipeline + Postgres streaming) auto-mark-stale when writing to an already-embedded dataset
|
||||
- `vectord::refresh::refresh_index` — reads dataset, diffs doc_ids vs existing embeddings, embeds only new rows, writes combined index, clears stale
|
||||
- `POST /vectors/refresh/{dataset}` + `GET /vectors/stale`
|
||||
- Id columns accept `Utf8`, `Int32`, `Int64`
|
||||
- End-to-end on threat_intel: initial 20-row embed 2.1s; re-ingest to 54 rows auto-marks stale; delta refresh embeds only 34 new in 970ms (6× faster than full re-embed); stale cleared
|
||||
- [x] Phase 16.2/16.5: Background autotune agent + ingest-triggered re-trials — 2026-04-17
|
||||
- `vectord::agent` — ε-greedy proposer, rate-limited, cooldown-gated, tokio background task
|
||||
- Ingest paths push `DatasetAppended` triggers to agent queue
|
||||
- Endpoints: `GET /vectors/agent/status`, `POST /vectors/agent/stop`, `POST /vectors/agent/enqueue/{idx}`
|
||||
- `[agent]` config section in lakehouse.toml (enabled, cycle_interval, cooldown, min_recall, max_trials/hr)
|
||||
- 3 unit tests
|
||||
- [x] Phase 17 VRAM gate: Two-profile sequential swap — 2026-04-17
|
||||
- Sidecar: `POST /admin/unload` (keep_alive=0), `POST /admin/preload`, `GET /admin/vram` (nvidia-smi + Ollama /api/ps)
|
||||
- `AiClient::unload_model / preload_model / vram_snapshot`
|
||||
- `VectorState.active_profile` singleton — activate swaps models, deactivate unloads
|
||||
- Verified: staffing-recruiter (qwen2.5) ↔ docs-assistant (mistral) — only one model in VRAM at a time
|
||||
- [x] MySQL streaming connector — 2026-04-17
|
||||
- `my_stream.rs` mirrors pg_stream: DSN parsing, OFFSET pagination, Arrow type mapping, Parquet streaming
|
||||
- `POST /ingest/mysql` with PII detection, lineage, agent trigger
|
||||
- Verified end-to-end on live MariaDB (10 rows, 9 columns, round-tripped all types)
|
||||
- 6 DSN + type-mapping unit tests
|
||||
- [x] Phase 18 hybrid: vectord-lance production crate — 2026-04-17
|
||||
- Firewall crate (Arrow 57 / Lance 4, separate from main Arrow 55 / DF 47 stack)
|
||||
- Public API: migrate_from_parquet, build_index (IVF_PQ), search, get_by_doc_id, append, build_scalar_index, stats
|
||||
- `lance_backend::LanceRegistry` resolves bucket → URI per index
|
||||
- `VectorBackend { Parquet | Lance }` enum on ModelProfile + IndexMeta
|
||||
- 8 HTTP endpoints under `/vectors/lance/*` (migrate, index, search, doc, append, stats, scalar-index, recall)
|
||||
- Profile-driven routing: `POST /vectors/profile/{id}/search` auto-routes to Lance when profile.vector_backend=lance
|
||||
- Auto-migrate + auto-index on activation
|
||||
- Measured on real 100K × 768d: migrate 0.57s, IVF_PQ build 16.2s (14× faster than HNSW 230s), search 23ms, append 100 rows 3.3ms, doc_id fetch 3.5ms (with scalar btree)
|
||||
- IVF_PQ recall@10 = 0.805 (HNSW = 1.000) — measured via `/vectors/lance/recall/{idx}` harness
|
||||
- [x] Phase E.3: Scheduled ingest — 2026-04-17
|
||||
- `ingestd::schedule` module: ScheduleDef, ScheduleStore (JSON at `_schedules/{id}.json`), Scheduler tokio task
|
||||
- Supports MySQL + Postgres sources on interval triggers (Cron variant defined, parsing stubbed)
|
||||
- 6 CRUD endpoints under `/ingest/schedules/*` + run-now manual trigger
|
||||
- Full catalog integration: PII, lineage, mark-stale, agent trigger
|
||||
- 6 unit tests
|
||||
- [x] PDF OCR via Tesseract — 2026-04-17
|
||||
- Two-tier: lopdf text extraction → Tesseract 5.5 fallback for scanned/image PDFs
|
||||
- Extracts embedded XObject /Image streams, shells to tesseract --oem 3 --psm 6
|
||||
- Same schema (source_file, page_number, text_content) — downstream unchanged
|
||||
- [ ] Fine-tuned domain models
|
||||
- [ ] Multi-node query distribution
|
||||
|
||||
---
|
||||
|
||||
**30 unit tests | 11 crates | 16 ADRs | 2.47M rows | 100K vectors | All built 2026-03-27**
|
||||
**52+ unit tests | 13 crates | 19 ADRs | 2.47M rows | 100K vectors | Hybrid Parquet+HNSW ⊕ Lance**
|
||||
**Latest: 2026-04-17 — 8 commits shipping Phase 16.2 through Phase 18**
|
||||
|
||||
+149
-24
@@ -1,22 +1,41 @@
|
||||
# PRD: Lakehouse — Rust-First Object Storage System
|
||||
# PRD: Lakehouse — Rust-First Substrate for Versioned Knowledge Stores
|
||||
|
||||
**Status:** Active — Phase 0-5 complete, entering production path
|
||||
**Status:** Active — Phases 0-18 shipped; hybrid Parquet+HNSW ⊕ Lance operational; scheduled ingest live; PDF OCR live; entering horizon items
|
||||
**Created:** 2026-03-27
|
||||
**Last reframed:** 2026-04-17 — hybrid architecture proven end-to-end on 100K vectors (see §Phase 18 + ADR-019)
|
||||
**Owner:** J
|
||||
|
||||
---
|
||||
|
||||
## Problem
|
||||
|
||||
### Use case 1 — Staffing analytics (reference implementation)
|
||||
|
||||
Legacy data systems silo information across CRMs, databases, spreadsheets, and file shares. Querying across them requires manual ETL, pre-defined schemas, and expensive database licenses. When AI enters the picture, these systems can't handle the dual requirement of fast analytical queries AND semantic retrieval over unstructured text.
|
||||
|
||||
A staffing company (our reference case) has candidate records in an ATS, client data in a CRM, timesheets in billing software, call logs from a phone system, and email records from Exchange. Answering "find every Java developer in Chicago who was called 5+ times but never placed" requires querying across all of them — and no single system can do it.
|
||||
|
||||
We need a system where:
|
||||
- Any data source (CSV, DB export, PDF, JSON) can be ingested without pre-defined schemas
|
||||
### Use case 2 — Local AI knowledge substrate (the second half)
|
||||
|
||||
Local LLM workloads need a substrate for ingesting, indexing, and retrieving large knowledge corpora. Each running model (or agent) has its own context — documents it cares about, a vector index tuned to its domain, a scoped view of the catalog. That infrastructure is architecturally identical to the staffing problem: ingest messy data, index it, query it, hand it to an AI. Building one substrate that serves both prevents fragmentation.
|
||||
|
||||
Concretely this means a running Ollama model like `qwen2.5:7b` or `claude-code-local` should be able to:
|
||||
- Bind to a named set of datasets
|
||||
- Get a scoped vector index pre-warmed for its domain
|
||||
- Issue searches that only see its bound data
|
||||
- Have its trial/tuning history isolated from other models
|
||||
- Swap between knowledge generations (today's, yesterday's) without rebuild
|
||||
|
||||
The same infrastructure that lets a recruiter query 2.47M rows of staffing data also lets a local 7B model answer questions grounded in a 500K-chunk documentation corpus. Same substrate, different tenant.
|
||||
|
||||
### Shared requirements
|
||||
|
||||
- Any data source (CSV, DB export, PDF, JSON, Postgres table) can be ingested without pre-defined schemas
|
||||
- Structured data is queryable via SQL at scale (millions of rows, sub-second)
|
||||
- Unstructured data is searchable via AI embeddings (semantic retrieval)
|
||||
- An LLM can answer natural language questions against all of it
|
||||
- Unstructured data is searchable via AI embeddings with per-profile indexes
|
||||
- An LLM can answer natural language questions against scoped data
|
||||
- Indexes can be hot-swapped between generations without rebuild downtime
|
||||
- Trials are first-class data — the system remembers how it was tuned
|
||||
- Everything runs locally — no cloud APIs, total data privacy
|
||||
- The system is rebuildable from repository + object storage alone
|
||||
|
||||
@@ -51,15 +70,22 @@ No new frameworks without documented ADR.
|
||||
|
||||
| Service | Responsibility |
|
||||
|---|---|
|
||||
| **gateway** | HTTP/gRPC ingress, routing, auth, CORS, body limits |
|
||||
| **catalogd** | Metadata control plane — dataset registry, schema versions, manifests |
|
||||
| **storaged** | Object I/O — read/write/list/delete via `object_store` |
|
||||
| **queryd** | SQL execution — DataFusion over Parquet, MemTable hot cache |
|
||||
| **ingestd** | *NEW* — Ingest pipeline: CSV/JSON/DB → normalize → Parquet → catalog |
|
||||
| **vectord** | *NEW* — Embedding store + vector index: chunk → embed → index → search |
|
||||
| **gateway** | HTTP/gRPC ingress, routing, auth, CORS, body limits, `X-Lakehouse-Bucket` header routing |
|
||||
| **catalogd** | Metadata control plane — dataset registry, schema versions, manifests, per-dataset resync from parquet footers |
|
||||
| **storaged** | Object I/O — `BucketRegistry` (multi-backend), rescue fallback, error journal, append-log batching pattern |
|
||||
| **queryd** | SQL execution — DataFusion over Parquet, MemTable hot cache, delta merge-on-read |
|
||||
| **ingestd** | Ingest pipeline: CSV / JSON / PDF / Postgres-stream → normalize → Parquet → catalog |
|
||||
| **vectord** | Embedding store + vector indexes + HNSW trial system (EmbeddingCache, trial journal, eval harness) |
|
||||
| **journald** | Append-only mutation event log (ADR-012) — distinct from storaged error journal |
|
||||
| **aibridge** | Rust↔Python boundary — HTTP client to FastAPI sidecar |
|
||||
| **ui** | Dioxus frontend — Ask, Explore, SQL, System tabs |
|
||||
| **shared** | Types, errors, Arrow helpers, config, protobuf definitions |
|
||||
| **shared** | Types, errors, Arrow helpers, config, protobuf definitions, **secrets provider trait**, **PII detection** |
|
||||
|
||||
**Federation building blocks** (shipped 2026-04-16):
|
||||
- `shared::secrets::SecretsProvider` trait + `FileSecretsProvider` reading `/etc/lakehouse/secrets.toml` (0600 enforced)
|
||||
- `storaged::registry::BucketRegistry` — multi-bucket resolution with `rescue_bucket` read fallback
|
||||
- `storaged::append_log::AppendLog` — write-once batched append pattern (no RMW, no small-file problem)
|
||||
- `storaged::error_journal::ErrorJournal` — bucket operation failure log at `primary://_errors/bucket_errors/batch_*.jsonl`
|
||||
|
||||
### Data Flow
|
||||
|
||||
@@ -92,10 +118,14 @@ User question → gateway
|
||||
1. Object storage = source of truth for all data
|
||||
2. catalogd = sole metadata authority
|
||||
3. No raw data in catalog — only pointers
|
||||
4. vectord stores embeddings AS Parquet (portable, not a proprietary format)
|
||||
4. vectord stores embeddings AS Parquet (portable, not a proprietary format) — see ADR-018 for the Parquet-vs-Lance trade review
|
||||
5. ingestd is idempotent — re-ingesting the same file is a no-op
|
||||
6. Hot cache is a performance layer, not a source of truth — eviction is safe
|
||||
7. All services modular and independently replaceable
|
||||
8. **Indexes are hot-swappable.** A new index generation can be built in the background while the current one serves queries. Promotion is atomic (pointer swap). Rollback to a prior generation is always possible. (Phase 16)
|
||||
9. **Every reader gets its own profile.** Human operators, AI agents, and local models are all clients of the same substrate. Each has a named profile with its own bucket, vector indexes, trial history, and dataset bindings. Profiles are a first-class architectural concept, not a tenancy afterthought. (Phase 17)
|
||||
10. **Trials are data, not logs.** Every index build is a trial with measurable metrics. The trial journal IS the agent's memory for how to tune itself. Stored as write-once batched JSONL per the ADR-018 append-log pattern.
|
||||
11. **Operational failures are findable in one HTTP call.** The bucket error journal, trial journal, and audit log all expose `/storage/errors`, `/hnsw/trials`, `/access/audit` with structured filter + aggregation. No `grep` archaeology to answer "what broke?"
|
||||
|
||||
---
|
||||
|
||||
@@ -262,19 +292,90 @@ Per-contract overlays with daily/weekly/monthly tiers and instant handoff.
|
||||
| 14.3 | Migration preview: show how old data maps to new schema before applying | Human approves before data transforms |
|
||||
| 14.4 | Versioned schemas in catalog: v1, v2, v3 coexist | Queries specify version or use latest |
|
||||
|
||||
### Phase 15+: Horizon
|
||||
### Phase 15: Infrastructure horizon items
|
||||
|
||||
- [x] HNSW vector index with trial system (shipped 2026-04-16)
|
||||
- [x] Federation foundation — ADR-017 (shipped 2026-04-16)
|
||||
- [x] Database connector ingest — Postgres batch with streaming (shipped 2026-04-16)
|
||||
- [x] Federation layer 2 — runtime bucket lifecycle, per-index bucket scoping, profile bucket auto-provisioning (shipped 2026-04-17)
|
||||
- [x] MySQL streaming connector — mirrors Postgres path, verified on live MariaDB (shipped 2026-04-17)
|
||||
- [x] PDF OCR for scanned documents — Tesseract 5.5 fallback when lopdf yields no text (shipped 2026-04-17)
|
||||
- [x] Scheduled ingest — interval-based per-source schedules with CRUD + run-now + auto-trigger agent (shipped 2026-04-17)
|
||||
- [ ] Multi-node query distribution (DataFusion supports this architecturally)
|
||||
|
||||
### Phase 16: Hot-Swap Index Generations
|
||||
|
||||
Make indexes upgrade-in-place without dropping queries.
|
||||
|
||||
| Step | Deliverable | Gate |
|
||||
|---|---|---|
|
||||
| 16.1 | "Active generation" pointer per logical index name | `/vectors/search` routes to current champion automatically |
|
||||
| 16.2 | Background trial runner: watches trial journal, proposes configs (random search / Bayesian), fires `/hnsw/trial` | Agent autonomously tunes without human POSTing each config |
|
||||
| 16.3 | Promotion endpoint: `POST /hnsw/promote/{index}/{trial_id}` atomically swaps active pointer | Next search hits new config, zero downtime |
|
||||
| 16.4 | Rollback: `POST /hnsw/rollback/{index}` reverts to previous generation | Bad promotion recoverable in milliseconds |
|
||||
| 16.5 | Dataset-append triggers: when `POST /ingest/file` writes to a dataset with attached vector indexes, schedule automatic re-trial (not full rebuild) | New docs get embedded + indexed without manual intervention |
|
||||
|
||||
**Gate:** Run the trial agent for 10 minutes against `resumes_100k_v2` with a fresh eval set. It explores the `ef_construction × ef_search` space, promotes the Pareto winner, continues running. Zero human clicks. All trials and promotions appear in `/hnsw/trials/resumes_100k_v2`.
|
||||
|
||||
**Risk:** Agent loops into a bad region (e.g. always proposes ef_construction=1). Mitigation: a hardcoded config space constraint + minimum-quality gate (don't promote anything with recall <0.9).
|
||||
|
||||
### Phase 17: Model Profiles + Dataset Bindings
|
||||
|
||||
Make "different models see different data" real instead of a config string.
|
||||
|
||||
| Step | Deliverable | Gate |
|
||||
|---|---|---|
|
||||
| 17.1 | `ModelProfile` manifest: id, ollama_name, bucket, bound_datasets[], hnsw_config, embed_model | `GET /models` lists profiles; `POST /models` creates one |
|
||||
| 17.2 | Profile activation endpoint: `POST /profile/{id}/activate` — warms EmbeddingCache for bound indexes, builds HNSW with profile's config | Next search against bound indexes is <1ms cold |
|
||||
| 17.3 | Model-scoped search: `POST /search?model=X` filters to bound datasets only | Model A can't see Model B's datasets unless explicitly shared |
|
||||
| 17.4 | VRAM-aware activation: only one (or small N) model loaded at a time on 16GB A4000 | Activating model B unloads model A via Ollama's keep_alive=0 |
|
||||
| 17.5 | Audit: every tool invocation by a model is logged with model identity | `GET /models/{id}/audit` shows exactly what each model touched |
|
||||
|
||||
**Gate:** Two model profiles defined: `staffing-recruiter` (bound to candidates/placements/timesheets) and `docs-assistant` (bound to a documentation corpus). Activate staffing-recruiter, search for candidates — works. Switch to docs-assistant, same search — returns zero from staffing (not bound) but finds docs. VRAM shows only one embedding model loaded at a time.
|
||||
|
||||
**VRAM reality:** 16GB A4000 realistically holds 1-2 loaded models concurrently. "Multi-model" in practice means sequential swap between profiles, not parallel serving. The profile abstraction makes this swap clean.
|
||||
|
||||
### Phase 18: Storage format decision (Lance evaluation)
|
||||
|
||||
The question raised 2026-04-16 after J's LLMS3 knowledge base identified Lance as `alternative_to` Parquet for vector workloads. Current stack: Parquet with binary-blob vector columns + in-RAM HNSW sidecar. Evaluated against: Lance native vector format with disk-resident indexes.
|
||||
|
||||
| Step | Deliverable | Decision criteria |
|
||||
|---|---|---|
|
||||
| 18.1 | ✅ Parallel Lance-backed vector index for `resumes_100k_v2` in standalone `crates/lance-bench` | Built 2026-04-16 |
|
||||
| 18.2 | ✅ Head-to-head benchmark across 8 dimensions (cold-load, search latency, disk, index build, random access, append) | Complete |
|
||||
| 18.3 | ✅ ADR-019 committed with measured data and decision | See `docs/ADR-019-vector-storage.md` |
|
||||
|
||||
**Outcome:** Hybrid architecture. Parquet+HNSW stays primary (2.55× faster search at 100K in-RAM). Lance joins as a second backend for Phase 16 hot-swap (14× faster index builds), Phase C/append workloads (0.08s vs full rewrite), RAG random-access retrieval (112× faster), and indexes past the ~5M RAM ceiling.
|
||||
|
||||
Per-profile `vector_backend: Parquet | Lance` becomes part of Phase 17 (model profiles). See ADR-019 for the full scorecard and caveats.
|
||||
|
||||
### Phase 19+: Further horizon
|
||||
|
||||
- Federated multi-bucket query (client A's S3 + client B's S3 + yours)
|
||||
- Database connector ingest (PostgreSQL, MySQL, MSSQL → Parquet via CDC)
|
||||
- PDF OCR for scanned documents (Tesseract integration)
|
||||
- Scheduled ingest (cron-based file watching, S3 event triggers)
|
||||
- Specialized fine-tuned models per domain (staffing matcher, resume parser)
|
||||
- Multi-node query distribution (DataFusion supports this architecturally)
|
||||
- Video/audio transcript ingest + multimodal embeddings
|
||||
- True distributed query (DataFusion multi-node) — only if single-machine ceilings bite
|
||||
|
||||
---
|
||||
|
||||
## Reference Dataset: Staffing Company
|
||||
## Known ceilings (honest)
|
||||
|
||||
The current stack has measurable limits. Documenting them so future decisions aren't based on wishful thinking.
|
||||
|
||||
| Dimension | Current ceiling | Breaks at | Escape hatch |
|
||||
|---|---|---|---|
|
||||
| Vector count per index (Parquet+HNSW in-RAM) | ~5M on 128GB | Past 5M | Switch that profile's `vector_backend` to Lance per ADR-019 — IVF_PQ stays on disk-resident quantized codes |
|
||||
| Concurrent active indexes | ~50-100 at 100K vectors each | 10M×50 configurations | Lance disk-resident + per-profile activation |
|
||||
| Rows per dataset | 2.47M proven, probably 100M+ fine | Approaches DataFusion memory limits | DataFusion predicate pushdown + partition pruning (existing) |
|
||||
| Concurrent loaded models | 1-2 on 16GB VRAM (A4000) | 3+ models simultaneous | Not our problem — architectural, driven by Ollama |
|
||||
| Trial journal growth per index | Thousands of trials, batched JSONL | High-frequency auto-tuning agent | Compaction via `/hnsw/trials/{idx}/compact` |
|
||||
| Error journal growth | Bounded by ring buffer (2000 events in-memory) + batched JSONL on disk | Continuous failure scenarios | Compaction + retention policy (TODO) |
|
||||
|
||||
---
|
||||
|
||||
## Reference Workloads
|
||||
|
||||
### Workload 1: Staffing Company
|
||||
|
||||
Scale-tested on 128GB RAM server:
|
||||
|
||||
@@ -315,6 +416,27 @@ Scale-tested on 128GB RAM server:
|
||||
- Instant handoff between agents — zero data copy
|
||||
- Full activity timeline preserved across handoffs
|
||||
|
||||
### Workload 2: Local LLM Knowledge Base
|
||||
|
||||
The second use case this substrate is built for. Reference corpus: the running `knowledge_base` Postgres database (586 team runs, response cache history, pipeline runs, threat intel) + LLMS3.com published corpus (~243 enriched documents).
|
||||
|
||||
Target scale on same 128GB server:
|
||||
- Documents: 10K-100K per model profile
|
||||
- Chunks after chunking: 500K-5M per profile
|
||||
- Embedding dimensions: 768 (nomic-embed-text)
|
||||
- Query latency: <100ms semantic search, <3s end-to-end RAG including LLM generation
|
||||
- Concurrent model profiles: 2-5 configured, 1-2 active at a time (VRAM-bound)
|
||||
|
||||
Measured to date (Phase 7 + Phase 16 prep):
|
||||
- 100K candidate-resume chunks embedded in 10 min via Ollama nomic-embed-text
|
||||
- HNSW search at 100% recall, ~1ms p50 on 100K vectors (ec=80 es=30 locked as default)
|
||||
- Trial journal instrumented and working for parameter tuning
|
||||
|
||||
Gaps still to close for this workload:
|
||||
- Model profiles (Phase 17) — today, "model" is a string, not a first-class entity
|
||||
- Hot-swap generations (Phase 16) — today, rebuild = downtime
|
||||
- Scale past 5M vectors — needs Phase 18 Lance evaluation to decide path
|
||||
|
||||
---
|
||||
|
||||
## Available Local Models
|
||||
@@ -331,12 +453,15 @@ Scale-tested on 128GB RAM server:
|
||||
|
||||
## Non-Goals
|
||||
|
||||
- Multi-tenancy (single-owner system)
|
||||
- Cloud deployment (local-first, always)
|
||||
- Full ACID transactions (single-writer model is sufficient)
|
||||
- Real-time streaming / CDC (batch ingest is the model)
|
||||
- Replacing the CRM (this is the analytical layer BEHIND the CRM)
|
||||
- Custom file formats (Parquet is the format, period)
|
||||
- Real-time streaming / CDC (batch ingest is the model; scheduled refresh, not transactional replication)
|
||||
- Replacing the CRM (this is the analytical + AI layer BEHIND the CRM)
|
||||
- Custom file formats — Parquet for datasets + sidecar indexes for vectors (see ADR-018 for why we stayed Parquet instead of migrating to Lance, and the ceilings that choice implies)
|
||||
- Hard multi-tenant isolation (profiles and federation provide soft isolation; this is not a SaaS platform with adversarial tenants — operator is single-trust)
|
||||
|
||||
Removed from prior non-goals (2026-04-16):
|
||||
- ~~Multi-tenancy (single-owner system)~~ — federation + profile buckets are now first-class; soft multi-tenancy is a design goal. Hard adversarial multi-tenancy (adversarial tenants on shared infrastructure) remains out of scope.
|
||||
|
||||
---
|
||||
|
||||
|
||||
+40
-1
@@ -6,7 +6,35 @@ port = 3100
|
||||
|
||||
[storage]
|
||||
root = "./data"
|
||||
# backend = "local" # local | s3 | rustfs
|
||||
profile_root = "./data/_profiles"
|
||||
rescue_bucket = "rescue"
|
||||
|
||||
[[storage.buckets]]
|
||||
name = "primary"
|
||||
backend = "local"
|
||||
root = "./data"
|
||||
|
||||
[[storage.buckets]]
|
||||
name = "rescue"
|
||||
backend = "local"
|
||||
root = "./data/_rescue"
|
||||
|
||||
[[storage.buckets]]
|
||||
name = "testing"
|
||||
backend = "local"
|
||||
root = "./data/_testing"
|
||||
|
||||
# S3 bucket via MinIO. The name "s3:lakehouse" is the convention
|
||||
# lance_backend.rs uses to emit s3:// URIs for Lance datasets.
|
||||
# Credentials resolved via environment (AWS_ACCESS_KEY_ID etc) or
|
||||
# the secrets provider.
|
||||
[[storage.buckets]]
|
||||
name = "s3:lakehouse"
|
||||
backend = "s3"
|
||||
bucket = "lakehouse"
|
||||
endpoint = "http://localhost:9000"
|
||||
region = "us-east-1"
|
||||
secret_ref = "minio-lakehouse"
|
||||
|
||||
[catalog]
|
||||
# Manifests persisted to object storage under this prefix
|
||||
@@ -31,3 +59,14 @@ enabled = false
|
||||
# Export traces to stdout (set to "otlp" for OpenTelemetry collector)
|
||||
exporter = "stdout"
|
||||
service_name = "lakehouse"
|
||||
|
||||
[agent]
|
||||
# Phase 16.2 — background autotune agent. Opt-in: set enabled = true to
|
||||
# let the agent continuously propose + trial HNSW configs and auto-promote
|
||||
# winners. Defaults are conservative so it stays out of the way of live
|
||||
# search traffic on shared Ollama.
|
||||
enabled = true
|
||||
cycle_interval_secs = 120 # periodic wake if no triggers
|
||||
cooldown_between_trials_secs = 10 # min gap between trials
|
||||
min_recall = 0.9 # never promote below this
|
||||
max_trials_per_hour = 20 # hard budget cap
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
"""Admin endpoints: model lifecycle + VRAM introspection.
|
||||
|
||||
Phase 17 / Phase C — the VRAM-aware profile swap story. Ollama loads
|
||||
models lazily and unloads after a TTL (default 5min). For predictable
|
||||
swaps between model profiles, we need explicit control:
|
||||
|
||||
- unload: send keep_alive=0 to force immediate unload
|
||||
- preload: send keep_alive=5m with empty prompt to proactively load
|
||||
- ps: query Ollama's /api/ps to see what's currently loaded
|
||||
|
||||
All three are thin wrappers over Ollama's own API. No state held here.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import shutil
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
from .ollama import client
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
class ModelRequest(BaseModel):
|
||||
model: str
|
||||
|
||||
|
||||
@router.post("/unload")
|
||||
async def unload(req: ModelRequest):
|
||||
"""Force Ollama to unload a model immediately (keep_alive=0 trick)."""
|
||||
async with client() as c:
|
||||
resp = await c.post(
|
||||
"/api/generate",
|
||||
json={"model": req.model, "prompt": "", "keep_alive": 0, "stream": False},
|
||||
)
|
||||
# Ollama returns 200 even on unload; anything else is abnormal.
|
||||
if resp.status_code not in (200,):
|
||||
raise HTTPException(502, f"Ollama unload error: {resp.text}")
|
||||
return {"unloaded": req.model}
|
||||
|
||||
|
||||
@router.post("/preload")
|
||||
async def preload(req: ModelRequest):
|
||||
"""Force Ollama to load a model into VRAM and keep it there."""
|
||||
async with client() as c:
|
||||
resp = await c.post(
|
||||
"/api/generate",
|
||||
json={
|
||||
"model": req.model,
|
||||
# "" can confuse some models; a single token is safer and still ~instant.
|
||||
"prompt": " ",
|
||||
"keep_alive": "5m",
|
||||
"stream": False,
|
||||
"options": {"num_predict": 1},
|
||||
},
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
raise HTTPException(502, f"Ollama preload error: {resp.text}")
|
||||
data = resp.json()
|
||||
return {
|
||||
"preloaded": req.model,
|
||||
"load_duration_ns": data.get("load_duration"),
|
||||
"total_duration_ns": data.get("total_duration"),
|
||||
}
|
||||
|
||||
|
||||
@router.get("/ps")
|
||||
async def list_loaded():
|
||||
"""What models does Ollama currently have in VRAM?"""
|
||||
async with client() as c:
|
||||
resp = await c.get("/api/ps")
|
||||
if resp.status_code != 200:
|
||||
raise HTTPException(502, f"Ollama ps error: {resp.text}")
|
||||
data = resp.json()
|
||||
# Flatten Ollama's fields into something small + useful.
|
||||
models = [
|
||||
{
|
||||
"name": m.get("name"),
|
||||
"size_bytes": m.get("size"),
|
||||
"size_vram_bytes": m.get("size_vram"),
|
||||
"expires_at": m.get("expires_at"),
|
||||
}
|
||||
for m in data.get("models", [])
|
||||
]
|
||||
return {"models": models}
|
||||
|
||||
|
||||
@router.get("/vram")
|
||||
async def vram_summary():
|
||||
"""Combined: nvidia-smi VRAM + Ollama loaded models.
|
||||
|
||||
Shells out to nvidia-smi; if it's not on PATH, returns just the
|
||||
Ollama view. Intentionally async-via-to_thread so the blocking
|
||||
subprocess doesn't stall the event loop.
|
||||
"""
|
||||
gpu = None
|
||||
if shutil.which("nvidia-smi"):
|
||||
gpu = await asyncio.to_thread(_nvidia_smi_snapshot)
|
||||
|
||||
async with client() as c:
|
||||
resp = await c.get("/api/ps")
|
||||
loaded = resp.json().get("models", []) if resp.status_code == 200 else []
|
||||
|
||||
return {
|
||||
"gpu": gpu,
|
||||
"ollama_loaded": [
|
||||
{
|
||||
"name": m.get("name"),
|
||||
"size_vram_mib": (m.get("size_vram", 0) or 0) // (1024 * 1024),
|
||||
"expires_at": m.get("expires_at"),
|
||||
}
|
||||
for m in loaded
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _nvidia_smi_snapshot():
|
||||
"""One-shot nvidia-smi poll. Returns None on failure."""
|
||||
import subprocess
|
||||
try:
|
||||
out = subprocess.check_output(
|
||||
[
|
||||
"nvidia-smi",
|
||||
"--query-gpu=memory.used,memory.total,utilization.gpu,name",
|
||||
"--format=csv,noheader,nounits",
|
||||
],
|
||||
timeout=2,
|
||||
).decode().strip()
|
||||
used_mib, total_mib, util_pct, name = [s.strip() for s in out.split(",")]
|
||||
return {
|
||||
"name": name,
|
||||
"used_mib": int(used_mib),
|
||||
"total_mib": int(total_mib),
|
||||
"utilization_pct": int(util_pct),
|
||||
}
|
||||
except Exception:
|
||||
return None
|
||||
@@ -2,6 +2,7 @@ import os
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
from .admin import router as admin_router
|
||||
from .embed import router as embed_router
|
||||
from .generate import router as generate_router
|
||||
from .rerank import router as rerank_router
|
||||
@@ -11,6 +12,7 @@ app = FastAPI(title="Lakehouse AI Sidecar")
|
||||
app.include_router(embed_router, prefix="/embed", tags=["embed"])
|
||||
app.include_router(generate_router, prefix="/generate", tags=["generate"])
|
||||
app.include_router(rerank_router, prefix="/rerank", tags=["rerank"])
|
||||
app.include_router(admin_router, prefix="/admin", tags=["admin"])
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
|
||||
Reference in New Issue
Block a user