feat(datagen): close the fold/overview gaps and raise typed store errors

crates/lance-context-api/src/lib.rs

@@ -327,6 +327,16 @@ pub trait DatagenStoreApi {
         item_id: &str,
     ) -> impl Future<Output = ContextResult<Option<FoldedDatagenItemDto>>> + Send;
 
+    /// Like [`DatagenStoreApi::fold_item`], but `load_blobs` selects the blob projection:
+    /// `false` (the `fold_item` default) leaves blob fields lazy — bytes absent, resolved later
+    /// through `get_blob` — while `true` materializes them inline, at the cost of reading the
+    /// payload column.
+    fn fold_item_with_blobs(
+        &self,
+        item_id: &str,
+        load_blobs: bool,
+    ) -> impl Future<Output = ContextResult<Option<FoldedDatagenItemDto>>> + Send;
+
     fn root_item_statuses(
         &self,
         root_item_ids: &[String],
@@ -346,6 +356,10 @@ pub trait DatagenStoreApi {
         root_item_id: &str,
     ) -> impl Future<Output = ContextResult<Vec<DatagenEventDto>>> + Send;
 
+    /// Aggregate the whole log into a run overview: per-status root item counts,
+    /// failure counts by error type, and completed-step counts.
+    fn overview(&self) -> impl Future<Output = ContextResult<DatagenRunOverviewDto>> + Send;
+
     /// Materialize one FIELD_* event's offloaded blob bytes by event id.
     /// Returns `None` when the event or its payload is absent.
     fn get_blob(
@@ -1155,6 +1169,19 @@ pub struct DatagenStepCursorDto {
     pub item_seq: i64,
 }
 
+/// A position within one stream's step tree, without the `item_seq` a cursor carries.
+/// Mirrors the Python `StepPosition` wire dict.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DatagenStreamPositionDto {
+    pub step_name: String,
+    pub step_kind: String,
+    pub step_index: i64,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub enclosing_step: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub selector_step: Option<String>,
+}
+
 /// An item reconstructed by folding its events into latest state.
 /// Mirrors the Python `FoldedItem` wire dict.
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -1168,6 +1195,13 @@ pub struct FoldedDatagenItemDto {
     pub last_attempt: i32,
     pub fields: std::collections::BTreeMap<String, DatagenFieldStateDto>,
     pub trajectory: Vec<DatagenStepCursorDto>,
+    /// Positions with a STEP_STARTED — gates driver-frame (re-)opening on resume. `started`
+    /// minus `completed` = frames that were open when the process died.
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub started: Vec<DatagenStreamPositionDto>,
+    /// Positions with a STEP_COMPLETED — gates STEP_COMPLETED re-emission on resume.
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub completed: Vec<DatagenStreamPositionDto>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub query_tags: Option<Value>,
     /// `field_name -> event_id` for the folded blob fields, so a caller can resolve a blob by field
@@ -1201,6 +1235,38 @@ pub struct DatagenFailureDto {
     pub traceback: Option<String>,
 }
 
+/// Whole-run aggregation over a datagen log. `items` counts root items only
+/// (`running + completed + filtered`); `failures` counts FAILED events, which are
+/// non-terminal, so a failed-then-retried item still counts as `running`.
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct DatagenRunOverviewDto {
+    pub items: usize,
+    pub running: usize,
+    pub completed: usize,
+    pub filtered: usize,
+    pub failures: usize,
+    /// FAILED-event count per `error_type`.
+    #[serde(default)]
+    pub failures_by_error_type: std::collections::BTreeMap<String, usize>,
+    /// Failure roll-up grouped by the `run_id` that emitted the FAILED event.
+    #[serde(default)]
+    pub failures_by_run: std::collections::BTreeMap<String, DatagenFailureBucketDto>,
+    /// STEP_COMPLETED count per step name, across every item.
+    #[serde(default)]
+    pub completed_steps: std::collections::BTreeMap<String, usize>,
+}
+
+/// One `run_id`'s slice of an overview's failure roll-up, with a capped sample of failing root
+/// item ids to drill into.
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct DatagenFailureBucketDto {
+    pub failures: usize,
+    #[serde(default)]
+    pub failures_by_error_type: std::collections::BTreeMap<String, usize>,
+    #[serde(default)]
+    pub sample_root_item_ids: Vec<String>,
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct ListDatagenFailuresResponse {
     pub failures: Vec<DatagenFailureDto>,

crates/lance-context-client/src/lib.rs

@@ -548,6 +548,19 @@ impl DatagenStoreApi for RemoteDatagenStore {
         Ok(resp.item)
     }
 
+    async fn fold_item_with_blobs(
+        &self,
+        item_id: &str,
+        load_blobs: bool,
+    ) -> ContextResult<Option<FoldedDatagenItemDto>> {
+        let resp = self
+            .client
+            .fold_datagen_item_with_blobs(&self.store_name, item_id, load_blobs)
+            .await
+            .map_err(to_ctx_err)?;
+        Ok(resp.item)
+    }
+
     async fn root_item_statuses(
         &self,
         root_item_ids: &[String],
@@ -558,6 +571,13 @@ impl DatagenStoreApi for RemoteDatagenStore {
             .map_err(to_ctx_err)
     }
 
+    async fn overview(&self) -> ContextResult<DatagenRunOverviewDto> {
+        self.client
+            .datagen_overview(&self.store_name)
+            .await
+            .map_err(to_ctx_err)
+    }
+
     async fn item_failures(&self, item_id: &str) -> ContextResult<Vec<DatagenFailureDto>> {
         let resp = self
             .client
@@ -1286,10 +1306,23 @@ impl ContextClient {
         &self,
         name: &str,
         item_id: &str,
+    ) -> Result<GetFoldedDatagenItemResponse, ClientError> {
+        self.fold_datagen_item_with_blobs(name, item_id, false)
+            .await
+    }
+
+    /// Fold an item, choosing the blob projection: `load_blobs` materializes blob bytes inline
+    /// instead of leaving them to a later `get_blob`.
+    pub async fn fold_datagen_item_with_blobs(
+        &self,
+        name: &str,
+        item_id: &str,
+        load_blobs: bool,
     ) -> Result<GetFoldedDatagenItemResponse, ClientError> {
         let resp = self
             .http
             .get(self.url(&format!("/datagen/{}/items/{}", name, item_id)))
+            .query(&[("load_blobs", load_blobs)])
             .send()
             .await?;
         Self::handle_response(resp).await
@@ -1308,6 +1341,16 @@ impl ContextClient {
         Self::handle_response(resp).await
     }
 
+    /// Aggregate the whole datagen log into a run overview.
+    pub async fn datagen_overview(&self, name: &str) -> Result<DatagenRunOverviewDto, ClientError> {
+        let resp = self
+            .http
+            .get(self.url(&format!("/datagen/{}/overview", name)))
+            .send()
+            .await?;
+        Self::handle_response(resp).await
+    }
+
     pub async fn datagen_root_item_statuses(
         &self,
         name: &str,

crates/lance-context-core/src/api_impl.rs

@@ -5,8 +5,9 @@ use uuid::Uuid;
 use lance_context_api::{
     AddDatagenEventsResponse, AddRecordRequest, AddRecordsResponse, AddRolloutRequest,
     AddRolloutsResponse, AddRowsResponse, CompactRequest, CompactResponse, CompactStatsResponse,
-    ContextError, ContextResult, ContextStoreApi, DatagenEventDto, DatagenFailureDto,
-    DatagenFieldStateDto, DatagenRootItemStatusesResponse, DatagenStepCursorDto, DatagenStoreApi,
+    ContextError, ContextResult, ContextStoreApi, DatagenEventDto, DatagenFailureBucketDto,
+    DatagenFailureDto, DatagenFieldStateDto, DatagenRootItemStatusesResponse,
+    DatagenRunOverviewDto, DatagenStepCursorDto, DatagenStoreApi, DatagenStreamPositionDto,
     DatagenValueDto, DeleteRecordResponse, FoldedDatagenItemDto, GenericStoreApi, RecordDto,
     RecordPatchDto, RelationshipDto, RetrieveRequest, RetrieveResultDto, RolloutRecordDto,
     RolloutStoreApi, SchemaSpec, SearchRequest, SearchResultDto, StateMetadataDto,
@@ -15,9 +16,10 @@ use lance_context_api::{
 };
 
 use crate::datagen::{
-    DatagenBlobValue, DatagenEvent, DatagenEventType, DatagenFailure, DatagenFieldState,
-    DatagenItemLookup, DatagenItemStatus, DatagenRootItemStatuses, DatagenStepCursor,
-    DatagenStepKind, DatagenValue, FoldedDatagenItem,
+    DatagenBlobProjection, DatagenBlobValue, DatagenEvent, DatagenEventType, DatagenFailure,
+    DatagenFieldState, DatagenItemLookup, DatagenItemStatus, DatagenRootItemStatuses,
+    DatagenRunOverview, DatagenStepCursor, DatagenStepKind, DatagenStreamPosition, DatagenValue,
+    FoldedDatagenItem,
 };
 use crate::datagen_store::DatagenStore;
 use crate::generic_codec::Row;
@@ -771,6 +773,25 @@ impl DatagenStoreApi for DatagenStore {
         })
     }
 
+    async fn fold_item_with_blobs(
+        &self,
+        item_id: &str,
+        load_blobs: bool,
+    ) -> ContextResult<Option<FoldedDatagenItemDto>> {
+        let blobs = if load_blobs {
+            DatagenBlobProjection::Eager
+        } else {
+            DatagenBlobProjection::Lazy
+        };
+        let lookup = DatagenStore::fold_item_with(self, item_id, blobs)
+            .await
+            .map_err(to_ctx_err)?;
+        Ok(match lookup {
+            DatagenItemLookup::NeverStarted => None,
+            DatagenItemLookup::Found(item) => Some(folded_item_to_dto(&item)),
+        })
+    }
+
     async fn root_item_statuses(
         &self,
         root_item_ids: &[String],
@@ -782,6 +803,11 @@ impl DatagenStoreApi for DatagenStore {
         Ok(root_item_statuses_to_dto(&statuses))
     }
 
+    async fn overview(&self) -> ContextResult<DatagenRunOverviewDto> {
+        let overview = DatagenStore::overview(self).await.map_err(to_ctx_err)?;
+        Ok(run_overview_to_dto(&overview))
+    }
+
     async fn item_failures(&self, item_id: &str) -> ContextResult<Vec<DatagenFailureDto>> {
         let failures = DatagenStore::item_failures(self, item_id)
             .await
@@ -972,6 +998,8 @@ pub fn folded_item_to_dto(item: &FoldedDatagenItem) -> FoldedDatagenItemDto {
             .map(|(name, state)| (name.clone(), field_state_to_dto(state)))
             .collect(),
         trajectory: item.trajectory.ordered.iter().map(cursor_to_dto).collect(),
+        started: position_set_to_dto(&item.trajectory.started),
+        completed: position_set_to_dto(&item.trajectory.completed),
         query_tags: item.query_tags.clone(),
         blob_event_ids: item.blob_event_ids.clone(),
     }
@@ -1003,6 +1031,38 @@ fn cursor_to_dto(cursor: &DatagenStepCursor) -> DatagenStepCursorDto {
     }
 }
 
+fn position_to_dto(position: &DatagenStreamPosition) -> DatagenStreamPositionDto {
+    DatagenStreamPositionDto {
+        step_name: position.step.name.clone(),
+        step_kind: position.step.kind.as_str().to_string(),
+        step_index: position.index,
+        enclosing_step: position.enclosing.clone(),
+        selector_step: position.selector.clone(),
+    }
+}
+
+/// Project a fold position set to DTOs in a deterministic order (the sets are unordered).
+fn position_set_to_dto(
+    positions: &std::collections::HashSet<DatagenStreamPosition>,
+) -> Vec<DatagenStreamPositionDto> {
+    let mut dtos: Vec<DatagenStreamPositionDto> = positions.iter().map(position_to_dto).collect();
+    dtos.sort_by(|a, b| {
+        (
+            &a.step_name,
+            a.step_index,
+            &a.enclosing_step,
+            &a.selector_step,
+        )
+            .cmp(&(
+                &b.step_name,
+                b.step_index,
+                &b.enclosing_step,
+                &b.selector_step,
+            ))
+    });
+    dtos
+}
+
 fn root_item_statuses_to_dto(
     statuses: &DatagenRootItemStatuses,
 ) -> DatagenRootItemStatusesResponse {
@@ -1014,6 +1074,32 @@ fn root_item_statuses_to_dto(
     }
 }
 
+fn run_overview_to_dto(overview: &DatagenRunOverview) -> DatagenRunOverviewDto {
+    DatagenRunOverviewDto {
+        items: overview.items,
+        running: overview.running,
+        completed: overview.completed,
+        filtered: overview.filtered,
+        failures: overview.failures,
+        failures_by_error_type: overview.failures_by_error_type.clone(),
+        failures_by_run: overview
+            .failures_by_run
+            .iter()
+            .map(|(run_id, bucket)| {
+                (
+                    run_id.clone(),
+                    DatagenFailureBucketDto {
+                        failures: bucket.failures,
+                        failures_by_error_type: bucket.failures_by_error_type.clone(),
+                        sample_root_item_ids: bucket.sample_root_item_ids.clone(),
+                    },
+                )
+            })
+            .collect(),
+        completed_steps: overview.completed_steps.clone(),
+    }
+}
+
 fn failure_to_dto(failure: &DatagenFailure) -> DatagenFailureDto {
     DatagenFailureDto {
         at: cursor_to_dto(&failure.at),