diff --git a/crates/lance-context-core/src/datagen.rs b/crates/lance-context-core/src/datagen.rs new file mode 100644 index 0000000..8014353 --- /dev/null +++ b/crates/lance-context-core/src/datagen.rs @@ -0,0 +1,599 @@ +use std::collections::{BTreeMap, BTreeSet, HashMap}; + +use chrono::{DateTime, Utc}; +use serde_json::Value; +use uuid::Uuid; + +/// Current schema version for the append-only datagen checkpoint log. +pub const DATAGEN_SCHEMA_VERSION: i32 = 1; + +/// One lifecycle or field-level event in a datagen item's checkpoint history. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DatagenEventType { + ItemCreated, + FieldSet, + FieldAppend, + StepCompleted, + Failed, + Terminal, +} + +impl DatagenEventType { + #[must_use] + pub fn as_str(self) -> &'static str { + match self { + Self::ItemCreated => "ITEM_CREATED", + Self::FieldSet => "FIELD_SET", + Self::FieldAppend => "FIELD_APPEND", + Self::StepCompleted => "STEP_COMPLETED", + Self::Failed => "FAILED", + Self::Terminal => "TERMINAL", + } + } + + pub fn parse(value: &str) -> Result { + match value { + "ITEM_CREATED" => Ok(Self::ItemCreated), + "FIELD_SET" => Ok(Self::FieldSet), + "FIELD_APPEND" => Ok(Self::FieldAppend), + "STEP_COMPLETED" => Ok(Self::StepCompleted), + "FAILED" => Ok(Self::Failed), + "TERMINAL" => Ok(Self::Terminal), + other => Err(format!("unsupported datagen event type '{other}'")), + } + } +} + +/// Terminal outcome of an item. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DatagenTerminal { + Completed, + Filtered, +} + +impl DatagenTerminal { + #[must_use] + pub fn as_str(self) -> &'static str { + match self { + Self::Completed => "completed", + Self::Filtered => "filtered", + } + } + + pub fn parse(value: &str) -> Result { + match value { + "completed" => Ok(Self::Completed), + "filtered" => Ok(Self::Filtered), + other => Err(format!("unsupported datagen terminal value '{other}'")), + } + } +} + +/// Current status derived exclusively by folding the event log. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DatagenItemStatus { + Pending, + Running, + Completed, + Filtered, + Failed, +} + +/// Lazy reference to an inline blob event. `bytes` is absent on normal fold and +/// trajectory reads; callers materialize it through `DatagenStore::get_blob`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DatagenBlobValue { + pub bytes: Option>, + pub size: i64, + pub checksum: Option, +} + +/// Canonical value stored by a FIELD_SET or FIELD_APPEND event. +#[derive(Debug, Clone, PartialEq)] +pub enum DatagenValue { + Int(i64), + Float(f64), + Bool(bool), + String(String), + Json(Value), + Blob(DatagenBlobValue), +} + +impl DatagenValue { + #[must_use] + pub fn kind(&self) -> &'static str { + match self { + Self::Int(_) => "int", + Self::Float(_) => "float", + Self::Bool(_) => "bool", + Self::String(_) => "str", + Self::Json(_) => "json", + Self::Blob(_) => "blob", + } + } +} + +/// A single append-only row in `log.lance`. +#[derive(Debug, Clone, PartialEq)] +pub struct DatagenEvent { + /// Deterministic idempotency key. The MemWAL read path de-duplicates by it. + pub event_id: String, + pub item_id: String, + pub root_item_id: String, + pub parent_item_id: Option, + /// Strictly increasing per item. A collision between different event ids is + /// treated as split-brain corruption during fold. + pub item_seq: i64, + /// Shared by every event emitted for one checkpoint boundary. + pub checkpoint_id: String, + pub event_type: DatagenEventType, + pub step_name: Option, + pub step_index: Option, + pub step_instance_id: Option, + pub iteration: Option, + pub attempt: i32, + pub run_id: String, + /// Fencing identity for the writer/lease that owned this item. + pub writer_epoch: String, + pub field_name: Option, + /// Stable codec id, not a Python class name. + pub field_type: Option, + pub codec_version: Option, + pub value: Option, + /// Query tags captured on ITEM_CREATED. They are not part of correctness. + pub query_tags: Option, + pub terminal: Option, + pub error_type: Option, + pub error_dump: Option, + pub traceback: Option, + pub event_ts: DateTime, + pub schema_version: i32, +} + +impl DatagenEvent { + pub fn validate(&self) -> Result<(), String> { + for (name, value) in [ + ("event_id", self.event_id.as_str()), + ("item_id", self.item_id.as_str()), + ("root_item_id", self.root_item_id.as_str()), + ("checkpoint_id", self.checkpoint_id.as_str()), + ("run_id", self.run_id.as_str()), + ("writer_epoch", self.writer_epoch.as_str()), + ] { + if value.is_empty() { + return Err(format!("{name} must not be empty")); + } + } + if self.item_seq < 0 { + return Err("item_seq must be non-negative".to_string()); + } + if self.attempt < 0 { + return Err("attempt must be non-negative".to_string()); + } + if self.schema_version <= 0 { + return Err("schema_version must be positive".to_string()); + } + + match self.event_type { + DatagenEventType::FieldSet | DatagenEventType::FieldAppend => { + if self.field_name.as_deref().is_none_or(str::is_empty) { + return Err("field events require field_name".to_string()); + } + if self.field_type.as_deref().is_none_or(str::is_empty) { + return Err("field events require field_type".to_string()); + } + if self.codec_version.is_none() { + return Err("field events require codec_version".to_string()); + } + if self.value.is_none() { + return Err("field events require a value".to_string()); + } + if self.step_name.as_deref().is_none_or(str::is_empty) + || self.step_index.is_none() + || self.step_instance_id.as_deref().is_none_or(str::is_empty) + { + return Err( + "field events require step_name, step_index, and step_instance_id" + .to_string(), + ); + } + } + DatagenEventType::StepCompleted => { + if self.step_name.as_deref().is_none_or(str::is_empty) + || self.step_index.is_none() + || self.step_instance_id.as_deref().is_none_or(str::is_empty) + { + return Err( + "STEP_COMPLETED requires step_name, step_index, and step_instance_id" + .to_string(), + ); + } + } + DatagenEventType::Failed => { + if self.error_type.as_deref().is_none_or(str::is_empty) { + return Err("FAILED requires error_type".to_string()); + } + } + DatagenEventType::Terminal => { + if self.terminal.is_none() { + return Err("TERMINAL requires terminal".to_string()); + } + } + DatagenEventType::ItemCreated => {} + } + Ok(()) + } +} + +/// Generate a deterministic event id for retry-safe checkpoint ingestion. +#[must_use] +pub fn datagen_event_id(item_id: &str, checkpoint_id: &str, ordinal: u32) -> String { + let input = format!("{item_id}\0{checkpoint_id}\0{ordinal}"); + Uuid::new_v5(&Uuid::NAMESPACE_OID, input.as_bytes()).to_string() +} + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct DatagenStepCursor { + pub checkpoint_id: String, + pub step_name: String, + pub step_index: i64, + pub step_instance_id: String, + pub iteration: Option, + pub attempt: i32, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum DatagenFieldState { + Set(DatagenValue), + Appended(Vec), +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DatagenFailure { + pub event_id: String, + pub run_id: String, + pub checkpoint_id: String, + pub item_seq: i64, + pub step_name: Option, + pub error_type: String, + pub error_dump: Option, + pub traceback: Option, + pub failed_at: DateTime, +} + +/// Current item state reconstructed solely from the append-only log. +#[derive(Debug, Clone, PartialEq)] +pub struct FoldedDatagenItem { + pub item_id: String, + pub root_item_id: String, + pub parent_item_id: Option, + pub fields: BTreeMap, + pub completed_steps: BTreeSet, + pub status: DatagenItemStatus, + pub terminal: Option, + pub failure: Option, + pub query_tags: Option, + pub current_run_id: String, + pub last_item_seq: i64, + pub last_checkpoint_id: String, +} + +/// State captured immediately after a STEP_COMPLETED event. +#[derive(Debug, Clone, PartialEq)] +pub struct DatagenTrajectoryPoint { + pub cursor: DatagenStepCursor, + pub item: FoldedDatagenItem, +} + +pub fn fold_datagen_events(events: &[DatagenEvent]) -> Result { + let ordered = normalize_events(events)?; + let first = ordered + .first() + .ok_or_else(|| "cannot fold an empty datagen event list".to_string())?; + let mut item = initial_item(first); + for event in ordered { + apply_event(&mut item, event)?; + } + Ok(item) +} + +pub fn datagen_trajectory(events: &[DatagenEvent]) -> Result, String> { + let ordered = normalize_events(events)?; + let first = ordered + .first() + .ok_or_else(|| "cannot build a trajectory from an empty event list".to_string())?; + let mut item = initial_item(first); + let mut trajectory = Vec::new(); + for event in ordered { + apply_event(&mut item, event)?; + if event.event_type == DatagenEventType::StepCompleted { + let cursor = step_cursor(event)?; + trajectory.push(DatagenTrajectoryPoint { + cursor, + item: item.clone(), + }); + } + } + Ok(trajectory) +} + +fn normalize_events(events: &[DatagenEvent]) -> Result, String> { + let mut by_id: HashMap<&str, &DatagenEvent> = HashMap::new(); + for event in events { + event.validate()?; + match by_id.insert(&event.event_id, event) { + Some(previous) if previous != event => { + return Err(format!( + "event_id '{}' was reused with different content", + event.event_id + )); + } + _ => {} + } + } + + let mut ordered: Vec<_> = by_id.into_values().collect(); + ordered.sort_by(|left, right| { + left.item_seq + .cmp(&right.item_seq) + .then_with(|| left.event_id.cmp(&right.event_id)) + }); + for pair in ordered.windows(2) { + if pair[0].item_id != pair[1].item_id { + return Err("all events in a fold must belong to one item".to_string()); + } + if pair[0].item_seq == pair[1].item_seq { + return Err(format!( + "item '{}' has conflicting events at item_seq {}", + pair[0].item_id, pair[0].item_seq + )); + } + } + Ok(ordered) +} + +fn initial_item(first: &DatagenEvent) -> FoldedDatagenItem { + FoldedDatagenItem { + item_id: first.item_id.clone(), + root_item_id: first.root_item_id.clone(), + parent_item_id: first.parent_item_id.clone(), + fields: BTreeMap::new(), + completed_steps: BTreeSet::new(), + status: DatagenItemStatus::Pending, + terminal: None, + failure: None, + query_tags: None, + current_run_id: first.run_id.clone(), + last_item_seq: first.item_seq, + last_checkpoint_id: first.checkpoint_id.clone(), + } +} + +fn apply_event(item: &mut FoldedDatagenItem, event: &DatagenEvent) -> Result<(), String> { + if event.item_id != item.item_id { + return Err(format!( + "event '{}' belongs to item '{}', expected '{}'", + event.event_id, event.item_id, item.item_id + )); + } + if event.root_item_id != item.root_item_id { + return Err(format!( + "item '{}' changed root_item_id from '{}' to '{}'", + item.item_id, item.root_item_id, event.root_item_id + )); + } + if event.parent_item_id != item.parent_item_id { + return Err(format!( + "item '{}' changed parent_item_id during its trajectory", + item.item_id + )); + } + + item.current_run_id = event.run_id.clone(); + item.last_item_seq = event.item_seq; + item.last_checkpoint_id = event.checkpoint_id.clone(); + + match event.event_type { + DatagenEventType::ItemCreated => { + item.status = DatagenItemStatus::Pending; + item.terminal = None; + item.failure = None; + if event.query_tags.is_some() { + item.query_tags = event.query_tags.clone(); + } + } + DatagenEventType::FieldSet => { + let field_name = event.field_name.clone().unwrap(); + item.fields.insert( + field_name, + DatagenFieldState::Set(event.value.clone().unwrap()), + ); + item.status = DatagenItemStatus::Running; + item.terminal = None; + item.failure = None; + } + DatagenEventType::FieldAppend => { + let field_name = event.field_name.clone().unwrap(); + let value = event.value.clone().unwrap(); + match item.fields.entry(field_name) { + std::collections::btree_map::Entry::Vacant(entry) => { + entry.insert(DatagenFieldState::Appended(vec![value])); + } + std::collections::btree_map::Entry::Occupied(mut entry) => match entry.get_mut() { + DatagenFieldState::Appended(values) => values.push(value), + DatagenFieldState::Set(_) => { + return Err(format!( + "field '{}' mixes FIELD_SET and FIELD_APPEND", + entry.key() + )); + } + }, + } + item.status = DatagenItemStatus::Running; + item.terminal = None; + item.failure = None; + } + DatagenEventType::StepCompleted => { + item.completed_steps.insert(step_cursor(event)?); + item.status = DatagenItemStatus::Running; + item.terminal = None; + item.failure = None; + } + DatagenEventType::Failed => { + item.status = DatagenItemStatus::Failed; + item.terminal = None; + item.failure = Some(DatagenFailure { + event_id: event.event_id.clone(), + run_id: event.run_id.clone(), + checkpoint_id: event.checkpoint_id.clone(), + item_seq: event.item_seq, + step_name: event.step_name.clone(), + error_type: event.error_type.clone().unwrap(), + error_dump: event.error_dump.clone(), + traceback: event.traceback.clone(), + failed_at: event.event_ts, + }); + } + DatagenEventType::Terminal => { + let terminal = event.terminal.unwrap(); + item.status = match terminal { + DatagenTerminal::Completed => DatagenItemStatus::Completed, + DatagenTerminal::Filtered => DatagenItemStatus::Filtered, + }; + item.terminal = Some(terminal); + item.failure = None; + } + } + Ok(()) +} + +fn step_cursor(event: &DatagenEvent) -> Result { + Ok(DatagenStepCursor { + checkpoint_id: event.checkpoint_id.clone(), + step_name: event + .step_name + .clone() + .ok_or_else(|| "STEP_COMPLETED missing step_name".to_string())?, + step_index: event + .step_index + .ok_or_else(|| "STEP_COMPLETED missing step_index".to_string())?, + step_instance_id: event + .step_instance_id + .clone() + .ok_or_else(|| "STEP_COMPLETED missing step_instance_id".to_string())?, + iteration: event.iteration, + attempt: event.attempt, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::TimeZone; + use serde_json::json; + + fn event(seq: i64, event_type: DatagenEventType) -> DatagenEvent { + let checkpoint_id = format!("checkpoint-{seq}"); + DatagenEvent { + event_id: datagen_event_id("item-1", &checkpoint_id, 0), + item_id: "item-1".to_string(), + root_item_id: "item-1".to_string(), + parent_item_id: None, + item_seq: seq, + checkpoint_id, + event_type, + step_name: None, + step_index: None, + step_instance_id: None, + iteration: None, + attempt: 0, + run_id: "run-1".to_string(), + writer_epoch: "writer-1".to_string(), + field_name: None, + field_type: None, + codec_version: None, + value: None, + query_tags: None, + terminal: None, + error_type: None, + error_dump: None, + traceback: None, + event_ts: Utc.timestamp_micros(1_700_000_000_000_000 + seq).unwrap(), + schema_version: DATAGEN_SCHEMA_VERSION, + } + } + + fn completed_step(seq: i64) -> DatagenEvent { + let mut event = event(seq, DatagenEventType::StepCompleted); + event.step_name = Some("noop".to_string()); + event.step_index = Some(3); + event.step_instance_id = Some("loop/2/noop".to_string()); + event.iteration = Some(2); + event + } + + #[test] + fn no_op_step_is_present_in_fold_and_trajectory() { + let created = event(0, DatagenEventType::ItemCreated); + let completed = completed_step(1); + + let folded = fold_datagen_events(&[completed.clone(), created]).unwrap(); + assert_eq!(folded.status, DatagenItemStatus::Running); + assert_eq!(folded.completed_steps.len(), 1); + assert!(folded.fields.is_empty()); + + let trajectory = datagen_trajectory(&[completed]).unwrap(); + assert_eq!(trajectory.len(), 1); + assert_eq!(trajectory[0].cursor.step_instance_id, "loop/2/noop"); + } + + #[test] + fn retry_duplicate_event_is_folded_once() { + let created = event(0, DatagenEventType::ItemCreated); + let mut append = event(1, DatagenEventType::FieldAppend); + append.field_name = Some("messages".to_string()); + append.field_type = Some("json".to_string()); + append.codec_version = Some(1); + append.value = Some(DatagenValue::Json(json!({"role": "assistant"}))); + append.step_name = Some("generate".to_string()); + append.step_index = Some(1); + append.step_instance_id = Some("generate/0".to_string()); + + let folded = fold_datagen_events(&[created, append.clone(), append.clone()]).unwrap(); + assert_eq!( + folded.fields.get("messages"), + Some(&DatagenFieldState::Appended(vec![DatagenValue::Json( + json!({"role": "assistant"}) + )])) + ); + } + + #[test] + fn sequence_collision_is_rejected() { + let created = event(0, DatagenEventType::ItemCreated); + let first = completed_step(1); + let mut second = first.clone(); + second.event_id = "different-event".to_string(); + second.checkpoint_id = "different-checkpoint".to_string(); + + let error = fold_datagen_events(&[created, first, second]).unwrap_err(); + assert!(error.contains("conflicting events at item_seq 1")); + } + + #[test] + fn later_run_can_supersede_a_failure() { + let created = event(0, DatagenEventType::ItemCreated); + let mut failed = event(1, DatagenEventType::Failed); + failed.error_type = Some("RuntimeError".to_string()); + + let mut retried = event(2, DatagenEventType::ItemCreated); + retried.run_id = "run-2".to_string(); + retried.checkpoint_id = "retry-created".to_string(); + retried.event_id = datagen_event_id("item-1", "retry-created", 0); + + let folded = fold_datagen_events(&[created, failed, retried]).unwrap(); + assert_eq!(folded.status, DatagenItemStatus::Pending); + assert_eq!(folded.current_run_id, "run-2"); + assert!(folded.failure.is_none()); + } +} diff --git a/crates/lance-context-core/src/datagen_store.rs b/crates/lance-context-core/src/datagen_store.rs new file mode 100644 index 0000000..228a3ef --- /dev/null +++ b/crates/lance-context-core/src/datagen_store.rs @@ -0,0 +1,1419 @@ +//! Lance-backed append-only checkpoint log for datagen pipelines. +//! +//! A [`DatagenStore`] owns exactly one Lance dataset. Item state, failures, and +//! trajectories are all derived by folding immutable events from that dataset; +//! there are no cross-dataset writes or correctness-critical projections. +//! +//! Concurrent writers append through per-instance MemWAL shards. Reads union +//! the base table with every flushed shard and de-duplicate by deterministic +//! `event_id`, making a retried checkpoint batch idempotent. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use arrow_array::builder::{ + BooleanBuilder, Float64Builder, Int32Builder, Int64Builder, LargeBinaryBuilder, + LargeStringBuilder, StringBuilder, TimestampMicrosecondBuilder, +}; +use arrow_array::{ + Array, ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, LargeBinaryArray, + LargeStringArray, RecordBatch, RecordBatchIterator, StringArray, TimestampMicrosecondArray, + UInt64Array, +}; +use arrow_schema::{ArrowError, DataType, Field, Schema, TimeUnit}; +use futures::{stream, StreamExt, TryStreamExt}; +use lance::dataset::mem_wal::{ + DatasetMemWalExt, LsmScanner, ShardManifestStore, ShardSnapshot, ShardWriter, ShardWriterConfig, +}; +use lance::dataset::{builder::DatasetBuilder, Dataset, WriteMode, WriteParams}; +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 tokio::task::JoinHandle; +use tracing::{info, warn}; +use uuid::Uuid; + +use crate::datagen::{ + datagen_trajectory, fold_datagen_events, DatagenBlobValue, DatagenEvent, DatagenEventType, + DatagenTrajectoryPoint, DatagenValue, FoldedDatagenItem, +}; +use crate::rollout_store::derive_shard_id; +use crate::store::{column_as, column_as_optional, timestamp_from_micros}; + +const DEFAULT_MANIFEST_SCAN_BATCH_SIZE: usize = 16; +const DEFAULT_SHARD_SCAN_CONCURRENCY: usize = 16; + +/// Configuration for opening a [`DatagenStore`]. +#[derive(Debug, Clone, Default)] +pub struct DatagenStoreOptions { + pub storage_options: Option>, + /// Stable identity of this writer instance. Multi-writer deployments must + /// assign a distinct value to each live instance. + pub shard_id: Option, + /// Merge this writer's flushed generations into the base table after the + /// threshold is reached. `None` or zero disables count-triggered merging. + pub merge_after_generations: Option, + /// Periodically merge this writer's pending generations. `None` or zero + /// disables the timer. + pub cleanup_interval_secs: Option, +} + +/// A single append-only Lance dataset for datagen checkpoint events. +pub struct DatagenStore { + dataset: Dataset, + write_shard: Uuid, + storage_options: Option>, + merge_after_generations: usize, + cleanup_interval_secs: u64, + write_writer: Option, +} + +impl DatagenStore { + pub async fn open(uri: &str) -> LanceResult { + Self::open_with_options(uri, DatagenStoreOptions::default()).await + } + + pub async fn open_with_options(uri: &str, options: DatagenStoreOptions) -> LanceResult { + Self::open_inner(uri, options, true).await + } + + pub async fn open_existing_with_options( + uri: &str, + options: DatagenStoreOptions, + ) -> LanceResult { + Self::open_inner(uri, options, false).await + } + + async fn open_inner( + uri: &str, + options: DatagenStoreOptions, + create_if_missing: bool, + ) -> LanceResult { + let storage_options = options.storage_options.clone(); + let dataset = match Self::load_with_options(uri, storage_options.clone()).await { + Ok(dataset) => dataset, + Err(LanceError::DatasetNotFound { .. }) if create_if_missing => { + Self::create_with_options(uri, storage_options.clone()).await? + } + Err(error) => return Err(error), + }; + + Ok(Self { + dataset, + write_shard: derive_shard_id(options.shard_id.as_deref()), + storage_options, + merge_after_generations: options.merge_after_generations.unwrap_or(0), + cleanup_interval_secs: options.cleanup_interval_secs.unwrap_or(0), + write_writer: None, + }) + } + + #[must_use] + pub fn uri(&self) -> &str { + self.dataset.uri() + } + + #[must_use] + pub fn version(&self) -> u64 { + self.dataset.manifest.version + } + + /// Append one or more complete checkpoint batches. + /// + /// The supplied slice is persisted as one MemWAL generation. Callers should + /// include FIELD_* events and the corresponding STEP_COMPLETED marker in + /// the same call so a crash cannot expose a partially checkpointed step. + pub async fn append(&mut self, events: &[DatagenEvent]) -> LanceResult { + if events.is_empty() { + return Ok(self.dataset.manifest.version); + } + validate_write_batch(events)?; + let batch = events_to_batch(events)?; + + self.ensure_mem_wal().await?; + self.write_with_resident_writer(&batch).await?; + if self.merge_after_generations > 0 { + self.merge_own_shard_if_ready(self.merge_after_generations) + .await?; + } + Ok(self.dataset.manifest.version) + } + + /// Append one completed step boundary atomically. + /// + /// The batch may contain any number of FIELD_SET/FIELD_APPEND events, + /// including zero, but must contain exactly one STEP_COMPLETED marker and + /// must belong to one item/checkpoint/writer attempt. + pub async fn append_checkpoint(&mut self, events: &[DatagenEvent]) -> LanceResult { + validate_checkpoint_batch(events)?; + self.append(events).await + } + + async fn write_with_resident_writer(&mut self, batch: &RecordBatch) -> LanceResult<()> { + match self.put_seal_drain(batch).await { + Ok(()) => Ok(()), + Err(error) if is_fenced_error(&error) => { + self.write_writer = None; + self.put_seal_drain(batch).await + } + Err(error) => Err(error), + } + } + + async fn put_seal_drain(&mut self, batch: &RecordBatch) -> LanceResult<()> { + self.ensure_write_writer().await?; + let writer = self + .write_writer + .as_ref() + .expect("ensure_write_writer set the writer"); + writer.put(vec![batch.clone()]).await?; + writer.force_seal_active().await?; + writer.wait_for_flush_drain().await?; + Ok(()) + } + + async fn ensure_write_writer(&mut self) -> LanceResult<()> { + if self.write_writer.is_some() { + return Ok(()); + } + let config = ShardWriterConfig { + shard_id: self.write_shard, + ..Default::default() + }; + self.write_writer = Some( + self.dataset + .mem_wal_writer(self.write_shard, config) + .await?, + ); + Ok(()) + } + + /// Gracefully stop this store's resident MemWAL writer. + pub async fn close(&mut self) -> LanceResult<()> { + if let Some(writer) = self.write_writer.take() { + writer.close().await?; + } + Ok(()) + } + + /// Read one item's event history without materializing blob bytes. + pub async fn events_for_item(&self, item_id: &str) -> LanceResult> { + self.filtered_events(&format!("item_id = '{}'", escape_sql_literal(item_id))) + .await + } + + /// Read a root item and every projected descendant without blob bytes. + pub async fn events_for_root(&self, root_item_id: &str) -> LanceResult> { + self.filtered_events(&format!( + "root_item_id = '{}'", + escape_sql_literal(root_item_id) + )) + .await + } + + /// Read failure events directly from the source-of-truth log. + pub async fn failures(&self, run_id: Option<&str>) -> LanceResult> { + let filter = match run_id { + Some(run_id) => format!( + "event_type = 'FAILED' AND run_id = '{}'", + escape_sql_literal(run_id) + ), + None => "event_type = 'FAILED'".to_string(), + }; + self.filtered_events(&filter).await + } + + /// Reconstruct one item's latest state exclusively from its event log. + pub async fn fold_item(&self, item_id: &str) -> LanceResult> { + let events = self.events_for_item(item_id).await?; + if events.is_empty() { + return Ok(None); + } + fold_datagen_events(&events) + .map(Some) + .map_err(invalid_input) + } + + /// Reconstruct state after every completed step without loading blob bytes. + pub async fn trajectory(&self, item_id: &str) -> LanceResult> { + let events = self.events_for_item(item_id).await?; + if events.is_empty() { + return Ok(Vec::new()); + } + datagen_trajectory(&events).map_err(invalid_input) + } + + /// Materialize one FIELD_* event's blob by id. + /// + /// Each candidate dataset is first filtered using only `event_id`; the + /// matching `_rowid` is then passed to `take_rows` for an O(single blob) + /// payload read. + pub async fn get_blob(&self, event_id: &str) -> LanceResult>> { + let snapshots = self.wal_shard_snapshots().await?; + let mut generations: Vec<(u64, String)> = snapshots + .iter() + .flat_map(|snapshot| { + snapshot.flushed_generations.iter().map(|generation| { + ( + generation.generation, + self.flushed_generation_uri(snapshot.shard_id, &generation.path), + ) + }) + }) + .collect(); + generations.sort_by_key(|(generation, _)| std::cmp::Reverse(*generation)); + + for (_, uri) in generations { + let dataset = self.open_flushed_dataset(&uri).await?; + if let Some(payload) = Self::get_blob_from_dataset(&dataset, event_id).await? { + return Ok(payload); + } + } + + Ok(Self::get_blob_from_dataset(&self.dataset, event_id) + .await? + .flatten()) + } + + /// Number of flushed generations waiting across all writer shards. + pub async fn pending_wal_generations(&self) -> LanceResult { + Ok(self + .wal_shard_snapshots() + .await? + .iter() + .map(|snapshot| snapshot.flushed_generations.len()) + .sum()) + } + + /// Merge every currently flushed generation owned by this writer into the + /// base table. + pub async fn cleanup_own_shard(&mut self) -> LanceResult { + self.merge_own_shard_if_ready(1).await + } + + /// Start periodic cleanup after the configured interval. The task keeps + /// only a weak reference and exits when the store is dropped. + pub fn spawn_periodic_cleanup(store: Arc>) -> Option> { + let interval_secs = store.try_read().ok()?.cleanup_interval_secs; + if interval_secs == 0 { + return None; + } + + let weak = Arc::downgrade(&store); + Some(tokio::spawn(async move { + let interval = std::time::Duration::from_secs(interval_secs); + let pass_timeout = interval + .saturating_mul(5) + .max(std::time::Duration::from_secs(30)); + let mut ticker = tokio::time::interval(interval); + ticker.tick().await; + loop { + ticker.tick().await; + let Some(store) = weak.upgrade() else { + return; + }; + let mut guard = store.write().await; + match tokio::time::timeout(pass_timeout, guard.cleanup_own_shard()).await { + Ok(Ok(0)) => {} + Ok(Ok(reclaimed)) => info!( + shard = %guard.write_shard, + reclaimed, + "datagen WAL cleanup merged flushed generations" + ), + Ok(Err(error)) => warn!( + shard = %guard.write_shard, + error = %error, + "datagen WAL cleanup failed" + ), + Err(_) => warn!( + shard = %guard.write_shard, + timeout_secs = pass_timeout.as_secs(), + "datagen WAL cleanup timed out" + ), + } + } + })) + } + + async fn filtered_events(&self, filter: &str) -> LanceResult> { + let columns = self.non_blob_columns(); + let refs: Vec<&str> = columns.iter().map(String::as_str).collect(); + let scanner = self.lsm_scanner().await?.project(&refs).filter(filter)?; + let mut stream = scanner.try_into_stream().await?; + let mut events = Vec::new(); + while let Some(batch) = stream.try_next().await? { + events.extend(batch_to_events(&batch)?); + } + events.sort_by(|left, right| { + left.item_id + .cmp(&right.item_id) + .then_with(|| left.item_seq.cmp(&right.item_seq)) + .then_with(|| left.event_id.cmp(&right.event_id)) + }); + Ok(events) + } + + async fn merge_own_shard_if_ready(&mut self, threshold: usize) -> LanceResult { + let object_store = self.dataset.object_store(None).await?; + let branch_location = self.dataset.branch_location(); + let manifest_store = ShardManifestStore::new( + object_store, + &branch_location.path, + self.write_shard, + DEFAULT_MANIFEST_SCAN_BATCH_SIZE, + ); + let Some(manifest) = manifest_store.read_latest().await? else { + return Ok(0); + }; + let pending = manifest.flushed_generations.len(); + if pending == 0 || pending < threshold.max(1) { + return Ok(0); + } + self.merge_own_shard(&manifest_store, &manifest).await?; + Ok(pending) + } + + async fn merge_own_shard( + &mut self, + manifest_store: &ShardManifestStore, + manifest: &ShardManifest, + ) -> LanceResult<()> { + if manifest.flushed_generations.is_empty() { + return Ok(()); + } + + // claim_epoch below fences the resident writer. Drain it first and + // reopen lazily on the next append. + self.close().await?; + + let base_uri = self.dataset.uri().trim_end_matches('/').to_string(); + let mut merged_generations = HashSet::new(); + let mut merged_paths = Vec::new(); + let mut batches = Vec::new(); + for flushed in &manifest.flushed_generations { + let generation_uri = format!( + "{}/_mem_wal/{}/{}", + base_uri, self.write_shard, flushed.path + ); + let generation = + Self::load_with_options(&generation_uri, self.storage_options.clone()).await?; + let mut stream = generation.scan().try_into_stream().await?; + while let Some(batch) = stream.try_next().await? { + if batch.num_rows() > 0 { + batches.push(batch); + } + } + merged_generations.insert(flushed.generation); + merged_paths.push(flushed.path.clone()); + } + + if !batches.is_empty() { + let schema = Arc::new(datagen_log_schema()); + let reader = RecordBatchIterator::new( + batches.into_iter().map(Ok::), + schema, + ); + let mut params = WriteParams { + mode: WriteMode::Append, + ..Default::default() + }; + if let Some(options) = &self.storage_options { + params.store_params = Some(ObjectStoreParams { + storage_options_accessor: Some(Arc::new( + StorageOptionsAccessor::with_static_options(options.clone()), + )), + ..Default::default() + }); + } + self.dataset.append(reader, Some(params)).await?; + } + + let (epoch, _) = manifest_store.claim_epoch(manifest.shard_spec_id).await?; + manifest_store + .commit_update(epoch, |current| ShardManifest { + version: current.version + 1, + flushed_generations: current + .flushed_generations + .iter() + .filter(|generation| !merged_generations.contains(&generation.generation)) + .cloned() + .collect(), + ..current.clone() + }) + .await?; + + let object_store = self.dataset.object_store(None).await?; + let branch_path = self.dataset.branch_location().path.clone(); + for path in merged_paths { + let generation_path = branch_path + .clone() + .join("_mem_wal") + .join(self.write_shard.to_string().as_str()) + .join(path.as_str()); + if let Err(error) = object_store.remove_dir_all(generation_path).await { + warn!( + shard = %self.write_shard, + generation_path = %path, + error = %error, + "failed to delete merged datagen WAL generation" + ); + } + } + Ok(()) + } + + async fn get_blob_from_dataset( + dataset: &Dataset, + event_id: &str, + ) -> LanceResult>>> { + let mut scanner = dataset.scan(); + scanner + .project(&["event_id"])? + .filter(&format!("event_id = '{}'", escape_sql_literal(event_id)))? + .with_row_id() + .limit(Some(1), None)?; + + let mut stream = scanner.try_into_stream().await?; + while let Some(batch) = stream.try_next().await? { + let event_id_array = column_as::(&batch, "event_id")?; + let row_id_array = column_as::(&batch, "_rowid")?; + for row in 0..batch.num_rows() { + if event_id_array.value(row) != event_id { + continue; + } + let projection = dataset.schema().project(&["value_blob"])?; + let payload_batch = dataset + .take_rows(&[row_id_array.value(row)], projection) + .await?; + let payload = column_as_optional::(&payload_batch, "value_blob"); + return Ok(Some(match payload { + Some(array) if !array.is_null(0) => Some(array.value(0).to_vec()), + _ => None, + })); + } + } + Ok(None) + } + + fn non_blob_columns(&self) -> Vec { + self.dataset + .schema() + .fields + .iter() + .map(|field| field.name.clone()) + .filter(|name| name != "value_blob") + .collect() + } + + async fn lsm_scanner(&self) -> LanceResult { + Ok(LsmScanner::new( + Arc::new(self.dataset.clone()), + self.wal_shard_snapshots().await?, + vec!["event_id".to_string()], + )) + } + + async fn wal_shard_snapshots(&self) -> LanceResult> { + let object_store = self.dataset.object_store(None).await?; + let branch_path = self.dataset.branch_location().path.clone(); + let shard_ids = self.dataset.list_mem_wal_latest_shard_ids().await?; + + let snapshots: Vec> = stream::iter(shard_ids) + .map(|shard_id| { + let object_store = object_store.clone(); + let branch_path = branch_path.clone(); + async move { + let manifest_store = ShardManifestStore::new( + object_store, + &branch_path, + shard_id, + DEFAULT_MANIFEST_SCAN_BATCH_SIZE, + ); + let Some(manifest) = manifest_store.read_latest().await? else { + return Ok(None); + }; + let mut snapshot = ShardSnapshot::new(shard_id) + .with_spec_id(manifest.shard_spec_id) + .with_current_generation(manifest.current_generation); + for flushed in manifest.flushed_generations { + snapshot = + snapshot.with_flushed_generation(flushed.generation, flushed.path); + } + Ok::<_, LanceError>(Some(snapshot)) + } + }) + .buffer_unordered(DEFAULT_SHARD_SCAN_CONCURRENCY) + .try_collect() + .await?; + Ok(snapshots.into_iter().flatten().collect()) + } + + async fn ensure_mem_wal(&mut self) -> LanceResult<()> { + if self.mem_wal_index_present().await? { + return Ok(()); + } + match self + .dataset + .initialize_mem_wal() + .unsharded() + .execute() + .await + { + Ok(()) => Ok(()), + Err(error) => { + let uri = self.dataset.uri().to_string(); + self.dataset = Self::load_with_options(&uri, self.storage_options.clone()).await?; + if self.mem_wal_index_present().await? { + Ok(()) + } else { + Err(error) + } + } + } + } + + async fn mem_wal_index_present(&self) -> LanceResult { + let indices = self.dataset.load_indices().await?; + Ok(indices.iter().any(|index| index.name == MEM_WAL_INDEX_NAME)) + } + + fn flushed_generation_uri(&self, shard_id: Uuid, path: &str) -> String { + format!( + "{}/_mem_wal/{shard_id}/{path}", + self.dataset.uri().trim_end_matches('/') + ) + } + + async fn open_flushed_dataset(&self, uri: &str) -> LanceResult { + let mut builder = DatasetBuilder::from_uri(uri).with_session(self.dataset.session()); + if let Some(options) = self.storage_options.clone() { + builder = builder.with_storage_options(options); + } + builder.load().await + } + + async fn load_with_options( + uri: &str, + storage_options: Option>, + ) -> LanceResult { + if let Some(options) = storage_options { + DatasetBuilder::from_uri(uri) + .with_storage_options(options) + .load() + .await + } else { + Dataset::open(uri).await + } + } + + async fn create_with_options( + uri: &str, + storage_options: Option>, + ) -> LanceResult { + let schema = Arc::new(datagen_log_schema()); + let batches = RecordBatchIterator::new( + vec![Ok::(RecordBatch::new_empty( + schema.clone(), + ))] + .into_iter(), + schema, + ); + let mut params = WriteParams { + mode: WriteMode::Create, + ..Default::default() + }; + if let Some(options) = storage_options { + params.store_params = Some(ObjectStoreParams { + storage_options_accessor: Some(Arc::new( + StorageOptionsAccessor::with_static_options(options), + )), + ..Default::default() + }); + } + Dataset::write(batches, uri, Some(params)).await + } +} + +/// Arrow schema for the single append-only datagen checkpoint log. +#[must_use] +pub fn datagen_log_schema() -> Schema { + let mut event_id_metadata = HashMap::new(); + event_id_metadata.insert( + "lance-schema:unenforced-primary-key".to_string(), + "true".to_string(), + ); + + Schema::new(vec![ + Field::new("event_id", DataType::Utf8, false).with_metadata(event_id_metadata), + Field::new("item_id", DataType::Utf8, false), + Field::new("root_item_id", DataType::Utf8, false), + Field::new("parent_item_id", DataType::Utf8, true), + Field::new("item_seq", DataType::Int64, false), + Field::new("checkpoint_id", DataType::Utf8, false), + Field::new("event_type", DataType::Utf8, false), + Field::new("step_name", DataType::Utf8, true), + Field::new("step_index", DataType::Int64, true), + Field::new("step_instance_id", DataType::Utf8, true), + Field::new("iteration", DataType::Int64, true), + Field::new("attempt", DataType::Int32, false), + Field::new("run_id", DataType::Utf8, false), + Field::new("writer_epoch", DataType::Utf8, false), + Field::new("field_name", DataType::Utf8, true), + Field::new("field_type", DataType::Utf8, true), + Field::new("codec_version", DataType::Int32, true), + Field::new("value_kind", DataType::Utf8, true), + Field::new("value_i64", DataType::Int64, true), + Field::new("value_f64", DataType::Float64, true), + Field::new("value_bool", DataType::Boolean, true), + Field::new("value_str", DataType::LargeUtf8, true), + Field::new("value_json", DataType::LargeUtf8, true), + // Inline LargeBinary is required while MemWAL's LSM scanner does not + // materialize blob-v2 columns. + Field::new("value_blob", DataType::LargeBinary, true), + Field::new("payload_size", DataType::Int64, true), + Field::new("payload_checksum", DataType::Utf8, true), + Field::new("query_tags_json", DataType::LargeUtf8, true), + Field::new("terminal", DataType::Utf8, true), + Field::new("error_type", DataType::Utf8, true), + Field::new("error_dump", DataType::LargeUtf8, true), + Field::new("traceback", DataType::LargeUtf8, true), + Field::new( + "event_ts", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + false, + ), + Field::new("schema_version", DataType::Int32, false), + ]) +} + +fn validate_write_batch(events: &[DatagenEvent]) -> LanceResult<()> { + let mut by_id: HashMap<&str, &DatagenEvent> = HashMap::new(); + let mut item_sequences: HashMap<(&str, i64), &str> = HashMap::new(); + for event in events { + event.validate().map_err(invalid_input)?; + if let Some(previous) = by_id.insert(&event.event_id, event) { + if previous != event { + return Err(invalid_input(format!( + "event_id '{}' was reused with different content", + event.event_id + ))); + } + } + if let Some(previous_event_id) = + item_sequences.insert((&event.item_id, event.item_seq), &event.event_id) + { + if previous_event_id != event.event_id { + return Err(invalid_input(format!( + "item '{}' has duplicate item_seq {} in one batch", + event.item_id, event.item_seq + ))); + } + } + if let Some(DatagenValue::Blob(blob)) = &event.value { + let bytes = blob.bytes.as_ref().ok_or_else(|| { + invalid_input(format!( + "blob event '{}' has no bytes on the write path", + event.event_id + )) + })?; + if blob.size != bytes.len() as i64 { + return Err(invalid_input(format!( + "blob event '{}' declares {} bytes but contains {}", + event.event_id, + blob.size, + bytes.len() + ))); + } + } + } + Ok(()) +} + +fn validate_checkpoint_batch(events: &[DatagenEvent]) -> LanceResult<()> { + let first = events + .first() + .ok_or_else(|| invalid_input("checkpoint batch must not be empty"))?; + let mut completed = 0; + for event in events { + if event.item_id != first.item_id + || event.checkpoint_id != first.checkpoint_id + || event.run_id != first.run_id + || event.writer_epoch != first.writer_epoch + || event.attempt != first.attempt + { + return Err(invalid_input( + "checkpoint batch must share item_id, checkpoint_id, run_id, writer_epoch, and attempt", + )); + } + match event.event_type { + DatagenEventType::FieldSet | DatagenEventType::FieldAppend => {} + DatagenEventType::StepCompleted => completed += 1, + _ => { + return Err(invalid_input( + "checkpoint batch may contain only FIELD_SET, FIELD_APPEND, and STEP_COMPLETED events", + )); + } + } + } + if completed != 1 { + return Err(invalid_input(format!( + "checkpoint batch requires exactly one STEP_COMPLETED event, found {completed}" + ))); + } + let completion = events + .iter() + .find(|event| event.event_type == DatagenEventType::StepCompleted) + .unwrap(); + for event in events.iter().filter(|event| { + matches!( + event.event_type, + DatagenEventType::FieldSet | DatagenEventType::FieldAppend + ) + }) { + if event.step_name != completion.step_name + || event.step_index != completion.step_index + || event.step_instance_id != completion.step_instance_id + || event.iteration != completion.iteration + { + return Err(invalid_input( + "all field events must share the STEP_COMPLETED step identity", + )); + } + } + validate_write_batch(events) +} + +fn events_to_batch(events: &[DatagenEvent]) -> LanceResult { + let mut event_id = StringBuilder::new(); + let mut item_id = StringBuilder::new(); + let mut root_item_id = StringBuilder::new(); + let mut parent_item_id = StringBuilder::new(); + let mut item_seq = Int64Builder::new(); + let mut checkpoint_id = StringBuilder::new(); + let mut event_type = StringBuilder::new(); + let mut step_name = StringBuilder::new(); + let mut step_index = Int64Builder::new(); + let mut step_instance_id = StringBuilder::new(); + let mut iteration = Int64Builder::new(); + let mut attempt = Int32Builder::new(); + let mut run_id = StringBuilder::new(); + let mut writer_epoch = StringBuilder::new(); + let mut field_name = StringBuilder::new(); + let mut field_type = StringBuilder::new(); + let mut codec_version = Int32Builder::new(); + let mut value_kind = StringBuilder::new(); + let mut value_i64 = Int64Builder::new(); + let mut value_f64 = Float64Builder::new(); + let mut value_bool = BooleanBuilder::new(); + let mut value_str = LargeStringBuilder::new(); + let mut value_json = LargeStringBuilder::new(); + let mut value_blob = LargeBinaryBuilder::new(); + let mut payload_size = Int64Builder::new(); + let mut payload_checksum = StringBuilder::new(); + let mut query_tags_json = LargeStringBuilder::new(); + let mut terminal = StringBuilder::new(); + let mut error_type = StringBuilder::new(); + let mut error_dump = LargeStringBuilder::new(); + let mut traceback = LargeStringBuilder::new(); + let mut event_ts = + TimestampMicrosecondBuilder::with_capacity(events.len()).with_timezone("UTC"); + let mut schema_version = Int32Builder::new(); + + for event in events { + event_id.append_value(&event.event_id); + item_id.append_value(&event.item_id); + root_item_id.append_value(&event.root_item_id); + parent_item_id.append_option(event.parent_item_id.as_deref()); + item_seq.append_value(event.item_seq); + checkpoint_id.append_value(&event.checkpoint_id); + event_type.append_value(event.event_type.as_str()); + step_name.append_option(event.step_name.as_deref()); + step_index.append_option(event.step_index); + step_instance_id.append_option(event.step_instance_id.as_deref()); + iteration.append_option(event.iteration); + attempt.append_value(event.attempt); + run_id.append_value(&event.run_id); + writer_epoch.append_value(&event.writer_epoch); + field_name.append_option(event.field_name.as_deref()); + field_type.append_option(event.field_type.as_deref()); + codec_version.append_option(event.codec_version); + + value_kind.append_option(event.value.as_ref().map(DatagenValue::kind)); + value_i64.append_option(match &event.value { + Some(DatagenValue::Int(value)) => Some(*value), + _ => None, + }); + value_f64.append_option(match &event.value { + Some(DatagenValue::Float(value)) => Some(*value), + _ => None, + }); + value_bool.append_option(match &event.value { + Some(DatagenValue::Bool(value)) => Some(*value), + _ => None, + }); + value_str.append_option(match &event.value { + Some(DatagenValue::String(value)) => Some(value.as_str()), + _ => None, + }); + match &event.value { + Some(DatagenValue::Json(value)) => value_json.append_value(value.to_string()), + _ => value_json.append_null(), + } + match &event.value { + Some(DatagenValue::Blob(blob)) => { + value_blob.append_option(blob.bytes.as_deref()); + payload_size.append_value(blob.size); + payload_checksum.append_option(blob.checksum.as_deref()); + } + _ => { + value_blob.append_null(); + payload_size.append_null(); + payload_checksum.append_null(); + } + } + match &event.query_tags { + Some(tags) => query_tags_json.append_value(tags.to_string()), + None => query_tags_json.append_null(), + } + terminal.append_option(event.terminal.map(|value| value.as_str())); + error_type.append_option(event.error_type.as_deref()); + error_dump.append_option(event.error_dump.as_deref()); + traceback.append_option(event.traceback.as_deref()); + event_ts.append_value(event.event_ts.timestamp_micros()); + schema_version.append_value(event.schema_version); + } + + let schema = Arc::new(datagen_log_schema()); + let arrays: Vec = vec![ + Arc::new(event_id.finish()), + Arc::new(item_id.finish()), + Arc::new(root_item_id.finish()), + Arc::new(parent_item_id.finish()), + Arc::new(item_seq.finish()), + Arc::new(checkpoint_id.finish()), + Arc::new(event_type.finish()), + Arc::new(step_name.finish()), + Arc::new(step_index.finish()), + Arc::new(step_instance_id.finish()), + Arc::new(iteration.finish()), + Arc::new(attempt.finish()), + Arc::new(run_id.finish()), + Arc::new(writer_epoch.finish()), + Arc::new(field_name.finish()), + Arc::new(field_type.finish()), + Arc::new(codec_version.finish()), + Arc::new(value_kind.finish()), + Arc::new(value_i64.finish()), + Arc::new(value_f64.finish()), + Arc::new(value_bool.finish()), + Arc::new(value_str.finish()), + Arc::new(value_json.finish()), + Arc::new(value_blob.finish()), + Arc::new(payload_size.finish()), + Arc::new(payload_checksum.finish()), + Arc::new(query_tags_json.finish()), + Arc::new(terminal.finish()), + Arc::new(error_type.finish()), + Arc::new(error_dump.finish()), + Arc::new(traceback.finish()), + Arc::new(event_ts.finish()), + Arc::new(schema_version.finish()), + ]; + Ok(RecordBatch::try_new(schema, arrays)?) +} + +fn batch_to_events(batch: &RecordBatch) -> LanceResult> { + let event_id = column_as::(batch, "event_id")?; + let item_id = column_as::(batch, "item_id")?; + let root_item_id = column_as::(batch, "root_item_id")?; + let parent_item_id = column_as_optional::(batch, "parent_item_id"); + let item_seq = column_as::(batch, "item_seq")?; + let checkpoint_id = column_as::(batch, "checkpoint_id")?; + let event_type = column_as::(batch, "event_type")?; + let step_name = column_as_optional::(batch, "step_name"); + let step_index = column_as_optional::(batch, "step_index"); + let step_instance_id = column_as_optional::(batch, "step_instance_id"); + let iteration = column_as_optional::(batch, "iteration"); + let attempt = column_as::(batch, "attempt")?; + let run_id = column_as::(batch, "run_id")?; + let writer_epoch = column_as::(batch, "writer_epoch")?; + let field_name = column_as_optional::(batch, "field_name"); + let field_type = column_as_optional::(batch, "field_type"); + let codec_version = column_as_optional::(batch, "codec_version"); + let value_kind = column_as_optional::(batch, "value_kind"); + let value_i64 = column_as_optional::(batch, "value_i64"); + let value_f64 = column_as_optional::(batch, "value_f64"); + let value_bool = column_as_optional::(batch, "value_bool"); + let value_str = column_as_optional::(batch, "value_str"); + let value_json = column_as_optional::(batch, "value_json"); + let value_blob = column_as_optional::(batch, "value_blob"); + let payload_size = column_as_optional::(batch, "payload_size"); + let payload_checksum = column_as_optional::(batch, "payload_checksum"); + let query_tags_json = column_as_optional::(batch, "query_tags_json"); + let terminal = column_as_optional::(batch, "terminal"); + let error_type = column_as_optional::(batch, "error_type"); + let error_dump = column_as_optional::(batch, "error_dump"); + let traceback = column_as_optional::(batch, "traceback"); + let event_ts = column_as::(batch, "event_ts")?; + let schema_version = column_as::(batch, "schema_version")?; + + let mut events = Vec::with_capacity(batch.num_rows()); + for row in 0..batch.num_rows() { + let event_id_value = event_id.value(row).to_string(); + let value = match optional_string(value_kind, row).as_deref() { + None => None, + Some("int") => Some(DatagenValue::Int(required_i64( + value_i64, + row, + "value_i64", + )?)), + Some("float") => Some(DatagenValue::Float(required_f64( + value_f64, + row, + "value_f64", + )?)), + Some("bool") => Some(DatagenValue::Bool(required_bool( + value_bool, + row, + "value_bool", + )?)), + Some("str") => Some(DatagenValue::String( + optional_large_string(value_str, row) + .ok_or_else(|| invalid_input("value_kind=str requires value_str"))?, + )), + Some("json") => { + let json = optional_large_string(value_json, row) + .ok_or_else(|| invalid_input("value_kind=json requires value_json"))?; + Some(DatagenValue::Json(serde_json::from_str(&json).map_err( + |error| { + invalid_input(format!( + "event '{}' contains invalid value_json: {}", + event_id_value, error + )) + }, + )?)) + } + Some("blob") => Some(DatagenValue::Blob(DatagenBlobValue { + bytes: optional_bytes(value_blob, row), + size: required_i64(payload_size, row, "payload_size")?, + checksum: optional_string(payload_checksum, row), + })), + Some(other) => { + return Err(invalid_input(format!( + "event '{}' has unsupported value_kind '{}'", + event_id_value, other + ))); + } + }; + + let query_tags = match optional_large_string(query_tags_json, row) { + Some(json) => Some(serde_json::from_str(&json).map_err(|error| { + invalid_input(format!( + "event '{}' contains invalid query_tags_json: {}", + event_id_value, error + )) + })?), + None => None, + }; + let event = DatagenEvent { + event_id: event_id_value, + item_id: item_id.value(row).to_string(), + root_item_id: root_item_id.value(row).to_string(), + parent_item_id: optional_string(parent_item_id, row), + item_seq: item_seq.value(row), + checkpoint_id: checkpoint_id.value(row).to_string(), + event_type: DatagenEventType::parse(event_type.value(row)).map_err(invalid_input)?, + step_name: optional_string(step_name, row), + step_index: optional_i64(step_index, row), + step_instance_id: optional_string(step_instance_id, row), + iteration: optional_i64(iteration, row), + attempt: attempt.value(row), + run_id: run_id.value(row).to_string(), + writer_epoch: writer_epoch.value(row).to_string(), + field_name: optional_string(field_name, row), + field_type: optional_string(field_type, row), + codec_version: optional_i32(codec_version, row), + value, + query_tags, + terminal: match optional_string(terminal, row) { + Some(value) => { + Some(crate::datagen::DatagenTerminal::parse(&value).map_err(invalid_input)?) + } + None => None, + }, + error_type: optional_string(error_type, row), + error_dump: optional_large_string(error_dump, row), + traceback: optional_large_string(traceback, row), + event_ts: timestamp_from_micros(event_ts.value(row), "event_ts")?, + schema_version: schema_version.value(row), + }; + event.validate().map_err(invalid_input)?; + events.push(event); + } + Ok(events) +} + +fn optional_string(array: Option<&StringArray>, row: usize) -> Option { + array + .filter(|array| !array.is_null(row)) + .map(|array| array.value(row).to_string()) +} + +fn optional_large_string(array: Option<&LargeStringArray>, row: usize) -> Option { + array + .filter(|array| !array.is_null(row)) + .map(|array| array.value(row).to_string()) +} + +fn optional_bytes(array: Option<&LargeBinaryArray>, row: usize) -> Option> { + array + .filter(|array| !array.is_null(row)) + .map(|array| array.value(row).to_vec()) +} + +fn optional_i64(array: Option<&Int64Array>, row: usize) -> Option { + array + .filter(|array| !array.is_null(row)) + .map(|array| array.value(row)) +} + +fn optional_i32(array: Option<&Int32Array>, row: usize) -> Option { + array + .filter(|array| !array.is_null(row)) + .map(|array| array.value(row)) +} + +fn required_i64(array: Option<&Int64Array>, row: usize, name: &str) -> LanceResult { + optional_i64(array, row).ok_or_else(|| invalid_input(format!("{name} must not be null"))) +} + +fn required_f64(array: Option<&Float64Array>, row: usize, name: &str) -> LanceResult { + array + .filter(|array| !array.is_null(row)) + .map(|array| array.value(row)) + .ok_or_else(|| invalid_input(format!("{name} must not be null"))) +} + +fn required_bool(array: Option<&BooleanArray>, row: usize, name: &str) -> LanceResult { + array + .filter(|array| !array.is_null(row)) + .map(|array| array.value(row)) + .ok_or_else(|| invalid_input(format!("{name} must not be null"))) +} + +fn escape_sql_literal(value: &str) -> String { + value.replace('\'', "''") +} + +fn invalid_input(message: impl Into) -> LanceError { + LanceError::from(ArrowError::InvalidArgumentError(message.into())) +} + +impl Drop for DatagenStore { + fn drop(&mut self) { + if let Some(writer) = self.write_writer.take() { + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(async move { + let _ = writer.close().await; + }); + } + } + } +} + +fn is_fenced_error(error: &LanceError) -> bool { + let text = error.to_string(); + text.contains("fenced") || text.contains("Fenced") +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::datagen::{ + datagen_event_id, DatagenFieldState, DatagenItemStatus, DatagenTerminal, + DATAGEN_SCHEMA_VERSION, + }; + use chrono::{TimeZone, Utc}; + use serde_json::json; + use tempfile::TempDir; + + fn event( + item_id: &str, + seq: i64, + checkpoint_id: &str, + ordinal: u32, + event_type: DatagenEventType, + ) -> DatagenEvent { + DatagenEvent { + event_id: datagen_event_id(item_id, checkpoint_id, ordinal), + item_id: item_id.to_string(), + root_item_id: item_id.split('/').next().unwrap().to_string(), + parent_item_id: None, + item_seq: seq, + checkpoint_id: checkpoint_id.to_string(), + event_type, + step_name: None, + step_index: None, + step_instance_id: None, + iteration: None, + attempt: 0, + run_id: "run-1".to_string(), + writer_epoch: "writer-1".to_string(), + field_name: None, + field_type: None, + codec_version: None, + value: None, + query_tags: None, + terminal: None, + error_type: None, + error_dump: None, + traceback: None, + event_ts: Utc.timestamp_micros(1_700_000_000_000_000 + seq).unwrap(), + schema_version: DATAGEN_SCHEMA_VERSION, + } + } + + fn field_event( + seq: i64, + ordinal: u32, + event_type: DatagenEventType, + field_name: &str, + field_type: &str, + value: DatagenValue, + ) -> DatagenEvent { + let mut event = event("item-1", seq, "grade-0", ordinal, event_type); + event.step_name = Some("grade".to_string()); + event.step_index = Some(2); + event.step_instance_id = Some("root/grade/0".to_string()); + event.field_name = Some(field_name.to_string()); + event.field_type = Some(field_type.to_string()); + event.codec_version = Some(1); + event.value = Some(value); + event + } + + fn completed_step(seq: i64, ordinal: u32) -> DatagenEvent { + let mut event = event( + "item-1", + seq, + "grade-0", + ordinal, + DatagenEventType::StepCompleted, + ); + event.step_name = Some("grade".to_string()); + event.step_index = Some(2); + event.step_instance_id = Some("root/grade/0".to_string()); + event + } + + #[test] + fn single_log_roundtrip_retry_fold_trajectory_and_blob() { + let directory = TempDir::new().unwrap(); + let uri = directory.path().to_string_lossy().to_string(); + let blob_bytes = b"small-screenshot".to_vec(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = DatagenStore::open(&uri).await.unwrap(); + + let mut created = event("item-1", 0, "created", 0, DatagenEventType::ItemCreated); + created.query_tags = Some(json!({"domain": "math"})); + store.append(&[created]).await.unwrap(); + + let checkpoint = vec![ + field_event( + 1, + 0, + DatagenEventType::FieldSet, + "score", + "int", + DatagenValue::Int(i64::MAX), + ), + field_event( + 2, + 1, + DatagenEventType::FieldAppend, + "messages", + "json", + DatagenValue::Json(json!({"role": "assistant", "content": "42"})), + ), + field_event( + 3, + 2, + DatagenEventType::FieldSet, + "screenshot", + "image", + DatagenValue::Blob(DatagenBlobValue { + bytes: Some(blob_bytes.clone()), + size: blob_bytes.len() as i64, + checksum: Some("sha256:test".to_string()), + }), + ), + completed_step(4, 3), + ]; + store.append_checkpoint(&checkpoint).await.unwrap(); + // Simulate a client retry after an ambiguous response. + store.append_checkpoint(&checkpoint).await.unwrap(); + + let mut terminal = event("item-1", 5, "terminal", 0, DatagenEventType::Terminal); + terminal.terminal = Some(DatagenTerminal::Completed); + store.append(&[terminal]).await.unwrap(); + + let events = store.events_for_item("item-1").await.unwrap(); + assert_eq!(events.len(), 6); + let blob_event = events + .iter() + .find(|event| event.field_name.as_deref() == Some("screenshot")) + .unwrap(); + let DatagenValue::Blob(blob) = blob_event.value.as_ref().unwrap() else { + panic!("screenshot should be a blob"); + }; + assert!(blob.bytes.is_none()); + assert_eq!(blob.size, blob_bytes.len() as i64); + assert_eq!( + store.get_blob(&blob_event.event_id).await.unwrap(), + Some(blob_bytes) + ); + + let folded = store.fold_item("item-1").await.unwrap().unwrap(); + assert_eq!(folded.status, DatagenItemStatus::Completed); + assert_eq!( + folded.fields.get("score"), + Some(&DatagenFieldState::Set(DatagenValue::Int(i64::MAX))) + ); + assert_eq!(folded.completed_steps.len(), 1); + assert_eq!(folded.query_tags, Some(json!({"domain": "math"}))); + + let trajectory = store.trajectory("item-1").await.unwrap(); + assert_eq!(trajectory.len(), 1); + assert_eq!(trajectory[0].cursor.step_name, "grade"); + }); + } + + #[test] + fn no_op_checkpoint_requires_and_persists_completion_marker() { + let directory = TempDir::new().unwrap(); + let uri = directory.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = DatagenStore::open(&uri).await.unwrap(); + let no_op = completed_step(0, 0); + store.append_checkpoint(&[no_op]).await.unwrap(); + let folded = store.fold_item("item-1").await.unwrap().unwrap(); + assert_eq!(folded.completed_steps.len(), 1); + + let field_only = field_event( + 1, + 0, + DatagenEventType::FieldSet, + "score", + "int", + DatagenValue::Int(1), + ); + let error = store + .append_checkpoint(&[field_only]) + .await + .unwrap_err() + .to_string(); + assert!(error.contains("exactly one STEP_COMPLETED")); + }); + } + + #[test] + fn all_wal_shards_are_visible_and_failures_stay_in_the_log() { + let directory = TempDir::new().unwrap(); + let uri = directory.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut writer_a = DatagenStore::open_with_options( + &uri, + DatagenStoreOptions { + storage_options: None, + shard_id: Some("writer-a".to_string()), + ..Default::default() + }, + ) + .await + .unwrap(); + let mut writer_b = DatagenStore::open_with_options( + &uri, + DatagenStoreOptions { + storage_options: None, + shard_id: Some("writer-b".to_string()), + ..Default::default() + }, + ) + .await + .unwrap(); + + writer_a + .append(&[event( + "root-a", + 0, + "created-a", + 0, + DatagenEventType::ItemCreated, + )]) + .await + .unwrap(); + + let mut failed = event("root-b", 0, "failed-b", 0, DatagenEventType::Failed); + failed.run_id = "run-failed".to_string(); + failed.writer_epoch = "writer-b".to_string(); + failed.error_type = Some("ValueError".to_string()); + failed.error_dump = Some("bad source item".to_string()); + writer_b.append(&[failed]).await.unwrap(); + + assert_eq!(writer_a.events_for_item("root-b").await.unwrap().len(), 1); + let failures = writer_a.failures(Some("run-failed")).await.unwrap(); + assert_eq!(failures.len(), 1); + assert_eq!(failures[0].error_type.as_deref(), Some("ValueError")); + }); + } + + #[test] + fn cleanup_moves_events_to_base_without_changing_logical_results() { + let directory = TempDir::new().unwrap(); + let uri = directory.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = DatagenStore::open_with_options( + &uri, + DatagenStoreOptions { + storage_options: None, + shard_id: Some("writer-a".to_string()), + ..Default::default() + }, + ) + .await + .unwrap(); + let created = event("item-1", 0, "created", 0, DatagenEventType::ItemCreated); + let blob = field_event( + 1, + 0, + DatagenEventType::FieldSet, + "image", + "image", + DatagenValue::Blob(DatagenBlobValue { + bytes: Some(b"payload".to_vec()), + size: 7, + checksum: None, + }), + ); + store.append(&[created]).await.unwrap(); + store + .append_checkpoint(&[blob.clone(), completed_step(2, 1)]) + .await + .unwrap(); + assert_eq!(store.pending_wal_generations().await.unwrap(), 2); + + assert_eq!(store.cleanup_own_shard().await.unwrap(), 2); + assert_eq!(store.pending_wal_generations().await.unwrap(), 0); + assert_eq!(store.events_for_item("item-1").await.unwrap().len(), 3); + assert_eq!( + store.get_blob(&blob.event_id).await.unwrap(), + Some(b"payload".to_vec()) + ); + }); + } +} diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index cf78999..d0d12bc 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -3,6 +3,8 @@ mod api_impl; mod context; +mod datagen; +mod datagen_store; mod eval; mod export; mod id; @@ -17,6 +19,13 @@ mod store; pub use api_impl::rollout_record_to_dto; pub use context::{Context, ContextEntry, Snapshot}; +pub use datagen::{ + datagen_event_id, datagen_trajectory, fold_datagen_events, DatagenBlobValue, DatagenEvent, + DatagenEventType, DatagenFailure, DatagenFieldState, DatagenItemStatus, DatagenStepCursor, + DatagenTerminal, DatagenTrajectoryPoint, DatagenValue, FoldedDatagenItem, + DATAGEN_SCHEMA_VERSION, +}; +pub use datagen_store::{datagen_log_schema, DatagenStore, DatagenStoreOptions}; pub use eval::{ AbReport, EvalConfig, EvalQuery, EvalQuerySet, EvalReport, MetricScores, QueryEval, RelevanceLabel, RetrievalMode, diff --git a/crates/lance-context/src/lib.rs b/crates/lance-context/src/lib.rs index d63f33c..2e4a4f4 100644 --- a/crates/lance-context/src/lib.rs +++ b/crates/lance-context/src/lib.rs @@ -3,10 +3,14 @@ // Explicit re-exports from core (no glob to avoid recursion depth overflow) pub use lance_context_core::serde; pub use lance_context_core::{ + datagen_event_id, datagen_log_schema, datagen_trajectory, fold_datagen_events, CompactionConfig, CompactionMetrics, CompactionStats, Context, ContextEntry, ContextNamespace, - ContextRecord, ContextStoreOptions, IdIndexType, LifecycleQueryOptions, MetadataFilter, - PartitionInfo, PartitionSelector, PartitionSpec, RecordFilters, Relationship, RetrieveResult, - RolloutFilters, RolloutRecord, SearchResult, Snapshot, StateMetadata, LIFECYCLE_ACTIVE, + ContextRecord, ContextStoreOptions, DatagenBlobValue, DatagenEvent, DatagenEventType, + DatagenFailure, DatagenFieldState, DatagenItemStatus, DatagenStepCursor, DatagenStore, + DatagenStoreOptions, DatagenTerminal, DatagenTrajectoryPoint, DatagenValue, FoldedDatagenItem, + IdIndexType, LifecycleQueryOptions, MetadataFilter, PartitionInfo, PartitionSelector, + PartitionSpec, RecordFilters, Relationship, RetrieveResult, RolloutFilters, RolloutRecord, + SearchResult, Snapshot, StateMetadata, DATAGEN_SCHEMA_VERSION, LIFECYCLE_ACTIVE, LIFECYCLE_CONTRADICTED, }; diff --git a/specs/datagen-checkpoint-schema.md b/specs/datagen-checkpoint-schema.md new file mode 100644 index 0000000..d999ee7 --- /dev/null +++ b/specs/datagen-checkpoint-schema.md @@ -0,0 +1,101 @@ +# Datagen Checkpoint Log + +## Decision + +Each datagen experiment uses one Lance dataset: + +```text +/datagen_checkpoints///log.lance +``` + +The log is the only durable source of truth. Current state, failures, resume +cursors, and trajectories are all derived by folding events. This avoids +cross-dataset consistency hazards because Lance does not provide atomic +transactions across multiple datasets. + +Optional query projections may be added later, but they must be disposable, +rebuildable, and excluded from checkpoint correctness. + +## Event model + +Every row is an immutable event: + +- `ITEM_CREATED` +- `FIELD_SET` +- `FIELD_APPEND` +- `STEP_COMPLETED` +- `FAILED` +- `TERMINAL` + +Every completed step emits a `STEP_COMPLETED` event, even when no field changed. +All field events and the completion marker for one step are written in the same +checkpoint batch. + +`event_id` is a deterministic idempotency key. Retrying an ambiguously +acknowledged batch writes the same event ids, and the MemWAL LSM read path +de-duplicates them. `item_seq` is strictly increasing per item; two different +events at the same sequence are treated as a writer-fencing violation. + +## Schema + +| Column | Type | Purpose | +|---|---|---| +| `event_id` | string | Deterministic event identity and LSM primary key | +| `item_id` | string | Scoped item identity | +| `root_item_id` | string | Root of the projected item tree | +| `parent_item_id` | string? | Direct parent item | +| `item_seq` | int64 | Per-item event ordering | +| `checkpoint_id` | string | Atomic step-boundary identity | +| `event_type` | string | Event kind | +| `step_name` | string? | Step provenance | +| `step_index` | int64? | Static step position | +| `step_instance_id` | string? | Runtime step identity | +| `iteration` | int64? | Loop/branch iteration | +| `attempt` | int32 | Execution attempt | +| `run_id` | string | Run attribution | +| `writer_epoch` | string | Item ownership/fencing identity | +| `field_name` | string? | Changed field | +| `field_type` | string? | Stable codec id | +| `codec_version` | int32? | Codec compatibility version | +| `value_kind` | string? | `int`, `float`, `bool`, `str`, `json`, or `blob` | +| `value_i64` | int64? | Exact integer value | +| `value_f64` | float64? | Floating-point value | +| `value_bool` | bool? | Boolean value | +| `value_str` | large_string? | String value | +| `value_json` | large_string? | Canonical JSON value | +| `value_blob` | large_binary? | Sparse inline bytes | +| `payload_size` | int64? | Blob size | +| `payload_checksum` | string? | Blob integrity | +| `query_tags_json` | large_string? | Non-authoritative query tags | +| `terminal` | string? | `completed` or `filtered` | +| `error_type` | string? | Failure type | +| `error_dump` | large_string? | Serialized failure | +| `traceback` | large_string? | Failure traceback | +| `event_ts` | timestamp(us, UTC) | Event time | +| `schema_version` | int32 | Log schema compatibility version | + +`value_blob` remains inline while MemWAL's LSM scanner cannot materialize +blob-v2 columns. Normal fold and trajectory reads project it out. Blob access +first locates `event_id` using lightweight columns and then calls `take_rows` +for the exact `_rowid`. + +## Write and recovery invariants + +1. A checkpoint batch is one durable MemWAL generation. +2. Log writes are append-only; retries reuse deterministic event ids. +3. One live owner writes a given item at a time. Ownership changes require a + new `writer_epoch`. +4. Resume reads `item_id = X`, orders by `item_seq`, and folds the events. +5. `FIELD_SET` replaces a field; `FIELD_APPEND` accumulates it. +6. `STEP_COMPLETED` reconstructs the resume cursor. +7. `FAILED` and `TERMINAL` are ordinary log events, not separate datasets. +8. Blob bytes are loaded only when the corresponding lazy reference is used. + +## Maintenance + +Writers use distinct MemWAL shards. Reads union the base table and all flushed +shards, so any instance sees every writer's events. Each writer periodically +merges only its own generations into the base table. Shared base-table +compaction and index refresh must be scheduled by one elected maintenance +worker per experiment. +