diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index 9aa7acf..fc1d131 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -927,6 +927,9 @@ pub struct ExperimentRecordsResponse { pub has_more: bool, pub limit: usize, pub offset: usize, + /// The data source that was scanned: `"fragments"` (base table only), + /// `"wal"` (pending MemWAL generations only), or `"all"` (the union). + pub source: String, } /// State of a manual or automatic compaction job for one experiment. diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 2018ac1..1d7e2c7 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -46,7 +46,7 @@ pub use record::{ pub use registry::{RegistryEntry, RolloutRegistry}; pub use rollout::{RolloutRecord, ROLE_ARTIFACT, ROLE_ASSISTANT, ROLE_GRADE, ROLE_TOOL}; pub use rollout_store::{ - rollout_schema, RolloutFilters, RolloutObservation, RolloutPage, RolloutStore, + rollout_schema, ListSource, RolloutFilters, RolloutObservation, RolloutPage, RolloutStore, RolloutStoreOptions, }; pub use storage::{create_local_dir_if_needed, join_uri, validate_store_name, MAX_STORE_NAME_LEN}; diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 093f2b8..1a34b42 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -188,6 +188,29 @@ pub struct RolloutPage { pub has_more: bool, } +/// Which data source a rollout list scan reads. +/// +/// A rollout store's rows live in two tiers: the compacted **base table** +/// (`self.dataset`) and the pending **MemWAL** generations that have been +/// flushed but not yet merged into the base table. The default browse path +/// reads only the base table so list latency is independent of WAL backlog +/// (each pending generation is a separate object-store open); callers that need +/// the not-yet-merged tail or full cross-tier consistency opt into `Wal`/`All`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum ListSource { + /// Scan only the base table (`self.dataset`), skipping all MemWAL + /// generations. Fast and bounded; may lag the most recent (un-merged) + /// writes. This is the default. + #[default] + Fragments, + /// Scan only the flushed MemWAL generations (the not-yet-merged tail), + /// excluding the base table. + Wal, + /// Scan the base table unioned with every flushed MemWAL generation — fully + /// consistent, and the behavior of the historical union read path. + All, +} + /// Configuration for opening a [`RolloutStore`]. #[derive(Debug, Clone, Default)] pub struct RolloutStoreOptions { @@ -972,7 +995,21 @@ impl RolloutStore { Ok(records) } - /// Filter and page rollout rows in the LSM execution plan. + /// Filter and page rollout rows over the full base ∪ WAL union. + /// + /// Thin wrapper over [`Self::list_filtered_source`] with [`ListSource::All`], + /// preserving the historical union semantics for existing callers. + pub async fn list_filtered( + &self, + filters: &RolloutFilters, + limit: usize, + offset: usize, + ) -> LanceResult { + self.list_filtered_source(filters, limit, offset, ListSource::All) + .await + } + + /// Filter and page rollout rows from a chosen [`ListSource`]. /// /// Reads one row beyond the requested page to report `has_more`, avoiding /// an unbounded full-table count on every UI request. Pagination is @@ -985,17 +1022,25 @@ impl RolloutStore { /// global limit. Keeping wide token/logprob/metadata columns out of that /// full-source sort makes browsing large rollout tables substantially /// cheaper while preserving the same LSM deduplication semantics. - pub async fn list_filtered( + /// + /// [`ListSource::Fragments`] skips MemWAL manifest discovery entirely, so its + /// latency is independent of how far the merge backlog has grown. + pub async fn list_filtered_source( &self, filters: &RolloutFilters, limit: usize, offset: usize, + source: ListSource, ) -> LanceResult { - let shard_snapshots = self.wal_shard_snapshots().await?; + // Fragments never touches the WAL, so skip the per-shard manifest reads. + let shard_snapshots = match source { + ListSource::Fragments => Vec::new(), + ListSource::Wal | ListSource::All => self.wal_shard_snapshots().await?, + }; let filter = filters.expression(); let mut page_scanner = self - .lsm_scanner_with_snapshots(shard_snapshots.clone()) + .lsm_scanner_for_source(source, shard_snapshots.clone()) .project(&["id"]); if let Some(filter) = &filter { page_scanner = page_scanner.filter(filter)?; @@ -1022,7 +1067,7 @@ impl RolloutStore { let id_refs: Vec<&str> = page_ids.iter().map(String::as_str).collect(); let id_filter = format!("id IN ({})", sql_quoted_list(&id_refs)); let record_scanner = self - .lsm_scanner_with_snapshots(shard_snapshots) + .lsm_scanner_for_source(source, shard_snapshots) .project(&refs) .filter(&id_filter)?; @@ -1184,6 +1229,43 @@ impl RolloutStore { ) } + /// Build a paginating scanner for the requested [`ListSource`], deduplicating + /// by `id`: + /// - `Fragments`: base table only (`shard_snapshots` is ignored — callers + /// pass an empty vec so no manifest reads happen); + /// - `All`: base table ∪ the flushed generations in `shard_snapshots`; + /// - `Wal`: only the flushed generations, via + /// [`LsmScanner::without_base_table`], resolving relative generation paths + /// against the dataset root (matching [`Self::flushed_generation_uri`]). + fn lsm_scanner_for_source( + &self, + source: ListSource, + shard_snapshots: Vec, + ) -> LsmScanner { + match source { + ListSource::Fragments => LsmScanner::new( + Arc::new(self.dataset.clone()), + Vec::new(), + vec!["id".to_string()], + ), + ListSource::All => LsmScanner::new( + Arc::new(self.dataset.clone()), + shard_snapshots, + vec!["id".to_string()], + ), + ListSource::Wal => { + let arrow_schema: Schema = self.dataset.schema().into(); + LsmScanner::without_base_table( + Arc::new(arrow_schema), + self.dataset.uri().trim_end_matches('/').to_string(), + shard_snapshots, + vec!["id".to_string()], + ) + .with_session(self.dataset.session()) + } + } + } + /// Read the latest manifest for every MemWAL shard. Manifest reads are /// bounded-concurrent so stores with many writer instances do not pay one /// object-store round trip per shard serially. @@ -3179,6 +3261,74 @@ mod tests { }); } + #[test] + fn list_source_splits_base_and_wal() { + // Rows appended but not yet merged live only in the WAL, not the base + // table. `Fragments` must omit them; `Wal` must show exactly them; `All` + // is the union. After a merge the rows move to the base table and the WAL + // empties, flipping which source sees them. + 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_with_options( + &uri, + RolloutStoreOptions { + storage_options: None, + shard_id: Some("rollout-0".to_string()), + ..Default::default() + }, + ) + .await + .unwrap(); + + async fn list_ids(store: &RolloutStore, source: ListSource) -> Vec { + let page = store + .list_filtered_source(&RolloutFilters::default(), 25, 0, source) + .await + .unwrap(); + let mut ids: Vec<_> = page.records.iter().map(|r| r.id.clone()).collect(); + ids.sort(); + ids + } + + // Two un-merged rows sit in the WAL only. + store.add(&[assistant_record("g-0")]).await.unwrap(); + store.add(&[assistant_record("g-1")]).await.unwrap(); + assert_eq!(flushed_generation_count(&store).await, 2); + + assert!( + list_ids(&store, ListSource::Fragments).await.is_empty(), + "fragments must not see un-merged WAL rows" + ); + assert_eq!(list_ids(&store, ListSource::Wal).await, vec!["g-0", "g-1"]); + assert_eq!(list_ids(&store, ListSource::All).await, vec!["g-0", "g-1"]); + // The default wrapper preserves the historical union semantics. + let default_page = store + .list_filtered(&RolloutFilters::default(), 25, 0) + .await + .unwrap(); + let mut default_ids: Vec<_> = + default_page.records.iter().map(|r| r.id.clone()).collect(); + default_ids.sort(); + assert_eq!(default_ids, vec!["g-0", "g-1"]); + + // Merge folds the WAL into the base table and drains it. + assert_eq!(store.cleanup_own_shard().await.unwrap(), 2); + assert_eq!(flushed_generation_count(&store).await, 0); + + assert_eq!( + list_ids(&store, ListSource::Fragments).await, + vec!["g-0", "g-1"] + ); + assert!( + list_ids(&store, ListSource::Wal).await.is_empty(), + "wal must be empty after a merge" + ); + assert_eq!(list_ids(&store, ListSource::All).await, vec!["g-0", "g-1"]); + }); + } + /// Reproduces the master data-browser workload at the reported scale and /// compares the former wide-row pagination plan with the late-materialized /// implementation. diff --git a/crates/lance-context-master/src/error.rs b/crates/lance-context-master/src/error.rs index 75d9ace..9275348 100644 --- a/crates/lance-context-master/src/error.rs +++ b/crates/lance-context-master/src/error.rs @@ -9,6 +9,7 @@ use serde_json::json; #[derive(Debug)] pub enum MasterError { NotFound(String), + InvalidRequest(String), Internal(String), } @@ -23,6 +24,7 @@ impl IntoResponse for MasterError { fn into_response(self) -> Response { let (status, msg) = match self { MasterError::NotFound(m) => (StatusCode::NOT_FOUND, m), + MasterError::InvalidRequest(m) => (StatusCode::BAD_REQUEST, m), MasterError::Internal(m) => (StatusCode::INTERNAL_SERVER_ERROR, m), }; (status, Json(json!({ "error": msg }))).into_response() diff --git a/crates/lance-context-master/src/routes.rs b/crates/lance-context-master/src/routes.rs index 2a9146e..34b2e2e 100644 --- a/crates/lance-context-master/src/routes.rs +++ b/crates/lance-context-master/src/routes.rs @@ -15,7 +15,7 @@ use lance_context_api::{ ExperimentRecordsResponse, ExperimentSummary, TaskKind, TaskListResponse, TaskRecord, TaskState, }; -use lance_context_core::{rollout_record_to_dto, RolloutFilters, RolloutStore}; +use lance_context_core::{rollout_record_to_dto, ListSource, RolloutFilters, RolloutStore}; use tokio::sync::RwLock; use crate::error::MasterError; @@ -76,6 +76,32 @@ pub struct RecordListParams { pub artifact_type: Option, #[serde(default)] pub include_in_training: Option, + /// Data source to scan: `fragments` (base table only, the default), `wal` + /// (pending MemWAL generations only), or `all` (the union). Absent/empty + /// defaults to `fragments`. + #[serde(default)] + pub source: Option, +} + +/// Parse the `source` query param into a [`ListSource`]. Absent or empty → +/// `Fragments` (the fast default); an explicitly unknown value is rejected. +fn parse_list_source(raw: Option<&str>) -> Result { + match raw.map(str::trim) { + None | Some("") | Some("fragments") => Ok(ListSource::Fragments), + Some("wal") => Ok(ListSource::Wal), + Some("all") => Ok(ListSource::All), + Some(other) => Err(MasterError::InvalidRequest(format!( + "invalid source '{other}': expected one of fragments, wal, all" + ))), + } +} + +fn list_source_label(source: ListSource) -> &'static str { + match source { + ListSource::Fragments => "fragments", + ListSource::Wal => "wal", + ListSource::All => "all", + } } /// `GET /api/v1/experiments` @@ -142,6 +168,7 @@ pub async fn list_experiment_records( .await .map_err(MasterError::from_lance)?; let limit = params.limit.clamp(1, 100); + let source = parse_list_source(params.source.as_deref())?; let filters = RolloutFilters { id: non_empty(params.id), rollout_id: non_empty(params.rollout_id), @@ -154,7 +181,7 @@ pub async fn list_experiment_records( include_in_training: params.include_in_training, }; let page = store - .list_filtered(&filters, limit, params.offset) + .list_filtered_source(&filters, limit, params.offset, source) .await .map_err(MasterError::from_lance)?; @@ -167,6 +194,7 @@ pub async fn list_experiment_records( has_more: page.has_more, limit, offset: params.offset, + source: list_source_label(source).to_string(), })) } @@ -703,6 +731,7 @@ mod tests { policy_version: Some("policy-a".to_string()), artifact_type: Some("screenshot".to_string()), include_in_training: Some(true), + source: Some("all".to_string()), }), ) .await @@ -711,6 +740,7 @@ mod tests { assert_eq!(page.limit, 1); assert_eq!(page.records[0].id, "artifact-1"); assert!(page.records[0].binary_payload.is_none()); + assert_eq!(page.source, "all"); let missing = list_experiment_records( State(state), @@ -727,12 +757,29 @@ mod tests { policy_version: None, artifact_type: None, include_in_training: None, + source: None, }), ) .await; assert!(matches!(missing, Err(MasterError::NotFound(_)))); } + #[test] + fn parse_list_source_maps_and_defaults() { + assert_eq!(parse_list_source(None).unwrap(), ListSource::Fragments); + assert_eq!(parse_list_source(Some("")).unwrap(), ListSource::Fragments); + assert_eq!( + parse_list_source(Some("fragments")).unwrap(), + ListSource::Fragments + ); + assert_eq!(parse_list_source(Some("wal")).unwrap(), ListSource::Wal); + assert_eq!(parse_list_source(Some("all")).unwrap(), ListSource::All); + assert!(matches!( + parse_list_source(Some("bogus")), + Err(MasterError::InvalidRequest(_)) + )); + } + #[tokio::test] #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn blob_endpoint_sets_download_headers_and_returns_bytes() { diff --git a/crates/lance-context-master/ui/src/App.tsx b/crates/lance-context-master/ui/src/App.tsx index 7c19863..22e8d08 100644 --- a/crates/lance-context-master/ui/src/App.tsx +++ b/crates/lance-context-master/ui/src/App.tsx @@ -20,6 +20,7 @@ import { type CompactJobStatus, type ExperimentSummary, type RecordFilters, + type RecordSource, type RolloutRecord, type TaskRecord, type TaskState, @@ -447,16 +448,23 @@ function RecordsView({ }) { const [draft, setDraft] = useState({ ...EMPTY_RECORD_FILTERS }); const [filters, setFilters] = useState({ ...EMPTY_RECORD_FILTERS }); + const [source, setSource] = useState("fragments"); const [page, setPage] = useState(0); const pageSize = 25; const records = useQuery({ - queryKey: ["experiment-records", name, filters, page], - queryFn: () => listExperimentRecords(name, filters, pageSize, page * pageSize), + queryKey: ["experiment-records", name, filters, source, page], + queryFn: () => + listExperimentRecords(name, filters, pageSize, page * pageSize, source), placeholderData: (previous) => previous, }); const rows = records.data?.records ?? []; const hasMore = records.data?.has_more ?? false; + const selectSource = (next: RecordSource) => { + setSource(next); + setPage(0); + onToggle(null); + }; const setFilter = (key: keyof RecordFilters, value: string) => { setDraft((current) => ({ ...current, [key]: value })); }; @@ -474,8 +482,51 @@ function RecordsView({ }; const hasFilters = Object.values(filters).some(Boolean); + const sourceTabs: { key: RecordSource; label: string; title: string }[] = [ + { + key: "fragments", + label: "Fragments", + title: + "Base table only. Fast and independent of WAL backlog, but may lag the most recent (un-merged) writes.", + }, + { + key: "wal", + label: "WAL", + title: "Pending MemWAL generations only — the not-yet-merged tail.", + }, + { + key: "all", + label: "All", + title: "Base table unioned with pending WAL (fully consistent).", + }, + ]; + return (
+
+ {sourceTabs.map((t) => ( + + ))} + {source === "fragments" && ( + + base table only — may lag un-merged writes + + )} +