From fe9bcac57e7f01774c0a925d197c411c70bdc014 Mon Sep 17 00:00:00 2001 From: Xu Zhang Date: Tue, 14 Jul 2026 14:10:45 -0700 Subject: [PATCH] feat(rollout): add filtered reads --- README.md | 7 ++ crates/lance-context-api/src/lib.rs | 7 ++ crates/lance-context-client/src/lib.rs | 32 +++++ crates/lance-context-core/src/api_impl.rs | 18 ++- crates/lance-context-core/src/lib.rs | 4 +- crates/lance-context-core/src/rollout.rs | 116 ++++++++++++++++++ .../lance-context-core/src/rollout_store.rs | 115 +++++++++++++++-- .../src/routes/rollouts.rs | 96 ++++++++++++++- crates/lance-context/src/lib.rs | 3 +- crates/lance-context/src/unified_rollout.rs | 9 ++ python/python/lance_context/api.py | 19 ++- python/src/lib.rs | 14 ++- python/tests/test_rollout.py | 52 ++++++++ python/tests/test_rollout_remote.py | 30 +++++ specs/rollout-deployment.md | 7 ++ 15 files changed, 510 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index d0d5dc0..de5a75c 100644 --- a/README.md +++ b/README.md @@ -93,6 +93,13 @@ store.add({ for row in store.list(): print(row["rollout_id"], row["reward"]) + +# Exact-match filters are combined with AND and applied before pagination. +training_rows = store.list(filters={ + "policy_version": "ckpt-100", + "include_in_training": True, + "role": "assistant", +}) ``` In a real training run, generation workers and the learner talk to a shared diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index 03530c7..82f7e01 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -130,6 +130,13 @@ pub trait RolloutStoreApi { offset: Option, ) -> impl Future>> + Send; + fn list_filtered( + &self, + limit: Option, + offset: Option, + filters: Option, + ) -> impl Future>> + Send; + fn get(&self, id: &str) -> impl Future>> + Send; diff --git a/crates/lance-context-client/src/lib.rs b/crates/lance-context-client/src/lib.rs index 1ece4c3..71c4e39 100644 --- a/crates/lance-context-client/src/lib.rs +++ b/crates/lance-context-client/src/lib.rs @@ -310,6 +310,25 @@ impl RolloutStoreApi for RemoteRolloutStore { Ok(resp.records) } + async fn list_filtered( + &self, + limit: Option, + offset: Option, + filters: Option, + ) -> ContextResult> { + let filters = filters + .as_ref() + .map(serde_json::to_string) + .transpose() + .map_err(|err| ContextError::InvalidRequest(err.to_string()))?; + let resp = self + .client + .list_rollouts_filtered(&self.store_name, limit, offset, filters.as_deref()) + .await + .map_err(to_ctx_err)?; + Ok(resp.records) + } + async fn get(&self, id: &str) -> ContextResult> { let resp = self .client @@ -776,6 +795,16 @@ impl ContextClient { name: &str, limit: Option, offset: Option, + ) -> Result { + self.list_rollouts_filtered(name, limit, offset, None).await + } + + pub async fn list_rollouts_filtered( + &self, + name: &str, + limit: Option, + offset: Option, + filters: Option<&str>, ) -> Result { let mut request = self .http @@ -786,6 +815,9 @@ impl ContextClient { if let Some(offset) = offset { request = request.query(&[("offset", offset)]); } + if let Some(filters) = filters { + request = request.query(&[("filters", filters)]); + } let resp = request.send().await?; Self::handle_response(resp).await } diff --git a/crates/lance-context-core/src/api_impl.rs b/crates/lance-context-core/src/api_impl.rs index 96084a6..c04a557 100644 --- a/crates/lance-context-core/src/api_impl.rs +++ b/crates/lance-context-core/src/api_impl.rs @@ -15,7 +15,7 @@ use crate::record::{ ContextRecord, LifecycleQueryOptions, RecordFilters, RecordPatch, Relationship, StateMetadata, LIFECYCLE_ACTIVE, }; -use crate::rollout::RolloutRecord; +use crate::rollout::{RolloutFilters, RolloutRecord}; use crate::rollout_store::RolloutStore; use crate::store::{CompactionConfig, ContextStore}; @@ -400,6 +400,22 @@ impl RolloutStoreApi for RolloutStore { Ok(records.into_iter().map(rollout_record_to_dto).collect()) } + async fn list_filtered( + &self, + limit: Option, + offset: Option, + filters: Option, + ) -> ContextResult> { + let filters = filters + .map(RolloutFilters::from_json_value) + .transpose() + .map_err(ContextError::InvalidRequest)?; + let records = RolloutStore::list_filtered(self, limit, offset, filters.as_ref()) + .await + .map_err(to_ctx_err)?; + Ok(records.into_iter().map(rollout_record_to_dto).collect()) + } + async fn get(&self, id: &str) -> ContextResult> { let record = RolloutStore::get_by_id(self, id) .await diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 9e0d655..47b1e43 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -33,7 +33,9 @@ pub use record::{ LIFECYCLE_CONTRADICTED, }; pub use registry::{RegistryEntry, RolloutRegistry}; -pub use rollout::{RolloutRecord, ROLE_ARTIFACT, ROLE_ASSISTANT, ROLE_GRADE, ROLE_TOOL}; +pub use rollout::{ + RolloutFilters, RolloutRecord, ROLE_ARTIFACT, ROLE_ASSISTANT, ROLE_GRADE, ROLE_TOOL, +}; pub use rollout_store::{rollout_schema, RolloutObservation, RolloutStore, RolloutStoreOptions}; pub use store::{ CompactionConfig, CompactionStats, ContextStore, ContextStoreOptions, DistanceMetric, diff --git a/crates/lance-context-core/src/rollout.rs b/crates/lance-context-core/src/rollout.rs index a6db415..d11b79b 100644 --- a/crates/lance-context-core/src/rollout.rs +++ b/crates/lance-context-core/src/rollout.rs @@ -9,6 +9,71 @@ pub const ROLE_TOOL: &str = "tool"; pub const ROLE_GRADE: &str = "grade"; pub const ROLE_ARTIFACT: &str = "artifact"; +/// Exact-match filters supported by rollout list queries. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct RolloutFilters { + pub rollout_id: Option, + pub problem_id: Option, + pub policy_version: Option, + pub role: Option, + pub include_in_training: Option, + pub artifact_type: Option, +} + +impl RolloutFilters { + pub fn from_json_value(value: Value) -> Result { + let Value::Object(object) = value else { + return Err("rollout filters must be a JSON object".to_string()); + }; + + let mut filters = Self::default(); + for (key, value) in object { + match key.as_str() { + "rollout_id" => filters.rollout_id = Some(filter_string(&key, value)?), + "problem_id" => filters.problem_id = Some(filter_string(&key, value)?), + "policy_version" => { + filters.policy_version = Some(filter_string(&key, value)?); + } + "role" => filters.role = Some(filter_string(&key, value)?), + "include_in_training" => { + filters.include_in_training = Some(value.as_bool().ok_or_else(|| { + "rollout filter 'include_in_training' must be a boolean".to_string() + })?); + } + "artifact_type" => filters.artifact_type = Some(filter_string(&key, value)?), + _ => return Err(format!("unsupported rollout filter '{key}'")), + } + } + Ok(filters) + } + + pub(crate) fn predicate(&self) -> Option { + let mut parts = Vec::new(); + push_string_predicate(&mut parts, "rollout_id", self.rollout_id.as_deref()); + push_string_predicate(&mut parts, "problem_id", self.problem_id.as_deref()); + push_string_predicate(&mut parts, "policy_version", self.policy_version.as_deref()); + push_string_predicate(&mut parts, "role", self.role.as_deref()); + if let Some(value) = self.include_in_training { + parts.push(format!("include_in_training = {value}")); + } + push_string_predicate(&mut parts, "artifact_type", self.artifact_type.as_deref()); + (!parts.is_empty()).then(|| parts.join(" AND ")) + } +} + +fn filter_string(name: &str, value: Value) -> Result { + value + .as_str() + .map(str::to_string) + .ok_or_else(|| format!("rollout filter '{name}' must be a string")) +} + +fn push_string_predicate(parts: &mut Vec, column: &str, value: Option<&str>) { + if let Some(value) = value { + parts.push(format!("{column} = '{}'", value.replace('\'', "''"))); + } +} + /// One row of a reinforcement-learning rollout dataset. /// /// A row is one message in a trajectory — an assistant turn, a tool call, a @@ -104,3 +169,54 @@ impl RolloutRecord { self.role == ROLE_ARTIFACT } } + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn rollout_filters_parse_supported_fields() { + let filters = RolloutFilters::from_json_value(json!({ + "rollout_id": "traj-1", + "problem_id": "problem-7", + "policy_version": "ckpt-42", + "role": "assistant", + "include_in_training": false, + "artifact_type": "screenshot" + })) + .unwrap(); + + assert_eq!(filters.rollout_id.as_deref(), Some("traj-1")); + assert_eq!(filters.problem_id.as_deref(), Some("problem-7")); + assert_eq!(filters.policy_version.as_deref(), Some("ckpt-42")); + assert_eq!(filters.role.as_deref(), Some("assistant")); + assert_eq!(filters.include_in_training, Some(false)); + assert_eq!(filters.artifact_type.as_deref(), Some("screenshot")); + } + + #[test] + fn rollout_filters_reject_unknown_and_wrong_types() { + assert!(RolloutFilters::from_json_value(json!({"reward": 1.0})).is_err()); + assert!(RolloutFilters::from_json_value(json!({"policy_version": 42})).is_err()); + assert!(RolloutFilters::from_json_value(json!({"include_in_training": "yes"})).is_err()); + assert!(RolloutFilters::from_json_value(json!([])).is_err()); + } + + #[test] + fn rollout_filter_predicate_escapes_strings_and_fields() { + let filters = RolloutFilters { + policy_version: Some("worker's-ckpt".to_string()), + role: Some("assistant".to_string()), + include_in_training: Some(true), + ..Default::default() + }; + + assert_eq!( + filters.predicate().as_deref(), + Some( + "policy_version = 'worker''s-ckpt' AND role = 'assistant' AND include_in_training = true" + ) + ); + } +} diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 2abb3a1..80f1e3e 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -74,7 +74,7 @@ use tokio::task::JoinHandle; use tracing::{info, warn}; use uuid::Uuid; -use crate::rollout::RolloutRecord; +use crate::rollout::{RolloutFilters, RolloutRecord}; use crate::store::{ column_as, column_as_optional, relationship_field, relationship_list_item_field, relationship_struct_builder, relationships_from_list, timestamp_from_micros, CompactionConfig, @@ -766,22 +766,41 @@ impl RolloutStore { &self, limit: Option, offset: Option, + ) -> LanceResult> { + self.list_filtered(limit, offset, None).await + } + + /// List rollout rows matching exact field filters. + /// + /// The filter is pushed into every LSM source before merge and + /// deduplication. Rollout rows are immutable by id, so filtering each + /// generation cannot reveal an older state of a mutable record. + pub async fn list_filtered( + &self, + limit: Option, + offset: Option, + filters: Option<&RolloutFilters>, ) -> 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); + let mut scanner = self.lsm_scanner().await?.project(&refs); + if let Some(predicate) = filters.and_then(RolloutFilters::predicate) { + scanner = scanner.filter(&predicate)?; + } + let post_scan_offset = if limit.is_none() { offset } else { None }; + if let Some(limit) = limit { + scanner = scanner.limit(limit, offset); + } + let mut stream = scanner.try_into_stream().await?; let mut results = Vec::new(); while let Some(batch) = stream.try_next().await? { results.extend(batch_to_rollout_records(&batch)?); } - if let Some(offset) = offset { + if let Some(offset) = post_scan_offset { results = results.into_iter().skip(offset).collect(); } - if let Some(limit) = limit { - results.truncate(limit); - } Ok(results) } @@ -1520,7 +1539,7 @@ fn optional_i8_list(array: Option<&ListArray>, row: usize) -> LanceResult