diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index f29cd1e..24301a9 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -928,6 +928,10 @@ pub enum TaskKind { /// task fans out to the configured worker endpoints and each worker merges /// its own shard. MergeWal, + /// Build a ZoneMap scalar index on the experiment base table's `id` column. + /// Runs on the master; serialized per experiment against `Compact` because + /// both mutate the shared base table. + IndexId, } /// Lifecycle state of a scheduled task, generalized from [`CompactJobStatus`] @@ -1123,4 +1127,17 @@ mod tests { assert_eq!(req.kind, TaskKind::Compact); assert_eq!(req.target, "exp-7"); } + + #[test] + fn task_kind_index_id_roundtrips_snake_case() { + assert_eq!( + serde_json::to_string(&TaskKind::IndexId).unwrap(), + r#""index_id""# + ); + let back: TaskKind = serde_json::from_str(r#""index_id""#).unwrap(); + assert_eq!(back, TaskKind::IndexId); + let req: EnqueueTaskRequest = + serde_json::from_str(r#"{"kind":"index_id","target":"exp-9"}"#).unwrap(); + assert_eq!(req.kind, TaskKind::IndexId); + } } diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index e387a7d..1791fd9 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -70,6 +70,8 @@ use lance::index::DatasetIndexExt; use lance::io::{ObjectStoreParams, StorageOptionsAccessor}; use lance::{Error as LanceError, Result as LanceResult}; use lance_index::mem_wal::{ShardManifest, MEM_WAL_INDEX_NAME}; +use lance_index::scalar::ScalarIndexParams; +use lance_index::IndexType; use tokio::task::JoinHandle; use tracing::{info, warn}; use uuid::Uuid; @@ -89,6 +91,10 @@ const DEFAULT_MANIFEST_SCAN_BATCH_SIZE: usize = 16; /// concurrently while collecting observability metrics. const DEFAULT_OBSERVE_CONCURRENCY: usize = 16; +/// Name of the scalar index on the base table's `id` column. Kept in sync with +/// `ContextStore`'s `ID_INDEX_NAME` so both tables index `id` under one name. +const ROLLOUT_ID_INDEX_NAME: &str = "id_idx"; + /// Read-only observability snapshot of a rollout store. /// /// Produced by [`RolloutStore::observe`] from base-table and MemWAL metadata. @@ -681,6 +687,37 @@ impl RolloutStore { } } + /// Build a ZoneMap scalar index on the base table's `id` column. + /// + /// `id` is the rollout table's (unenforced) primary key, so a lightweight + /// per-fragment min/max index accelerates id point-lookups and range scans + /// on the already-flushed base table. `replace(true)` makes this idempotent: + /// re-running simply rebuilds the index in place. + /// + /// # MemWAL interaction + /// + /// The rollout base table carries a fieldless MemWAL index, and Lance's + /// MemWAL does not *maintain* ZoneMap indices across WAL flushes (it only + /// keeps the indices named in `maintained_indexes`). That does not affect + /// correctness here: rollout rows are immutable and de-duplicated by `id` at + /// read time, so the ZoneMap only ever needs to describe the base table's + /// already-merged fragments — rows still living in unmerged WAL generations + /// are found by the normal full scan of those generations. Creating the + /// index is therefore safe alongside an existing MemWAL index. + pub async fn create_id_zonemap_index(&mut self) -> LanceResult<()> { + info!("Creating ZoneMap index on rollout id column"); + self.dataset + .create_index_builder(&["id"], IndexType::ZoneMap, &ScalarIndexParams::default()) + .name(ROLLOUT_ID_INDEX_NAME.to_string()) + .replace(true) + .await?; + // Reload the handle so subsequent reads on this instance observe the + // new index (mirrors the reload done after `compact`). + let uri = self.dataset.uri().to_string(); + self.dataset = Self::load_with_options(&uri, self.storage_options.clone()).await?; + Ok(()) + } + /// Whether the base table has accumulated at least `min_fragments` /// fragments (and is thus worth compacting). Quiet-hours gating from /// [`CompactionConfig`] is honored so an external scheduler can pass the @@ -2369,6 +2406,45 @@ mod tests { }); } + #[test] + fn create_id_zonemap_index_builds_and_is_idempotent() { + // Building the ZoneMap index on `id` must succeed even though the + // rollout table also carries a (fieldless) MemWAL index, and calling it + // twice must not error (replace(true) rebuilds in place). + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = RolloutStore::open(&uri).await.unwrap(); + // Add rows and fold them into the base table so there is data (and a + // MemWAL index) present when we build the scalar index. + store.add(&[assistant_record("a-0")]).await.unwrap(); + store.add(&[assistant_record("a-1")]).await.unwrap(); + store.cleanup_own_shard().await.unwrap(); + + store.create_id_zonemap_index().await.unwrap(); + let has_id_index = |s: &RolloutStore| { + let dataset = s.dataset.clone(); + async move { + dataset + .load_indices() + .await + .unwrap() + .iter() + .any(|i| i.name == ROLLOUT_ID_INDEX_NAME) + } + }; + assert!(has_id_index(&store).await, "id index should exist"); + + // Idempotent: a second build replaces in place without erroring. + store.create_id_zonemap_index().await.unwrap(); + assert!(has_id_index(&store).await, "id index should still exist"); + + // Rows remain readable exactly once after indexing. + assert_eq!(store.list(None, None).await.unwrap().len(), 2); + }); + } + #[test] fn spawn_periodic_cleanup_reclaims_on_a_timer() { // The background timer folds accumulated generations into the base table diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index dd0a82f..91c0c36 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -5,14 +5,18 @@ //! //! - **Compact** — rewrites an experiment's base-table fragments. Two `Rewrite`s //! on the *same* dataset conflict in Lance's conflict matrix, so compaction of -//! one experiment is serialized against itself via a per-name in-flight gate -//! ([`MasterState::inflight_compacts`]). Distinct experiments compact -//! concurrently. `Rewrite` vs the data-plane's `Append` is non-conflicting, so -//! this runs safely alongside live ingest. +//! one experiment is serialized against itself (and against `IndexId`) via a +//! per-name in-flight gate ([`MasterState::inflight_dataset_writes`]). Distinct +//! experiments compact concurrently. `Rewrite` vs the data-plane's `Append` is +//! non-conflicting, so this runs safely alongside live ingest. //! - **MergeWal** — folds flushed MemWAL generations back into the base table. //! The master cannot do this itself without fencing the live shard writer, so //! it fans out to every configured worker endpoint and each worker merges its //! own shard (`POST /api/v1/internal/merge-wal/{name}`). +//! - **IndexId** — builds a ZoneMap scalar index on the base table's `id` column +//! (runs locally on the master). It commits a `CreateIndex`, which can conflict +//! with a concurrent `Compact` `Rewrite` on the same dataset, so it shares the +//! per-name in-flight gate with `Compact`. //! //! Each task runs in its own `tokio::spawn`, so one task's failure never affects //! another. A global [`Semaphore`] bounds how many run at once. @@ -48,18 +52,20 @@ fn kind_label(kind: TaskKind) -> &'static str { match kind { TaskKind::Compact => "compact", TaskKind::MergeWal => "merge_wal", + TaskKind::IndexId => "index_id", } } -/// Enqueue a task and return its record. For [`TaskKind::Compact`] this de-dupes -/// against an existing non-terminal task for the same target: if one is already -/// `Queued` or `Running`, its record is returned unchanged and nothing new is -/// enqueued. `MergeWal` tasks are not de-duped (each fan-out is independent). +/// Enqueue a task and return its record. For [`TaskKind::Compact`] and +/// [`TaskKind::IndexId`] this de-dupes against an existing non-terminal task for +/// the same target: if one is already `Queued` or `Running`, its record is +/// returned unchanged and nothing new is enqueued. `MergeWal` tasks are not +/// de-duped (each fan-out is independent). pub async fn enqueue(state: &Arc, kind: TaskKind, target: &str) -> TaskRecord { - if matches!(kind, TaskKind::Compact) { + if matches!(kind, TaskKind::Compact | TaskKind::IndexId) { let tasks = state.tasks.lock().await; if let Some(existing) = tasks.values().find(|t| { - t.kind == TaskKind::Compact + t.kind == kind && t.target == target && matches!(t.state, TaskState::Queued | TaskState::Running) }) { @@ -115,6 +121,7 @@ async fn run_task(state: &Arc, id: String) { let outcome = match task.kind { TaskKind::Compact => run_compaction(state, &task.target).await, TaskKind::MergeWal => run_merge_wal(state, &task.target).await, + TaskKind::IndexId => run_index_id(state, &task.target).await, }; metrics::histogram!("master_task_duration_seconds", "kind" => kind_label(task.kind)) @@ -140,25 +147,59 @@ async fn run_task(state: &Arc, id: String) { .await; } -/// Compact one experiment. Serialized per experiment via the in-flight gate so -/// two `Rewrite`s on the same dataset never race. Returns a human-readable -/// outcome summary on success. +/// Compact one experiment. Serialized per experiment via the shared base-table +/// write gate so two `Rewrite`s (or a `Rewrite` racing a `CreateIndex`) on the +/// same dataset never conflict. Returns a human-readable outcome summary on +/// success. async fn run_compaction(state: &Arc, name: &str) -> Result { - // Per-name serial gate: refuse if a compaction for this experiment is - // already running. (De-dup at enqueue makes this rare, but the gate is the - // real guarantee against concurrent Rewrites.) + // Per-name serial gate: refuse if another base-table write for this + // experiment is already running. (De-dup at enqueue makes this rare, but + // the gate is the real guarantee against conflicting commits.) { - let mut inflight = state.inflight_compacts.lock().await; + let mut inflight = state.inflight_dataset_writes.lock().await; if !inflight.insert(name.to_string()) { - return Err(format!("a compaction for '{name}' is already in progress")); + return Err(format!( + "a base-table write for '{name}' is already in progress" + )); } } // Ensure the gate is released no matter how the inner call exits. let result = compact_inner(state, name).await; - state.inflight_compacts.lock().await.remove(name); + state.inflight_dataset_writes.lock().await.remove(name); result } +/// Build a ZoneMap scalar index on one experiment's `id` column. Shares the +/// per-name base-table write gate with [`run_compaction`] so an `IndexId` and a +/// `Compact` for the same experiment never commit concurrently (`CreateIndex` +/// vs `Rewrite` can conflict). Distinct experiments index concurrently. +async fn run_index_id(state: &Arc, name: &str) -> Result { + { + let mut inflight = state.inflight_dataset_writes.lock().await; + if !inflight.insert(name.to_string()) { + return Err(format!( + "a base-table write for '{name}' is already in progress" + )); + } + } + let result = index_id_inner(state, name).await; + state.inflight_dataset_writes.lock().await.remove(name); + result +} + +async fn index_id_inner(state: &Arc, name: &str) -> Result { + let uri = state.rollout_uri(name); + let opts = RolloutStoreOptions::default(); + let mut store = RolloutStore::open_existing_with_options(&uri, opts) + .await + .map_err(|e| e.to_string())?; + store + .create_id_zonemap_index() + .await + .map_err(|e| e.to_string())?; + Ok("built zonemap index on id".to_string()) +} + async fn compact_inner(state: &Arc, name: &str) -> Result { let uri = state.rollout_uri(name); let config = compaction_config(state); @@ -428,6 +469,40 @@ mod tests { worker.abort(); } + /// Manual enqueue of an `IndexId` task -> dispatcher builds the ZoneMap + /// index -> task reaches Done with the expected detail summary. + #[tokio::test] + async fn index_id_task_builds_index_and_reaches_done() { + let dir = TempDir::new().unwrap(); + let state = MasterState::new(config(&dir)).await.unwrap(); + let worker = spawn_scheduler(&state); + + let name = "exp"; + let uri = state.rollout_uri(name); + { + let mut store = RolloutStore::open(&uri).await.unwrap(); + for i in 0..3 { + let rec = rollout_record(&format!("r{i}")); + store.add(&[rec]).await.unwrap(); + store.cleanup_own_shard().await.unwrap(); + } + } + state + .registry + .write() + .await + .upsert(name, &uri) + .await + .unwrap(); + + let rec = enqueue(&state, TaskKind::IndexId, name).await; + let status = await_terminal(&state, &rec.id).await; + assert_eq!(status.state, TaskState::Done, "got {status:?}"); + assert_eq!(status.detail.as_deref(), Some("built zonemap index on id")); + + worker.abort(); + } + /// Enqueuing the same experiment twice while queued de-dupes to one task. #[tokio::test] async fn enqueue_dedupes_queued_compactions() { diff --git a/crates/lance-context-master/src/state.rs b/crates/lance-context-master/src/state.rs index 435fb1b..4a0b04b 100644 --- a/crates/lance-context-master/src/state.rs +++ b/crates/lance-context-master/src/state.rs @@ -36,10 +36,12 @@ pub struct MasterState { /// Every task (queued / running / terminal) keyed by task id, for the queue /// API and UI. Retained after completion so the UI can show recent history. pub tasks: Mutex>, - /// Experiment names with a compaction currently executing — the per-name - /// serial gate. A `Compact` task must not run while another `Compact` on the - /// same experiment is in flight (two `Rewrite`s on one dataset conflict). - pub inflight_compacts: Mutex>, + /// Experiment names with a base-table write currently executing — the + /// per-name serial gate. A `Compact` or `IndexId` task must not run while + /// another base-table-mutating task on the same experiment is in flight: + /// two `Rewrite`s (or a `Rewrite` and a `CreateIndex`) on one dataset can + /// conflict in Lance's commit matrix, forcing wasteful retries. + pub inflight_dataset_writes: Mutex>, /// Shared HTTP client for fanning `MergeWal` tasks out to worker endpoints. pub http: reqwest::Client, } @@ -68,7 +70,7 @@ impl MasterState { task_tx, task_rx: Mutex::new(Some(task_rx)), tasks: Mutex::new(HashMap::new()), - inflight_compacts: Mutex::new(HashSet::new()), + inflight_dataset_writes: Mutex::new(HashSet::new()), http: reqwest::Client::new(), })) } diff --git a/crates/lance-context-master/ui/src/App.tsx b/crates/lance-context-master/ui/src/App.tsx index c7fce4a..8d25ffa 100644 --- a/crates/lance-context-master/ui/src/App.tsx +++ b/crates/lance-context-master/ui/src/App.tsx @@ -654,6 +654,7 @@ function Drawer({ name }: { name: string }) {
+
)} @@ -727,6 +728,28 @@ function MergeWalButton({ name }: { name: string }) { ); } +/** "Index id" trigger — enqueues an index_id task (build ZoneMap on id). */ +function IndexIdButton({ name }: { name: string }) { + const qc = useQueryClient(); + const setView = useUiStore((s) => s.setView); + const trigger = useMutation({ + mutationFn: () => enqueueTask({ kind: "index_id", target: name }), + onSuccess: () => { + qc.invalidateQueries({ queryKey: ["tasks"] }); + setView("tasks"); + }, + }); + return ( + + ); +} + function TaskQueue() { const tasks = useQuery({ queryKey: ["tasks"], diff --git a/crates/lance-context-master/ui/src/api.ts b/crates/lance-context-master/ui/src/api.ts index d397e07..c344a6c 100644 --- a/crates/lance-context-master/ui/src/api.ts +++ b/crates/lance-context-master/ui/src/api.ts @@ -83,7 +83,7 @@ export type CompactJobStatus = | { state: "none" }; // Unified scheduler task DTOs (mirror `TaskRecord` etc. in `lance-context-api`). -export type TaskKind = "compact" | "merge_wal"; +export type TaskKind = "compact" | "merge_wal" | "index_id"; export type TaskState = "queued" | "running" | "done" | "failed"; export interface TaskRecord {