diff --git a/Cargo.lock b/Cargo.lock index 9cd03fc..219aea5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5457,6 +5457,7 @@ dependencies = [ name = "lance-context-api" version = "0.6.4" dependencies = [ + "arrow-schema 58.3.0", "base64", "chrono", "serde", diff --git a/crates/lance-context-api/Cargo.toml b/crates/lance-context-api/Cargo.toml index 5c35e15..9856488 100644 --- a/crates/lance-context-api/Cargo.toml +++ b/crates/lance-context-api/Cargo.toml @@ -9,6 +9,7 @@ description = "Shared request/response types for the lance-context REST API" keywords = ["context", "lance", "api"] [dependencies] +arrow-schema = "58" base64 = "0.22" chrono = { version = "0.4", default-features = false, features = ["clock", "serde"] } serde = { version = "1", features = ["derive"] } diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index 1fdb00b..9e30efb 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -1,7 +1,10 @@ +pub mod schema_spec; +pub use schema_spec::{ColumnSpec, ColumnType, SchemaSpec, ID_COLUMN}; + use base64::{engine::general_purpose::STANDARD as BASE64, Engine}; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; -use serde_json::Value; +use serde_json::{Map, Value}; use std::future::Future; // --------------------------------------------------------------------------- @@ -178,6 +181,118 @@ pub trait RolloutStoreApi { fn checkout(&mut self, version: u64) -> impl Future> + Send; } +// --------------------------------------------------------------------------- +// Generic stores (user-defined schemas) +// --------------------------------------------------------------------------- + +/// Create a store whose columns the caller declares. +#[derive(Debug, Serialize, Deserialize)] +pub struct CreateGenericStoreRequest { + /// Portable dataset name: 1-128 ASCII characters matching + /// `[A-Za-z0-9_][A-Za-z0-9._-]*`; `_registry` and `_stats` are reserved. + pub name: String, + /// The store's columns. Must declare an `id` column (non-nullable string); + /// see [`SchemaSpec`]. + pub schema: SchemaSpec, + #[serde(default)] + pub storage_options: Option>, + /// Whether an append seals before returning, making its rows immediately + /// readable. Defaults to `false` — durable but not visible until a flush, + /// which is what keeps concurrent appends from serializing. + #[serde(default)] + pub seal_on_add: bool, +} + +/// A generic store's identity and, for single-store lookups, its schema. +#[derive(Debug, Serialize, Deserialize)] +pub struct GenericStoreInfo { + pub name: String, + pub uri: String, + /// Dataset version. `None` in list responses, which are served without + /// opening each dataset. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub version: Option, + /// The store's schema. `None` in list responses for the same reason as + /// `version`: the schema is stored *in the dataset*, so reporting it means + /// opening every store. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub schema: Option, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct ListGenericStoresResponse { + pub stores: Vec, +} + +/// Rows to append, as JSON objects keyed by column name. +/// +/// There is deliberately no per-column DTO: the schema is declared at runtime, +/// so a static struct could not describe it. Values are validated against the +/// store's schema on the way in — an undeclared key is an error rather than +/// being dropped, and an omitted nullable column is written as null. +#[derive(Debug, Serialize, Deserialize)] +pub struct AddRowsRequest { + pub rows: Vec>, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct AddRowsResponse { + /// Base dataset version after the append. MemWAL appends do not advance it, + /// so this identifies the table, not the rows just written. + pub version: u64, + /// Number of rows accepted. + pub count: usize, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct ListRowsResponse { + pub rows: Vec>, +} + +/// Remote-capable surface of a store over a user-declared schema. +/// +/// Rows are `serde_json` maps in both directions, so this trait needs no DTO +/// conversion layer — the wire form *is* the row form. +pub trait GenericStoreApi { + /// The schema this store was created with. + fn spec(&self) -> &SchemaSpec; + + /// Append rows. Visibility on return depends on the store's `seal_on_add`. + fn add( + &self, + rows: &[Map], + ) -> impl Future> + Send; + + /// Read rows, newest-write-wins by `id`. Blob columns are projected out. + fn list( + &self, + limit: Option, + offset: Option, + ) -> impl Future>>> + Send; + + /// [`Self::list`], filtered by a SQL predicate over the store's columns. + fn list_filtered( + &self, + filter: &str, + limit: Option, + offset: Option, + ) -> impl Future>>> + Send; + + /// Fetch one row by `id`. `columns` selects what to read; `None` reads + /// everything except blob columns. Pass an explicit list to fetch a blob — + /// that is the intended way to read a large payload. + fn get( + &self, + id: &str, + columns: Option<&[String]>, + ) -> impl Future>>> + Send; + + /// Seal the active memtable so previously added rows become readable. + fn flush(&self) -> impl Future> + Send; + + fn version(&self) -> u64; +} + // --------------------------------------------------------------------------- // Datagen trait // --------------------------------------------------------------------------- diff --git a/crates/lance-context-core/src/schema_spec.rs b/crates/lance-context-api/src/schema_spec.rs similarity index 96% rename from crates/lance-context-core/src/schema_spec.rs rename to crates/lance-context-api/src/schema_spec.rs index d5da7ef..802b079 100644 --- a/crates/lance-context-core/src/schema_spec.rs +++ b/crates/lance-context-api/src/schema_spec.rs @@ -24,6 +24,13 @@ //! - **Vector columns carry their dimension and metric**, because an index //! cannot be built without them. //! +//! # This type lives in the API crate on purpose +//! +//! A schema is part of the wire contract — the server accepts one at store +//! creation and the client needs the same type to send it — so it belongs +//! alongside the other request/response types rather than in the storage crate. +//! It depends only on `arrow-schema` and `serde`, not on Lance. +//! //! # Schema is immutable after creation //! //! v1 deliberately has no evolution path: the schema is fixed when the store is @@ -37,8 +44,6 @@ use std::collections::{HashMap, HashSet}; use arrow_schema::{ArrowError, DataType, Field, Schema, TimeUnit}; use serde::{Deserialize, Serialize}; -use crate::store::DistanceMetric; - /// The mandatory primary-key column. Always the LSM merge key. pub const ID_COLUMN: &str = "id"; @@ -61,6 +66,11 @@ const BLOB_COLUMN_KEY: &str = "lance-context:blob"; /// a user column. const RESERVED_COLUMNS: &[&str] = &["_rowid", "_rowaddr", "_mem_wal", "_distance"]; +/// Distance metrics a vector column may declare. Kept in sync with +/// `lance_context_core::DistanceMetric`, which parses the same identifiers; +/// duplicated as strings so this crate stays free of a storage dependency. +const VALID_METRICS: &[&str] = &["l2", "euclidean", "cosine", "dot", "dot_product"]; + /// Maximum number of columns in a user schema. A guard against pathological /// declarations, not a Lance limit. const MAX_COLUMNS: usize = 1024; @@ -341,12 +351,12 @@ fn validate_column_type(name: &str, column_type: &ColumnType) -> Result<(), Stri )); } if let Some(metric) = metric { - DistanceMetric::parse(metric).map_err(|_| { - format!( + if !VALID_METRICS.contains(&metric.to_ascii_lowercase().as_str()) { + return Err(format!( "column '{name}': invalid distance metric '{metric}', expected one of \ - 'l2', 'cosine', 'dot'" - ) - })?; + {VALID_METRICS:?}" + )); + } } Ok(()) } diff --git a/crates/lance-context-client/src/lib.rs b/crates/lance-context-client/src/lib.rs index 2351fc1..61200d5 100644 --- a/crates/lance-context-client/src/lib.rs +++ b/crates/lance-context-client/src/lib.rs @@ -1,5 +1,6 @@ use lance_context_api::*; use reqwest::Client; +use serde_json::{Map, Value}; mod error; pub use error::ClientError; @@ -361,6 +362,121 @@ impl RolloutStoreApi for RemoteRolloutStore { } } +/// A [`GenericStoreApi`] backed by the REST server. +/// +/// Caches the schema at connect: unlike the fixed-schema stores there is no +/// compile-time type describing the columns, and `GenericStoreApi::spec` is +/// synchronous, so it cannot be fetched on demand. +pub struct RemoteGenericStore { + client: ContextClient, + store_name: String, + spec: SchemaSpec, + cached_version: u64, +} + +impl RemoteGenericStore { + pub async fn connect(base_url: &str, store_name: &str) -> Result { + let client = ContextClient::new(base_url); + let info = client.get_generic_store(store_name).await?; + let spec = info.schema.ok_or_else(|| ClientError::Api { + status: 500, + code: "MISSING_SCHEMA".to_string(), + message: "server did not report the store schema".to_string(), + })?; + Ok(Self { + client, + store_name: store_name.to_string(), + spec, + cached_version: info.version.unwrap_or(0), + }) + } + + pub async fn connect_or_create( + base_url: &str, + req: &CreateGenericStoreRequest, + ) -> Result { + let client = ContextClient::new(base_url); + let info = match client.get_generic_store(&req.name).await { + Ok(info) => info, + Err(ClientError::Api { status: 404, .. }) => client.create_generic_store(req).await?, + Err(e) => return Err(e), + }; + Ok(Self { + client, + store_name: req.name.clone(), + spec: info.schema.unwrap_or_else(|| req.schema.clone()), + cached_version: info.version.unwrap_or(0), + }) + } + + /// Seal the store's memtable so added rows become readable. + pub async fn flush_remote(&self) -> Result<(), ClientError> { + self.client.flush_generic_store(&self.store_name).await + } +} + +impl GenericStoreApi for RemoteGenericStore { + fn spec(&self) -> &SchemaSpec { + &self.spec + } + + async fn add(&self, rows: &[Map]) -> ContextResult { + self.client + .add_rows(&self.store_name, rows) + .await + .map_err(to_ctx_err) + } + + async fn list( + &self, + limit: Option, + offset: Option, + ) -> ContextResult>> { + Ok(self + .client + .list_rows(&self.store_name, None, limit, offset) + .await + .map_err(to_ctx_err)? + .rows) + } + + async fn list_filtered( + &self, + filter: &str, + limit: Option, + offset: Option, + ) -> ContextResult>> { + Ok(self + .client + .list_rows(&self.store_name, Some(filter), limit, offset) + .await + .map_err(to_ctx_err)? + .rows) + } + + async fn get( + &self, + id: &str, + columns: Option<&[String]>, + ) -> ContextResult>> { + self.client + .get_row(&self.store_name, id, columns) + .await + .map_err(to_ctx_err) + } + + async fn flush(&self) -> ContextResult<()> { + self.client + .flush_generic_store(&self.store_name) + .await + .map_err(to_ctx_err) + } + + fn version(&self) -> u64 { + self.cached_version + } +} + pub struct RemoteDatagenStore { client: ContextClient, store_name: String, @@ -972,6 +1088,128 @@ impl ContextClient { Self::handle_response(resp).await } + // ----------------------------------------------------------- generic + + pub async fn create_generic_store( + &self, + req: &CreateGenericStoreRequest, + ) -> Result { + let resp = self + .http + .post(self.url("/generic")) + .json(req) + .send() + .await?; + Self::handle_response(resp).await + } + + pub async fn list_generic_stores(&self) -> Result { + let resp = self.http.get(self.url("/generic")).send().await?; + Self::handle_response(resp).await + } + + pub async fn get_generic_store(&self, name: &str) -> Result { + let resp = self + .http + .get(self.url(&format!("/generic/{name}"))) + .send() + .await?; + Self::handle_response(resp).await + } + + pub async fn delete_generic_store(&self, name: &str) -> Result<(), ClientError> { + let resp = self + .http + .delete(self.url(&format!("/generic/{name}"))) + .send() + .await?; + if resp.status().is_success() { + Ok(()) + } else { + Err(Self::extract_error(resp).await) + } + } + + /// Append rows. They are validated server-side against the store's schema. + pub async fn add_rows( + &self, + name: &str, + rows: &[Map], + ) -> Result { + let req = AddRowsRequest { + rows: rows.to_vec(), + }; + let resp = self + .http + .post(self.url(&format!("/generic/{name}/rows"))) + .json(&req) + .send() + .await?; + Self::handle_response(resp).await + } + + /// List rows. Blob columns are projected out server-side. + pub async fn list_rows( + &self, + name: &str, + filter: Option<&str>, + limit: Option, + offset: Option, + ) -> Result { + let mut query: Vec<(String, String)> = Vec::new(); + if let Some(filter) = filter { + query.push(("filter".to_string(), filter.to_string())); + } + if let Some(limit) = limit { + query.push(("limit".to_string(), limit.to_string())); + } + if let Some(offset) = offset { + query.push(("offset".to_string(), offset.to_string())); + } + let resp = self + .http + .get(self.url(&format!("/generic/{name}/rows"))) + .query(&query) + .send() + .await?; + Self::handle_response(resp).await + } + + /// Fetch one row. `columns` selects what to read; `None` omits blob + /// columns. Name a blob column explicitly to fetch its bytes. + pub async fn get_row( + &self, + name: &str, + id: &str, + columns: Option<&[String]>, + ) -> Result>, ClientError> { + let mut request = self + .http + .get(self.url(&format!("/generic/{name}/rows/{id}"))); + if let Some(columns) = columns { + request = request.query(&[("columns", columns.join(","))]); + } + let resp = request.send().await?; + if resp.status() == 404 { + return Ok(None); + } + Self::handle_response(resp).await.map(Some) + } + + /// Seal the store's active memtable so added rows become readable. + pub async fn flush_generic_store(&self, name: &str) -> Result<(), ClientError> { + let resp = self + .http + .post(self.url(&format!("/generic/{name}/flush"))) + .send() + .await?; + if resp.status().is_success() { + Ok(()) + } else { + Err(Self::extract_error(resp).await) + } + } + pub async fn create_datagen_store( &self, req: &CreateDatagenStoreRequest, diff --git a/crates/lance-context-core/src/api_impl.rs b/crates/lance-context-core/src/api_impl.rs index a33e2de..fc44b0e 100644 --- a/crates/lance-context-core/src/api_impl.rs +++ b/crates/lance-context-core/src/api_impl.rs @@ -4,14 +4,14 @@ use uuid::Uuid; use lance_context_api::{ AddDatagenEventsResponse, AddRecordRequest, AddRecordsResponse, AddRolloutRequest, - AddRolloutsResponse, CompactRequest, CompactResponse, CompactStatsResponse, ContextError, - ContextResult, ContextStoreApi, DatagenEventDto, DatagenFailureDto, DatagenFieldStateDto, - DatagenRootItemStatusesResponse, DatagenStepCursorDto, DatagenStoreApi, DatagenValueDto, - DeleteRecordResponse, FoldedDatagenItemDto, RecordDto, RecordPatchDto, RelationshipDto, - RetrieveRequest, RetrieveResultDto, RolloutRecordDto, RolloutStoreApi, SearchRequest, - SearchResultDto, StateMetadataDto, UpdateRecordRequest, UpdateRecordResponse, - UpsertRecordRequest, UpsertRecordResponse, UpsertRecordsRequest, UpsertRecordsResponse, - UpsertResultDto, + AddRolloutsResponse, AddRowsResponse, CompactRequest, CompactResponse, CompactStatsResponse, + ContextError, ContextResult, ContextStoreApi, DatagenEventDto, DatagenFailureDto, + DatagenFieldStateDto, DatagenRootItemStatusesResponse, DatagenStepCursorDto, DatagenStoreApi, + DatagenValueDto, DeleteRecordResponse, FoldedDatagenItemDto, GenericStoreApi, RecordDto, + RecordPatchDto, RelationshipDto, RetrieveRequest, RetrieveResultDto, RolloutRecordDto, + RolloutStoreApi, SchemaSpec, SearchRequest, SearchResultDto, StateMetadataDto, + UpdateRecordRequest, UpdateRecordResponse, UpsertRecordRequest, UpsertRecordResponse, + UpsertRecordsRequest, UpsertRecordsResponse, UpsertResultDto, }; use crate::datagen::{ @@ -20,6 +20,8 @@ use crate::datagen::{ DatagenStepKind, DatagenValue, FoldedDatagenItem, }; use crate::datagen_store::DatagenStore; +use crate::generic_codec::Row; +use crate::generic_store::GenericStore; use crate::record::{ ContextRecord, LifecycleQueryOptions, RecordFilters, RecordPatch, Relationship, StateMetadata, LIFECYCLE_ACTIVE, @@ -381,6 +383,49 @@ impl ContextStoreApi for ContextStore { } } +impl GenericStoreApi for GenericStore { + fn spec(&self) -> &SchemaSpec { + GenericStore::spec(self) + } + + async fn add(&self, rows: &[Row]) -> ContextResult { + let count = rows.len(); + let version = GenericStore::add(self, rows).await.map_err(to_ctx_err)?; + Ok(AddRowsResponse { version, count }) + } + + async fn list(&self, limit: Option, offset: Option) -> ContextResult> { + GenericStore::list(self, limit, offset) + .await + .map_err(to_ctx_err) + } + + async fn list_filtered( + &self, + filter: &str, + limit: Option, + offset: Option, + ) -> ContextResult> { + GenericStore::list_filtered(self, filter, limit, offset) + .await + .map_err(to_ctx_err) + } + + async fn get(&self, id: &str, columns: Option<&[String]>) -> ContextResult> { + GenericStore::get(self, id, columns) + .await + .map_err(to_ctx_err) + } + + async fn flush(&self) -> ContextResult<()> { + GenericStore::flush(self).await.map_err(to_ctx_err) + } + + fn version(&self) -> u64 { + GenericStore::version(self) + } +} + impl RolloutStoreApi for RolloutStore { async fn add(&mut self, records: &[AddRolloutRequest]) -> ContextResult { let mut ids = Vec::with_capacity(records.len()); @@ -457,7 +502,7 @@ impl RolloutStoreApi for RolloutStore { } } -fn rollout_record_from_add_request(r: &AddRolloutRequest) -> RolloutRecord { +pub fn rollout_record_from_add_request(r: &AddRolloutRequest) -> RolloutRecord { RolloutRecord { id: r.id.clone(), rollout_id: r.rollout_id.clone(), @@ -549,7 +594,7 @@ pub fn rollout_record_to_dto(r: RolloutRecord) -> RolloutRecordDto { } } -fn dto_to_relationship(r: RelationshipDto) -> Relationship { +pub fn dto_to_relationship(r: RelationshipDto) -> Relationship { Relationship { target_id: r.target_id, relation: r.relation, @@ -557,7 +602,7 @@ fn dto_to_relationship(r: RelationshipDto) -> Relationship { } } -fn relationship_to_dto(r: Relationship) -> RelationshipDto { +pub fn relationship_to_dto(r: Relationship) -> RelationshipDto { RelationshipDto { target_id: r.target_id, relation: r.relation, @@ -565,7 +610,7 @@ fn relationship_to_dto(r: Relationship) -> RelationshipDto { } } -fn patch_from_dto(patch: &RecordPatchDto) -> RecordPatch { +pub fn patch_from_dto(patch: &RecordPatchDto) -> RecordPatch { RecordPatch { bot_id: patch.bot_id.clone(), session_id: patch.session_id.clone(), @@ -597,7 +642,7 @@ fn patch_from_dto(patch: &RecordPatchDto) -> RecordPatch { } } -fn record_from_add_request(r: &AddRecordRequest, id: String, run_id: String) -> ContextRecord { +pub fn record_from_add_request(r: &AddRecordRequest, id: String, run_id: String) -> ContextRecord { ContextRecord { id, external_id: r.external_id.clone(), @@ -638,7 +683,7 @@ fn record_from_add_request(r: &AddRecordRequest, id: String, run_id: String) -> } } -fn record_to_dto(r: ContextRecord) -> RecordDto { +pub fn record_to_dto(r: ContextRecord) -> RecordDto { RecordDto { id: r.id, external_id: r.external_id, diff --git a/crates/lance-context-core/src/generic_codec.rs b/crates/lance-context-core/src/generic_codec.rs index dcbbc0d..82aab46 100644 --- a/crates/lance-context-core/src/generic_codec.rs +++ b/crates/lance-context-core/src/generic_codec.rs @@ -28,7 +28,7 @@ use arrow_array::{ use arrow_schema::{ArrowError, Schema}; use serde_json::{Map, Number, Value}; -use crate::schema_spec::{ColumnSpec, ColumnType, SchemaSpec, ID_COLUMN}; +use lance_context_api::schema_spec::{ColumnSpec, ColumnType, SchemaSpec, ID_COLUMN}; /// One row: column name to value. Absent keys are written as null, which is why /// a nullable column can simply be omitted. @@ -620,7 +620,7 @@ fn number(value: f64) -> Value { #[cfg(test)] mod tests { use super::*; - use crate::schema_spec::ColumnSpec; + use lance_context_api::schema_spec::ColumnSpec; use serde_json::json; fn spec() -> SchemaSpec { diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs index 5f4d955..baaf146 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -30,9 +30,9 @@ use lance::session::Session; use lance::{Error as LanceError, Result as LanceResult}; use crate::generic_codec::{batch_to_rows, rows_to_batch, Row}; -use crate::schema_spec::{SchemaSpec, ID_COLUMN}; use crate::store::{CompactionConfig, CompactionStats}; use crate::store_base::{ListSource, StorageBase, StorageBaseOptions}; +use lance_context_api::schema_spec::{SchemaSpec, ID_COLUMN}; /// Schema-metadata key holding the serialized [`SchemaSpec`], so a store can be /// reopened without the caller re-declaring its columns. @@ -413,9 +413,12 @@ impl GenericStore { } #[cfg(test)] +// `GenericStore` owns a `Dataset` and is not `Debug`, so `expect_err` (which +// formats the Ok value) is unavailable; `.err().expect(..)` is the alternative. +#[allow(clippy::err_expect)] mod tests { use super::*; - use crate::schema_spec::{ColumnSpec, ColumnType}; + use lance_context_api::schema_spec::{ColumnSpec, ColumnType}; use serde_json::json; use tempfile::TempDir; diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 40f2d6d..114a6c2 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -16,13 +16,17 @@ mod record; mod registry; mod rollout; mod rollout_store; -pub mod schema_spec; pub mod serde; mod storage; mod store; mod store_base; -pub use api_impl::rollout_record_to_dto; +// Request/DTO conversions, exported so the server does not keep its own copies. +// These were duplicated verbatim between here and `routes/`; see #214. +pub use api_impl::{ + dto_to_relationship, patch_from_dto, record_from_add_request, record_to_dto, + relationship_to_dto, rollout_record_from_add_request, rollout_record_to_dto, +}; pub use context::{Context, ContextEntry, Snapshot}; pub use datagen::{ datagen_event_id, datagen_failures, datagen_trajectory, fold_datagen_events, DatagenBlobValue, @@ -58,7 +62,8 @@ pub use rollout_store::{ RolloutStore, RolloutStoreOptions, SqlQueryResult, SQL_MAX_RESULT_ROWS, SQL_MAX_SCAN_ROWS, SQL_TABLE_NAME, }; -pub use schema_spec::{ColumnSpec, ColumnType, SchemaSpec, ID_COLUMN}; +// Schema declaration lives in the API crate: it is part of the wire contract. +pub use lance_context_api::{ColumnSpec, ColumnType, SchemaSpec, ID_COLUMN}; pub use storage::{create_local_dir_if_needed, join_uri, validate_store_name, MAX_STORE_NAME_LEN}; pub use store::{ CompactionConfig, CompactionStats, ContextStore, ContextStoreOptions, DistanceMetric, diff --git a/crates/lance-context-server/src/routes/generic.rs b/crates/lance-context-server/src/routes/generic.rs new file mode 100644 index 0000000..959e1d6 --- /dev/null +++ b/crates/lance-context-server/src/routes/generic.rs @@ -0,0 +1,521 @@ +//! HTTP surface for stores over user-declared schemas. +//! +//! Unlike the fixed-schema routes, rows are opaque JSON objects in both +//! directions: the schema is declared at store creation, so there is no static +//! DTO to convert through. Validation happens in the core codec against the +//! store's own schema, which is where the column list actually lives. + +use std::sync::Arc; + +use axum::extract::{Path, Query, State}; +use axum::http::StatusCode; +use axum::Json; +use lance_context_api::{ + AddRowsRequest, AddRowsResponse, CreateGenericStoreRequest, GenericStoreInfo, + ListGenericStoresResponse, ListRowsResponse, +}; +use lance_context_core::{GenericStore, GenericStoreOptions}; +use serde::Deserialize; +use tokio::sync::RwLock; + +use crate::error::AppError; +use crate::state::AppState; + +/// Upper bound on one generic append request body. +/// +/// Generic stores hold blob columns inline — that is the point of them — so +/// this matches the rollout ceiling rather than the smaller datagen one. +pub const MAX_GENERIC_UPLOAD_BYTES: usize = 1024 * 1024 * 1024; + +/// `POST /api/v1/generic` +pub async fn create_generic_store( + State(state): State>, + Json(req): Json, +) -> Result<(StatusCode, Json), AppError> { + AppState::validate_name(&req.name)?; + // Reject an invalid schema before creating anything, so a bad declaration + // cannot leave an empty dataset behind. + req.schema.validate().map_err(AppError::InvalidRequest)?; + + if state + .generic_registry + .write() + .await + .contains(&req.name) + .await + .map_err(AppError::from_lance)? + { + return Err(AppError::AlreadyExists(format!( + "Generic store '{}' already exists", + req.name + ))); + } + + let uri = state.generic_uri(&req.name); + let options = GenericStoreOptions { + storage_options: req.storage_options, + shard_id: state.instance_id.clone(), + merge_after_generations: None, + session: None, + seal_on_add: req.seal_on_add, + }; + let store = GenericStore::open(&uri, req.schema.clone(), options) + .await + .map_err(AppError::from_lance)?; + let version = store.version(); + + let store = Arc::new(RwLock::new(store)); + state.register_generic(&req.name, &uri, store).await?; + + Ok(( + StatusCode::CREATED, + Json(GenericStoreInfo { + name: req.name, + uri, + version: Some(version), + schema: Some(req.schema), + }), + )) +} + +/// `GET /api/v1/generic` +/// +/// Served from the registry without opening each dataset, so `version` and +/// `schema` are omitted — both live in the dataset itself. +pub async fn list_generic_stores( + State(state): State>, +) -> Result, AppError> { + let entries = state + .generic_registry + .write() + .await + .list() + .await + .map_err(AppError::from_lance)?; + + Ok(Json(ListGenericStoresResponse { + stores: entries + .into_iter() + .map(|entry| GenericStoreInfo { + name: entry.name, + uri: entry.uri, + version: None, + schema: None, + }) + .collect(), + })) +} + +/// `GET /api/v1/generic/{name}` +pub async fn get_generic_store( + State(state): State>, + Path(name): Path, +) -> Result, AppError> { + let store = state.get_or_open_generic_store(&name).await?; + let guard = store.read().await; + Ok(Json(GenericStoreInfo { + name: name.clone(), + uri: state.generic_uri(&name), + version: Some(guard.version()), + schema: Some(guard.spec().clone()), + })) +} + +/// `DELETE /api/v1/generic/{name}` +pub async fn delete_generic_store( + State(state): State>, + Path(name): Path, +) -> Result { + if state.unregister_generic(&name).await? { + Ok(StatusCode::NO_CONTENT) + } else { + Err(AppError::NotFound(format!( + "Generic store '{name}' does not exist" + ))) + } +} + +/// `POST /api/v1/generic/{name}/rows` +/// +/// Rows are validated against the store's schema: an undeclared column is an +/// error rather than being dropped, and an omitted nullable column is null. +pub async fn add_rows( + State(state): State>, + Path(name): Path, + Json(req): Json, +) -> Result<(StatusCode, Json), AppError> { + let store = state.get_or_open_generic_store(&name).await?; + // Read lock: `add` takes `&self`, so appends do not serialize here. + let guard = store.read().await; + let count = req.rows.len(); + let version = guard.add(&req.rows).await.map_err(AppError::from_lance)?; + Ok(( + StatusCode::CREATED, + Json(AddRowsResponse { version, count }), + )) +} + +/// Query parameters for listing rows. +#[derive(Debug, Deserialize)] +pub struct ListRowsQuery { + #[serde(default)] + pub limit: Option, + #[serde(default)] + pub offset: Option, + /// SQL predicate over the store's own columns. + #[serde(default)] + pub filter: Option, +} + +/// `GET /api/v1/generic/{name}/rows` +/// +/// Blob columns are projected out, so a list never materializes a large +/// payload; fetch those per row via [`get_row`]. +pub async fn list_rows( + State(state): State>, + Path(name): Path, + Query(query): Query, +) -> Result, AppError> { + let store = state.get_or_open_generic_store(&name).await?; + let guard = store.read().await; + + let rows = match &query.filter { + Some(filter) => guard + .list_filtered(filter, query.limit, query.offset) + .await + .map_err(AppError::from_lance)?, + None => guard + .list(query.limit, query.offset) + .await + .map_err(AppError::from_lance)?, + }; + Ok(Json(ListRowsResponse { rows })) +} + +/// Query parameters for a point read. +#[derive(Debug, Deserialize)] +pub struct GetRowQuery { + /// Comma-separated columns to read. Absent reads everything except blob + /// columns; name a blob column explicitly to fetch it. + #[serde(default)] + pub columns: Option, +} + +/// `GET /api/v1/generic/{name}/rows/{id}` +pub async fn get_row( + State(state): State>, + Path((name, id)): Path<(String, String)>, + Query(query): Query, +) -> Result>, AppError> { + let store = state.get_or_open_generic_store(&name).await?; + let guard = store.read().await; + + let columns: Option> = query.columns.as_ref().map(|raw| { + raw.split(',') + .map(str::trim) + .filter(|column| !column.is_empty()) + .map(str::to_string) + .collect() + }); + + let row = guard + .get(&id, columns.as_deref()) + .await + .map_err(AppError::from_lance)? + .ok_or_else(|| AppError::NotFound(format!("Row '{id}' does not exist")))?; + Ok(Json(row)) +} + +/// `POST /api/v1/generic/{name}/flush` +/// +/// Seals the active memtable so previously added rows become readable. Needed +/// only when the store was created with `seal_on_add: false`. +pub async fn flush_generic_store( + State(state): State>, + Path(name): Path, +) -> Result { + let store = state.get_or_open_generic_store(&name).await?; + let guard = store.read().await; + guard.flush().await.map_err(AppError::from_lance)?; + Ok(StatusCode::NO_CONTENT) +} + +/// `POST /api/v1/generic/{name}/merge-wal` +/// +/// Folds this instance's flushed MemWAL generations into the base table. +/// Returns how many were reclaimed. +pub async fn merge_generic_wal( + State(state): State>, + Path(name): Path, +) -> Result, AppError> { + let store = state.get_or_open_generic_store(&name).await?; + let mut guard = store.write().await; + let reclaimed = guard.cleanup_wal().await.map_err(AppError::from_lance)?; + Ok(Json(serde_json::json!({ "reclaimed": reclaimed }))) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use lance_context_api::{ColumnSpec, ColumnType, SchemaSpec, ID_COLUMN}; + use tempfile::TempDir; + + use super::*; + use crate::state::AppState; + + fn spec() -> SchemaSpec { + SchemaSpec::new(vec![ + ( + ID_COLUMN.to_string(), + ColumnSpec::required(ColumnType::String { large: false }), + ), + ( + "user".to_string(), + ColumnSpec::new(ColumnType::String { large: false }), + ), + ( + "blob".to_string(), + ColumnSpec::new(ColumnType::Binary { blob: true }), + ), + ]) + } + + async fn test_state() -> (Arc, TempDir) { + let dir = TempDir::new().unwrap(); + let state = Arc::new(AppState::new_for_test(dir.path().to_path_buf()).await); + (state, dir) + } + + fn row(value: serde_json::Value) -> serde_json::Map { + value.as_object().unwrap().clone() + } + + #[tokio::test] + async fn create_list_add_and_read_round_trip() { + let (state, _dir) = test_state().await; + + let (status, Json(info)) = create_generic_store( + State(state.clone()), + Json(CreateGenericStoreRequest { + name: "s1".to_string(), + schema: spec(), + storage_options: None, + seal_on_add: true, + }), + ) + .await + .unwrap(); + assert_eq!(status, StatusCode::CREATED); + assert!(info.schema.is_some(), "create echoes the schema back"); + + let _ = add_rows( + State(state.clone()), + Path("s1".to_string()), + Json(AddRowsRequest { + rows: vec![ + row(serde_json::json!({"id": "r1", "user": "u1", "blob": [1, 2, 3]})), + row(serde_json::json!({"id": "r2", "user": "u2"})), + ], + }), + ) + .await + .unwrap(); + + let Json(listed) = list_rows( + State(state.clone()), + Path("s1".to_string()), + Query(ListRowsQuery { + limit: None, + offset: None, + filter: None, + }), + ) + .await + .unwrap(); + assert_eq!(listed.rows.len(), 2); + assert!( + !listed.rows[0].contains_key("blob"), + "listing must not materialize blob columns" + ); + + // A blob is fetchable per row by naming the column. + let Json(fetched) = get_row( + State(state.clone()), + Path(("s1".to_string(), "r1".to_string())), + Query(GetRowQuery { + columns: Some("blob".to_string()), + }), + ) + .await + .unwrap(); + assert_eq!(fetched["blob"], serde_json::json!("AQID")); + } + + #[tokio::test] + async fn invalid_schema_is_rejected_before_the_store_is_created() { + let (state, _dir) = test_state().await; + + // No `id` column. + let bad = SchemaSpec::new(vec![( + "name".to_string(), + ColumnSpec::new(ColumnType::String { large: false }), + )]); + let err = create_generic_store( + State(state.clone()), + Json(CreateGenericStoreRequest { + name: "bad".to_string(), + schema: bad, + storage_options: None, + seal_on_add: true, + }), + ) + .await + .expect_err("an invalid schema must be rejected"); + assert!(matches!(err, AppError::InvalidRequest(_)), "{err:?}"); + + // And nothing was registered. + let Json(listed) = list_generic_stores(State(state)).await.unwrap(); + assert!(listed.stores.is_empty()); + } + + #[tokio::test] + async fn duplicate_create_conflicts_and_missing_store_is_not_found() { + let (state, _dir) = test_state().await; + let request = || CreateGenericStoreRequest { + name: "s1".to_string(), + schema: spec(), + storage_options: None, + seal_on_add: true, + }; + + let _ = create_generic_store(State(state.clone()), Json(request())) + .await + .unwrap(); + let err = create_generic_store(State(state.clone()), Json(request())) + .await + .expect_err("a duplicate name must conflict"); + assert!(matches!(err, AppError::AlreadyExists(_)), "{err:?}"); + + let err = get_generic_store(State(state.clone()), Path("absent".to_string())) + .await + .expect_err("an unknown store must 404"); + assert!(matches!(err, AppError::NotFound(_)), "{err:?}"); + } + + #[tokio::test] + async fn undeclared_columns_are_rejected_rather_than_dropped() { + // A field-name typo must fail the write, not become missing data. + let (state, _dir) = test_state().await; + let _ = create_generic_store( + State(state.clone()), + Json(CreateGenericStoreRequest { + name: "s1".to_string(), + schema: spec(), + storage_options: None, + seal_on_add: true, + }), + ) + .await + .unwrap(); + + let err = add_rows( + State(state.clone()), + Path("s1".to_string()), + Json(AddRowsRequest { + rows: vec![row(serde_json::json!({"id": "r1", "usr": "typo"}))], + }), + ) + .await + .expect_err("an undeclared column must be rejected"); + assert!( + format!("{err:?}").contains("not declared"), + "expected an undeclared-column error, got {err:?}" + ); + } + + #[tokio::test] + async fn delete_removes_the_store_from_the_registry() { + let (state, _dir) = test_state().await; + let _ = create_generic_store( + State(state.clone()), + Json(CreateGenericStoreRequest { + name: "s1".to_string(), + schema: spec(), + storage_options: None, + seal_on_add: true, + }), + ) + .await + .unwrap(); + + assert_eq!( + delete_generic_store(State(state.clone()), Path("s1".to_string())) + .await + .unwrap(), + StatusCode::NO_CONTENT + ); + let Json(listed) = list_generic_stores(State(state.clone())).await.unwrap(); + assert!(listed.stores.is_empty()); + + let err = delete_generic_store(State(state), Path("s1".to_string())) + .await + .expect_err("deleting twice must 404"); + assert!(matches!(err, AppError::NotFound(_)), "{err:?}"); + } + + #[tokio::test] + async fn deferred_seal_needs_a_flush_and_wal_merges() { + let (state, _dir) = test_state().await; + let _ = create_generic_store( + State(state.clone()), + Json(CreateGenericStoreRequest { + name: "s1".to_string(), + schema: spec(), + storage_options: None, + // Default profile: durable on return, not yet visible. + seal_on_add: false, + }), + ) + .await + .unwrap(); + + let _ = add_rows( + State(state.clone()), + Path("s1".to_string()), + Json(AddRowsRequest { + rows: vec![row(serde_json::json!({"id": "r1"}))], + }), + ) + .await + .unwrap(); + + let query = || ListRowsQuery { + limit: None, + offset: None, + filter: None, + }; + let Json(before) = list_rows(State(state.clone()), Path("s1".to_string()), Query(query())) + .await + .unwrap(); + assert_eq!(before.rows.len(), 0, "deferred seal must not publish yet"); + + flush_generic_store(State(state.clone()), Path("s1".to_string())) + .await + .unwrap(); + let Json(after) = list_rows(State(state.clone()), Path("s1".to_string()), Query(query())) + .await + .unwrap(); + assert_eq!(after.rows.len(), 1); + + // And the WAL generation folds into the base table. + let Json(merged) = merge_generic_wal(State(state.clone()), Path("s1".to_string())) + .await + .unwrap(); + assert!(merged["reclaimed"].as_u64().unwrap() > 0); + let Json(kept) = list_rows(State(state), Path("s1".to_string()), Query(query())) + .await + .unwrap(); + assert_eq!(kept.rows.len(), 1, "rows survive the merge"); + } +} diff --git a/crates/lance-context-server/src/routes/mod.rs b/crates/lance-context-server/src/routes/mod.rs index 17ab77d..ec7e10c 100644 --- a/crates/lance-context-server/src/routes/mod.rs +++ b/crates/lance-context-server/src/routes/mod.rs @@ -1,6 +1,7 @@ pub mod compact; pub mod contexts; pub mod datagen; +pub mod generic; pub mod health; pub mod records; pub mod rollouts; @@ -146,6 +147,27 @@ pub fn router() -> Router> { "/api/v1/datagen/{name}/blobs/{event_id}", get(datagen::fetch_datagen_blob), ) + .route("/api/v1/generic", post(generic::create_generic_store)) + .route("/api/v1/generic", get(generic::list_generic_stores)) + .route("/api/v1/generic/{name}", get(generic::get_generic_store)) + .route( + "/api/v1/generic/{name}", + delete(generic::delete_generic_store), + ) + .route( + "/api/v1/generic/{name}/rows", + post(generic::add_rows).layer(DefaultBodyLimit::max(generic::MAX_GENERIC_UPLOAD_BYTES)), + ) + .route("/api/v1/generic/{name}/rows", get(generic::list_rows)) + .route("/api/v1/generic/{name}/rows/{id}", get(generic::get_row)) + .route( + "/api/v1/generic/{name}/flush", + post(generic::flush_generic_store), + ) + .route( + "/api/v1/generic/{name}/merge-wal", + post(generic::merge_generic_wal), + ) } #[cfg(test)] diff --git a/crates/lance-context-server/src/routes/records.rs b/crates/lance-context-server/src/routes/records.rs index 2365607..2698377 100644 --- a/crates/lance-context-server/src/routes/records.rs +++ b/crates/lance-context-server/src/routes/records.rs @@ -5,16 +5,14 @@ use axum::extract::{Path, Query, State}; use axum::http::header; use axum::response::Response; use axum::Json; -use chrono::Utc; use lance_context_api::{ AddRecordsRequest, AddRecordsResponse, DeleteRecordResponse, GetRecordResponse, - ListRecordsResponse, RecordDto, RecordPatchDto, RelationshipDto, StateMetadataDto, - UpdateRecordRequest, UpdateRecordResponse, UpsertRecordRequest, UpsertRecordResponse, - UpsertRecordsRequest, UpsertRecordsResponse, UpsertResultDto, + ListRecordsResponse, RecordDto, UpdateRecordRequest, UpdateRecordResponse, UpsertRecordRequest, + UpsertRecordResponse, UpsertRecordsRequest, UpsertRecordsResponse, UpsertResultDto, }; use lance_context_core::{ - ContextRecord, LifecycleQueryOptions, RecordFilters, RecordPatch, Relationship, StateMetadata, - LIFECYCLE_ACTIVE, + patch_from_dto, record_from_add_request, record_to_dto, ContextRecord, LifecycleQueryOptions, + RecordFilters, }; use uuid::Uuid; @@ -406,139 +404,6 @@ pub async fn related_records( Ok(Json(ListRecordsResponse { records: dtos })) } -pub fn record_to_dto(r: ContextRecord) -> RecordDto { - RecordDto { - id: r.id, - external_id: r.external_id, - run_id: r.run_id, - bot_id: r.bot_id, - session_id: r.session_id, - tenant: r.tenant, - source: r.source, - created_at: r.created_at, - role: r.role, - content_type: r.content_type, - text_payload: r.text_payload, - binary_payload: r.binary_payload, - payload_uri: r.payload_uri, - payload_size: r.payload_size, - payload_checksum: r.payload_checksum, - embedding: r.embedding, - state_metadata: r.state_metadata.map(|sm| StateMetadataDto { - step: sm.step, - active_plan_id: sm.active_plan_id, - tokens_used: sm.tokens_used, - custom: sm.custom, - }), - metadata: r.metadata, - relationships: r - .relationships - .into_iter() - .map(relationship_to_dto) - .collect(), - expires_at: r.expires_at, - retention_policy: r.retention_policy, - lifecycle_status: r.lifecycle_status, - retired_at: r.retired_at, - retired_reason: r.retired_reason, - supersedes_id: r.supersedes_id, - superseded_by_id: r.superseded_by_id, - } -} - -fn dto_to_relationship(r: RelationshipDto) -> Relationship { - Relationship { - target_id: r.target_id, - relation: r.relation, - weight: r.weight, - } -} - -fn relationship_to_dto(r: Relationship) -> RelationshipDto { - RelationshipDto { - target_id: r.target_id, - relation: r.relation, - weight: r.weight, - } -} - -fn patch_from_dto(patch: &RecordPatchDto) -> RecordPatch { - RecordPatch { - bot_id: patch.bot_id.clone(), - session_id: patch.session_id.clone(), - tenant: patch.tenant.clone(), - source: patch.source.clone(), - state_metadata: patch.state_metadata.as_ref().map(|sm| StateMetadata { - step: sm.step, - active_plan_id: sm.active_plan_id.clone(), - tokens_used: sm.tokens_used, - custom: sm.custom.clone(), - }), - metadata: patch.metadata.clone(), - relationships: patch.relationships.as_ref().map(|relationships| { - relationships - .iter() - .cloned() - .map(dto_to_relationship) - .collect() - }), - expires_at: patch.expires_at, - retention_policy: patch.retention_policy.clone(), - lifecycle_status: patch.lifecycle_status.clone(), - retired_at: patch.retired_at, - retired_reason: patch.retired_reason.clone(), - embedding: patch.embedding.clone(), - payload_uri: patch.payload_uri.clone(), - payload_size: patch.payload_size, - payload_checksum: patch.payload_checksum.clone(), - } -} - -fn record_from_add_request( - r: &lance_context_api::AddRecordRequest, - id: String, - run_id: String, -) -> ContextRecord { - ContextRecord { - id, - external_id: r.external_id.clone(), - run_id, - bot_id: r.bot_id.clone(), - session_id: r.session_id.clone(), - tenant: r.tenant.clone(), - source: r.source.clone(), - created_at: Utc::now(), - role: r.role.clone(), - state_metadata: r.state_metadata.as_ref().map(|sm| StateMetadata { - step: sm.step, - active_plan_id: sm.active_plan_id.clone(), - tokens_used: sm.tokens_used, - custom: sm.custom.clone(), - }), - metadata: r.metadata.clone(), - relationships: r - .relationships - .iter() - .cloned() - .map(dto_to_relationship) - .collect(), - expires_at: r.expires_at, - retention_policy: r.retention_policy.clone(), - lifecycle_status: LIFECYCLE_ACTIVE.to_string(), - retired_at: None, - retired_reason: None, - supersedes_id: r.supersedes_id.clone(), - superseded_by_id: None, - content_type: r.content_type.clone(), - text_payload: r.text_payload.clone(), - binary_payload: r.binary_payload.clone(), - payload_uri: r.payload_uri.clone(), - payload_size: r.payload_size, - payload_checksum: r.payload_checksum.clone(), - embedding: r.embedding.clone(), - } -} - #[cfg(test)] mod tests { use std::sync::Arc; @@ -547,7 +412,7 @@ mod tests { use axum::Json; use chrono::{Duration, Utc}; use lance_context_api::{ - AddRecordRequest, AddRecordsRequest, RecordPatchDto, UpdateRecordRequest, + AddRecordRequest, AddRecordsRequest, RecordPatchDto, RelationshipDto, UpdateRecordRequest, UpsertRecordRequest, UpsertRecordsRequest, }; use lance_context_core::ContextStore; diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index 745c087..d0ebc9d 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -6,16 +6,14 @@ use axum::extract::{FromRequest, Multipart, Path, Query, Request, State}; use axum::http::{header, StatusCode}; use axum::response::Response; use axum::Json; -use chrono::Utc; use lance_context_api::{ AddRolloutRequest, AddRolloutsRequest, AddRolloutsResponse, CheckoutRequest, CompactRequest, CompactResponse, CompactStatsResponse, CreateRolloutStoreRequest, GetRolloutResponse, - ListRolloutStoresResponse, ListRolloutsResponse, RelationshipDto, RolloutRecordDto, - RolloutStoreInfo, VersionResponse, + ListRolloutStoresResponse, ListRolloutsResponse, RolloutStoreInfo, VersionResponse, }; +use lance_context_core::{rollout_record_from_add_request, rollout_record_to_dto}; use lance_context_core::{ - CompactionConfig, Relationship, RolloutFilters, RolloutRecord, RolloutStore, - RolloutStoreOptions, + CompactionConfig, RolloutFilters, RolloutRecord, RolloutStore, RolloutStoreOptions, }; use tokio::sync::RwLock; @@ -692,113 +690,6 @@ fn content_type_is(content_type: &str, expected: &str) -> bool { .is_some_and(|mime| mime.eq_ignore_ascii_case(expected)) } -fn dto_to_relationship(r: RelationshipDto) -> Relationship { - Relationship { - target_id: r.target_id, - relation: r.relation, - weight: r.weight, - } -} - -fn relationship_to_dto(r: Relationship) -> RelationshipDto { - RelationshipDto { - target_id: r.target_id, - relation: r.relation, - weight: r.weight, - } -} - -fn rollout_record_from_add_request(r: &AddRolloutRequest) -> RolloutRecord { - RolloutRecord { - id: r.id.clone(), - rollout_id: r.rollout_id.clone(), - problem_id: r.problem_id.clone().unwrap_or_else(|| r.rollout_id.clone()), - dataset: r.dataset.clone(), - sequence_order: r.sequence_order, - role: r.role.clone(), - created_at: r.created_at.unwrap_or_else(Utc::now), - content: r.content.clone(), - content_type: r.content_type.clone(), - model_input_string: r.model_input_string.clone(), - model_output_string: r.model_output_string.clone(), - rationale: r.rationale.clone(), - problem_text: r.problem_text.clone(), - user_metadata: r.user_metadata.clone(), - input_tokens: r.input_tokens.clone(), - output_tokens: r.output_tokens.clone(), - num_input_tokens: r.num_input_tokens, - num_output_tokens: r.num_output_tokens, - output_logprobs: r.output_logprobs.clone(), - input_logprobs: r.input_logprobs.clone(), - ref_logprobs: r.ref_logprobs.clone(), - loss_mask: r.loss_mask.clone(), - advantage: r.advantage, - reward: r.reward, - raw_reward: r.raw_reward, - grader_id: r.grader_id.clone(), - score: r.score, - include_in_training: r.include_in_training, - exclude_reason: r.exclude_reason.clone(), - policy_version: r.policy_version.clone(), - relationships: r - .relationships - .iter() - .cloned() - .map(dto_to_relationship) - .collect(), - binary_payload: r.binary_payload.clone(), - payload_size: r.payload_size, - payload_checksum: r.payload_checksum.clone(), - artifact_type: r.artifact_type.clone(), - metadata: r.metadata.clone(), - } -} - -fn rollout_record_to_dto(r: RolloutRecord) -> RolloutRecordDto { - RolloutRecordDto { - id: r.id, - rollout_id: r.rollout_id, - problem_id: r.problem_id, - dataset: r.dataset, - sequence_order: r.sequence_order, - role: r.role, - created_at: r.created_at, - content: r.content, - content_type: r.content_type, - model_input_string: r.model_input_string, - model_output_string: r.model_output_string, - rationale: r.rationale, - problem_text: r.problem_text, - user_metadata: r.user_metadata, - input_tokens: r.input_tokens, - output_tokens: r.output_tokens, - num_input_tokens: r.num_input_tokens, - num_output_tokens: r.num_output_tokens, - output_logprobs: r.output_logprobs, - input_logprobs: r.input_logprobs, - ref_logprobs: r.ref_logprobs, - loss_mask: r.loss_mask, - advantage: r.advantage, - reward: r.reward, - raw_reward: r.raw_reward, - grader_id: r.grader_id, - score: r.score, - include_in_training: r.include_in_training, - exclude_reason: r.exclude_reason, - policy_version: r.policy_version, - relationships: r - .relationships - .into_iter() - .map(relationship_to_dto) - .collect(), - binary_payload: r.binary_payload, - payload_size: r.payload_size, - payload_checksum: r.payload_checksum, - artifact_type: r.artifact_type, - metadata: r.metadata, - } -} - #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/crates/lance-context-server/src/routes/search.rs b/crates/lance-context-server/src/routes/search.rs index 1efc79c..f08c348 100644 --- a/crates/lance-context-server/src/routes/search.rs +++ b/crates/lance-context-server/src/routes/search.rs @@ -9,8 +9,8 @@ use lance_context_api::{ use lance_context_core::{LifecycleQueryOptions, RecordFilters}; use crate::error::AppError; -use crate::routes::records::record_to_dto; use crate::state::AppState; +use lance_context_core::record_to_dto; pub async fn search( State(state): State>, diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index d0dd394..51b7b7a 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -6,7 +6,8 @@ use std::time::Duration; use lance_context_core::{ join_uri, validate_store_name, ContextStore, ContextStoreOptions, DatagenStore, - DatagenStoreOptions, RolloutRegistry, RolloutStore, RolloutStoreOptions, Session, + DatagenStoreOptions, GenericStore, GenericStoreOptions, RolloutRegistry, RolloutStore, + RolloutStoreOptions, Session, }; use lru::LruCache; use tokio::sync::{Mutex, RwLock}; @@ -83,6 +84,18 @@ pub struct AppState { /// dataset from [`Self::rollout_registry`] so the two store kinds never /// collide on a shared name. pub datagen_registry: RwLock, + /// Bounded LRU of resident generic-store handles, mirroring + /// [`Self::rollout_stores`]. + pub generic_stores: Mutex>>>, + /// Durable directory of which generic stores exist. + /// + /// Records only name and URI: a generic store's **schema is stored in its + /// own dataset**, so the registry deliberately does not carry a second + /// copy. Two writable copies of a schema is exactly the divergence this + /// refactor exists to remove. The cost is that listing stores cannot report + /// their schemas without opening each dataset, which is why + /// `GenericStoreInfo::schema` is `None` in list responses. + pub generic_registry: RwLock, } /// Process-wide admission control for the total artifact-blob payload held in @@ -194,6 +207,10 @@ impl AppState { let blob_budget = (config.rollout_max_inflight_blob_bytes > 0) .then(|| BlobBudget::new(config.rollout_max_inflight_blob_bytes)); let rollout_session = build_rollout_session(config.rollout_cache_bytes); + let generic_registry_uri = join_uri(&base_uri, "_registry.generic.lance"); + let generic_registry = RolloutRegistry::open_or_create(&generic_registry_uri, None) + .await + .map_err(AppError::from_lance)?; let datagen_registry_uri = join_uri(&base_uri, "_registry.datagen.lance"); let datagen_registry = RolloutRegistry::open_or_create(&datagen_registry_uri, None) .await @@ -211,6 +228,8 @@ impl AppState { rollout_session, datagen_stores: Mutex::new(LruCache::new(capacity)), datagen_registry: RwLock::new(datagen_registry), + generic_stores: Mutex::new(LruCache::new(capacity)), + generic_registry: RwLock::new(generic_registry), }) } @@ -237,6 +256,10 @@ impl AppState { let registry = RolloutRegistry::open_or_create(®istry_uri, None) .await .expect("open test registry"); + let generic_registry_uri = join_uri(&base_uri, "_registry.generic.lance"); + let generic_registry = RolloutRegistry::open_or_create(&generic_registry_uri, None) + .await + .expect("open test generic registry"); let datagen_registry_uri = join_uri(&base_uri, "_registry.datagen.lance"); let datagen_registry = RolloutRegistry::open_or_create(&datagen_registry_uri, None) .await @@ -254,6 +277,10 @@ impl AppState { rollout_flush_interval_secs: 0, blob_budget: None, rollout_session: build_rollout_session(2 * 1024 * 1024 * 1024), + generic_stores: Mutex::new(LruCache::new( + std::num::NonZeroUsize::new(DEFAULT_ROLLOUT_CACHE_CAPACITY).unwrap(), + )), + generic_registry: RwLock::new(generic_registry), datagen_stores: Mutex::new(LruCache::new( NonZeroUsize::new(DEFAULT_ROLLOUT_CACHE_CAPACITY).unwrap(), )), @@ -518,6 +545,113 @@ impl AppState { Ok(opened) } + /// Physical URI of a generic store's dataset. + #[must_use] + pub fn generic_uri(&self, name: &str) -> String { + join_uri(&self.base_uri, &format!("{}.generic.lance", name)) + } + + fn generic_store_options(&self, seal_on_add: bool) -> GenericStoreOptions { + GenericStoreOptions { + storage_options: None, + shard_id: self.instance_id.clone(), + merge_after_generations: (self.rollout_merge_after_generations > 0) + .then_some(self.rollout_merge_after_generations), + session: self.rollout_session.clone(), + seal_on_add, + } + } + + /// Record that a generic store exists, in both the durable registry and the + /// in-memory LRU. + pub async fn register_generic( + &self, + name: &str, + uri: &str, + store: Arc>, + ) -> Result<(), AppError> { + Self::validate_name(name)?; + self.generic_registry + .write() + .await + .upsert(name, uri) + .await + .map_err(AppError::from_lance)?; + self.generic_stores + .lock() + .await + .put(name.to_string(), store); + Ok(()) + } + + /// Remove a generic store from the registry and evict any resident handle. + /// Returns whether the store existed. + pub async fn unregister_generic(&self, name: &str) -> Result { + Self::validate_name(name)?; + let existed = self + .generic_registry + .write() + .await + .contains(name) + .await + .map_err(AppError::from_lance)?; + if !existed { + return Ok(false); + } + self.generic_registry + .write() + .await + .remove(name) + .await + .map_err(AppError::from_lance)?; + self.generic_stores.lock().await.pop(name); + Ok(true) + } + + /// Look up a generic store by name, lazily reopening it on a cache miss. + /// + /// The schema is not passed in: it is read back from the dataset, which is + /// where it durably lives. See [`Self::get_or_open_datagen_store`] for the + /// residency rationale. + pub async fn get_or_open_generic_store( + &self, + name: &str, + ) -> Result>, AppError> { + Self::validate_name(name)?; + if let Some(store) = self.generic_stores.lock().await.get(name) { + return Ok(store.clone()); + } + + let exists = self + .generic_registry + .write() + .await + .contains(name) + .await + .map_err(AppError::from_lance)?; + if !exists { + return Err(AppError::NotFound(format!( + "Generic store '{name}' does not exist" + ))); + } + + let uri = self.generic_uri(name); + // `seal_on_add` is a per-store property that is not persisted, so a + // reopened store takes the server default. Callers needing + // read-your-write should flush explicitly. + let opened = GenericStore::open_existing(&uri, self.generic_store_options(false)) + .await + .map_err(AppError::from_lance)?; + let opened = Arc::new(RwLock::new(opened)); + + let mut cache = self.generic_stores.lock().await; + if let Some(existing) = cache.get(name) { + return Ok(existing.clone()); + } + cache.put(name.to_string(), opened.clone()); + Ok(opened) + } + /// Spawn the single, process-wide WAL-cleanup sweeper. /// /// This replaces the former one-timer-per-store model, which does not scale @@ -788,6 +922,44 @@ impl AppState { ); } } + + // Datagen and generic stores hold resident MemWAL writers too. Skipping + // them left an unsealed memtable to the best-effort detached close in + // `Drop`, so rows could stay durable-but-invisible until the next WAL + // replay. Closing here makes shutdown deterministic for every kind. + let datagen: Vec<(String, Arc>)> = { + let cache = self.datagen_stores.lock().await; + cache + .iter() + .map(|(name, store)| (name.clone(), store.clone())) + .collect() + }; + for (name, store) in datagen { + if let Err(e) = store.write().await.close().await { + tracing::warn!( + store = %name, + error = %e, + "failed to close datagen writer during shutdown" + ); + } + } + + let generic: Vec<(String, Arc>)> = { + let cache = self.generic_stores.lock().await; + cache + .iter() + .map(|(name, store)| (name.clone(), store.clone())) + .collect() + }; + for (name, store) in generic { + if let Err(e) = store.write().await.close().await { + tracing::warn!( + store = %name, + error = %e, + "failed to close generic writer during shutdown" + ); + } + } } } diff --git a/crates/lance-context/src/lib.rs b/crates/lance-context/src/lib.rs index 063ccfe..c60616d 100644 --- a/crates/lance-context/src/lib.rs +++ b/crates/lance-context/src/lib.rs @@ -15,18 +15,20 @@ pub use lance_context_core::{ pub use lance_context_api::{ AddDatagenEventsRequest, AddDatagenEventsResponse, AddRecordRequest, AddRecordsResponse, - AddRolloutRequest, AddRolloutsResponse, CompactRequest, CompactResponse, CompactStatsResponse, - ContextError, ContextResult, ContextStoreApi, CreateDatagenStoreRequest, + AddRolloutRequest, AddRolloutsResponse, AddRowsRequest, AddRowsResponse, ColumnSpec, + ColumnType, CompactRequest, CompactResponse, CompactStatsResponse, ContextError, ContextResult, + ContextStoreApi, CreateDatagenStoreRequest, CreateGenericStoreRequest, CreateRolloutStoreRequest, DatagenEventDto, DatagenFailureDto, DatagenFieldStateDto, DatagenRootItemStatusesResponse, DatagenStepCursorDto, DatagenStoreApi, DatagenValueDto, - DeleteRecordResponse, FoldedDatagenItemDto, RecordDto, RelationshipDto, RetrieveRequest, - RetrieveResponse, RetrieveResultDto, RolloutRecordDto, RolloutStoreApi, SearchResultDto, - UpsertRecordRequest, UpsertRecordResponse, + DeleteRecordResponse, FoldedDatagenItemDto, GenericStoreApi, GenericStoreInfo, RecordDto, + RelationshipDto, RetrieveRequest, RetrieveResponse, RetrieveResultDto, RolloutRecordDto, + RolloutStoreApi, SchemaSpec, SearchResultDto, UpsertRecordRequest, UpsertRecordResponse, + ID_COLUMN, }; #[cfg(feature = "remote")] pub use lance_context_client::{ - ClientError, RemoteContextStore, RemoteDatagenStore, RemoteRolloutStore, + ClientError, RemoteContextStore, RemoteDatagenStore, RemoteGenericStore, RemoteRolloutStore, }; mod unified; @@ -37,3 +39,6 @@ pub use unified_rollout::RolloutStore; mod unified_datagen; pub use unified_datagen::DatagenStore; + +mod unified_generic; +pub use unified_generic::GenericStore; diff --git a/crates/lance-context/src/unified_generic.rs b/crates/lance-context/src/unified_generic.rs new file mode 100644 index 0000000..86119d8 --- /dev/null +++ b/crates/lance-context/src/unified_generic.rs @@ -0,0 +1,178 @@ +use std::collections::HashMap; + +use lance_context_api::{ + AddRowsResponse, ContextError, ContextResult, CreateGenericStoreRequest, GenericStoreApi, + SchemaSpec, +}; +use lance_context_core::{GenericStore as LocalStore, GenericStoreOptions}; +use serde_json::{Map, Value}; + +#[cfg(feature = "remote")] +use lance_context_client::RemoteGenericStore; + +/// A store over a user-declared schema, either an in-process Lance dataset +/// (`Local`) or a handle to a remote server (`Remote`). Mirrors +/// [`crate::RolloutStore`] but for schemas the caller defines. +pub enum GenericStore { + Local(Box), + #[cfg(feature = "remote")] + Remote(RemoteGenericStore), +} + +impl GenericStore { + /// Open an embedded store at `uri`, creating it with `schema` if absent. + pub async fn open(uri: &str, schema: SchemaSpec) -> Result { + Self::open_with_options(uri, schema, None, false).await + } + + pub async fn open_with_options( + uri: &str, + schema: SchemaSpec, + storage_options: Option>, + seal_on_add: bool, + ) -> Result { + let options = GenericStoreOptions { + storage_options, + // Embedded single-process use writes to the fallback shard; a + // multi-writer embedded deployment threads a per-writer id through + // the core `GenericStore` directly. + shard_id: None, + merge_after_generations: None, + session: None, + seal_on_add, + }; + let store = LocalStore::open(uri, schema, options) + .await + .map_err(|e| ContextError::Internal(e.to_string()))?; + Ok(Self::Local(Box::new(store))) + } + + /// Open an existing embedded store, reading its schema from the dataset. + pub async fn open_existing( + uri: &str, + storage_options: Option>, + seal_on_add: bool, + ) -> Result { + let options = GenericStoreOptions { + storage_options, + shard_id: None, + merge_after_generations: None, + session: None, + seal_on_add, + }; + let store = LocalStore::open_existing(uri, options) + .await + .map_err(|e| ContextError::Internal(e.to_string()))?; + Ok(Self::Local(Box::new(store))) + } + + #[cfg(feature = "remote")] + pub async fn connect(base_url: &str, store_name: &str) -> Result { + let store = RemoteGenericStore::connect(base_url, store_name) + .await + .map_err(|e| ContextError::Internal(e.to_string()))?; + Ok(Self::Remote(store)) + } + + #[cfg(feature = "remote")] + pub async fn connect_or_create( + base_url: &str, + req: &CreateGenericStoreRequest, + ) -> Result { + let store = RemoteGenericStore::connect_or_create(base_url, req) + .await + .map_err(|e| ContextError::Internal(e.to_string()))?; + Ok(Self::Remote(store)) + } + + /// Merge pending WAL generations into the base table. Local stores only — + /// a remote store's server owns its own merge schedule. + pub async fn cleanup_wal(&mut self) -> Result { + match self { + Self::Local(store) => store + .cleanup_wal() + .await + .map_err(|e| ContextError::Internal(e.to_string())), + #[cfg(feature = "remote")] + Self::Remote(_) => Err(ContextError::InvalidRequest( + "cleanup_wal is not available on a remote store; the server merges its own shard" + .to_string(), + )), + } + } +} + +impl GenericStoreApi for GenericStore { + fn spec(&self) -> &SchemaSpec { + match self { + Self::Local(store) => GenericStoreApi::spec(store.as_ref()), + #[cfg(feature = "remote")] + Self::Remote(store) => GenericStoreApi::spec(store), + } + } + + async fn add(&self, rows: &[Map]) -> ContextResult { + match self { + Self::Local(store) => GenericStoreApi::add(store.as_ref(), rows).await, + #[cfg(feature = "remote")] + Self::Remote(store) => GenericStoreApi::add(store, rows).await, + } + } + + async fn list( + &self, + limit: Option, + offset: Option, + ) -> ContextResult>> { + match self { + Self::Local(store) => GenericStoreApi::list(store.as_ref(), limit, offset).await, + #[cfg(feature = "remote")] + Self::Remote(store) => GenericStoreApi::list(store, limit, offset).await, + } + } + + async fn list_filtered( + &self, + filter: &str, + limit: Option, + offset: Option, + ) -> ContextResult>> { + match self { + Self::Local(store) => { + GenericStoreApi::list_filtered(store.as_ref(), filter, limit, offset).await + } + #[cfg(feature = "remote")] + Self::Remote(store) => { + GenericStoreApi::list_filtered(store, filter, limit, offset).await + } + } + } + + async fn get( + &self, + id: &str, + columns: Option<&[String]>, + ) -> ContextResult>> { + match self { + Self::Local(store) => GenericStoreApi::get(store.as_ref(), id, columns).await, + #[cfg(feature = "remote")] + Self::Remote(store) => GenericStoreApi::get(store, id, columns).await, + } + } + + async fn flush(&self) -> ContextResult<()> { + match self { + Self::Local(store) => GenericStoreApi::flush(store.as_ref()).await, + #[cfg(feature = "remote")] + Self::Remote(store) => GenericStoreApi::flush(store).await, + } + } + + fn version(&self) -> u64 { + match self { + Self::Local(store) => GenericStoreApi::version(store.as_ref()), + #[cfg(feature = "remote")] + Self::Remote(store) => GenericStoreApi::version(store), + } + } +} diff --git a/python/python/lance_context/__init__.py b/python/python/lance_context/__init__.py index 7072b2b..0929286 100644 --- a/python/python/lance_context/__init__.py +++ b/python/python/lance_context/__init__.py @@ -7,6 +7,7 @@ ContextNamespace, DatagenStore, EmbeddingProvider, + GenericStore, RemoteContext, RolloutStore, __version__, @@ -23,6 +24,7 @@ "Context", "ContextNamespace", "DatagenStore", + "GenericStore", "EmbeddingProvider", "MultiModalEmbeddingProvider", "RemoteContext", diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index 90f3b75..5320625 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -18,6 +18,9 @@ from ._internal import ( # pyright: ignore[reportMissingImports] DatagenStore as _DatagenStore, ) +from ._internal import ( # pyright: ignore[reportMissingImports] + GenericStore as _GenericStore, +) from ._internal import ( # pyright: ignore[reportMissingImports] RemoteContext as _RemoteContext, ) @@ -38,6 +41,7 @@ "AsyncRolloutStore", "Context", "ContextNamespace", + "GenericStore", "DatagenStore", "EmbeddingProvider", "RemoteContext", @@ -2841,3 +2845,145 @@ def get_blob(self, event_id: str) -> bytes | None: def __repr__(self) -> str: return f"DatagenStore(version={self._sync.version()})" + + +class GenericStore: + """Synchronous store over a schema you declare. + + Unlike the fixed-schema stores, the columns come from you: pass a ``schema`` + mapping column names to types at creation, then read and write plain dicts. + An ``id`` column is required — it is the primary key and the key rows are + deduplicated by. + + Types may be written in full form or shorthand:: + + store = GenericStore.open(uri, schema={ + "id": {"type": "string", "nullable": False}, + "user": "string", + "score": "float32", + "tags": {"type": "list", "item": {"type": "string"}}, + "vec": {"type": "vector", "dim": 768, "metric": "cosine"}, + "video": {"type": "binary", "blob": True}, + }) + + ``blob: True`` does not change where the bytes live — they are stored inline + like any other column — it excludes the column from bulk reads. :meth:`list` + never materializes a blob column; fetch one per row with + ``get(id, columns=["video"])``. That is what keeps a listing cheap on a + store holding large payloads. + """ + + def __init__(self, sync_store: _GenericStore) -> None: + self._sync = sync_store + + @classmethod + def open( + cls, + uri: str, + schema: Mapping[str, Any], + *, + storage_options: Mapping[str, str] | None = None, + seal_on_add: bool = False, + ) -> "GenericStore": + """Open (or create) an embedded store at ``uri`` with ``schema``. + + ``seal_on_add=True`` makes rows readable as soon as :meth:`add` returns, + at the cost of serializing concurrent appends. The default defers the + seal — durable on return, readable after :meth:`flush`. + """ + opts = dict(storage_options) if storage_options else None + return cls(_GenericStore.open(uri, dict(schema), opts, seal_on_add)) + + @classmethod + def open_existing( + cls, + uri: str, + *, + storage_options: Mapping[str, str] | None = None, + seal_on_add: bool = False, + ) -> "GenericStore": + """Open an existing embedded store, reading its schema from the dataset.""" + opts = dict(storage_options) if storage_options else None + return cls(_GenericStore.open_existing(uri, opts, seal_on_add)) + + @classmethod + def connect(cls, base_url: str, name: str) -> "GenericStore": + """Connect to an existing store on a remote server.""" + return cls(_GenericStore.connect(base_url, name)) + + @classmethod + def connect_or_create( + cls, + base_url: str, + name: str, + schema: Mapping[str, Any], + *, + storage_options: Mapping[str, str] | None = None, + seal_on_add: bool = False, + ) -> "GenericStore": + """Connect to a remote store, creating it with ``schema`` if absent.""" + opts = dict(storage_options) if storage_options else None + return cls( + _GenericStore.connect_or_create( + base_url, name, dict(schema), opts, seal_on_add + ) + ) + + def schema(self) -> dict[str, Any]: + """The store's schema, as declared at creation.""" + return self._sync.schema() + + def version(self) -> int: + """Base dataset version.""" + return self._sync.version() + + def add(self, rows: Iterable[Mapping[str, Any]]) -> int: + """Append rows. + + Omitted nullable columns are stored as null; a column the schema does + not declare raises, rather than being silently dropped. + """ + return self._sync.add(list(rows)) + + def list( + self, + *, + limit: int | None = None, + offset: int | None = None, + filter: str | None = None, + ) -> list[dict[str, Any]]: + """Read rows, newest write wins per ``id``. Blob columns are excluded. + + ``filter`` is a SQL predicate over your own columns, e.g. + ``"score > 0.5 AND user = 'u1'"``. + """ + return self._sync.list(limit, offset, filter) + + def get( + self, id: str, *, columns: Sequence[str] | None = None + ) -> dict[str, Any] | None: + """Fetch one row by ``id``, or ``None``. + + ``columns`` selects what to read; omit it to read everything except blob + columns. Name a blob column explicitly to fetch its bytes, which come + back base64-encoded. + """ + return self._sync.get(id, list(columns) if columns is not None else None) + + def flush(self) -> None: + """Seal buffered rows so they become readable. + + Unnecessary when the store was opened with ``seal_on_add=True``. + """ + self._sync.flush() + + def cleanup_wal(self) -> int: + """Merge pending WAL generations into the base table (embedded only). + + Returns how many were reclaimed. Without this, generations accumulate + and every read unions all of them. + """ + return self._sync.cleanup_wal() + + def __repr__(self) -> str: + return f"GenericStore(version={self._sync.version()})" diff --git a/python/src/lib.rs b/python/src/lib.rs index 79fff0a..fe3eddc 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -4,18 +4,19 @@ use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::Arc; use chrono::{DateTime, SecondsFormat, Utc}; -use pyo3::exceptions::{PyRuntimeError, PyTypeError}; +use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyBytes, PyDict, PyList, PyModule, PyType}; use pyo3::IntoPyObject; -use serde_json::Value; +use serde_json::{Map, Value}; use tokio::runtime::Runtime; use lance_context::{ - AddRolloutRequest, CreateDatagenStoreRequest, CreateRolloutStoreRequest, DatagenEventDto, - DatagenFailureDto, DatagenFieldStateDto, DatagenStepCursorDto, - DatagenStore as UnifiedDatagenStore, DatagenStoreApi, DatagenValueDto, FoldedDatagenItemDto, - RolloutRecordDto, RolloutStore as UnifiedRolloutStore, RolloutStoreApi, + AddRolloutRequest, ColumnSpec, CreateDatagenStoreRequest, CreateGenericStoreRequest, + CreateRolloutStoreRequest, DatagenEventDto, DatagenFailureDto, DatagenFieldStateDto, + DatagenStepCursorDto, DatagenStore as UnifiedDatagenStore, DatagenStoreApi, DatagenValueDto, + FoldedDatagenItemDto, GenericStore as UnifiedGenericStore, GenericStoreApi, RolloutRecordDto, + RolloutStore as UnifiedRolloutStore, RolloutStoreApi, SchemaSpec, }; use lance_context_api::{ AddRecordRequest, CompactRequest, CompactResponse, CompactStatsResponse, ContextStoreApi, @@ -2954,6 +2955,251 @@ fn failure_to_py(py: Python<'_>, failure: &DatagenFailureDto) -> PyResult, +} + +impl GenericStore { + fn from_store(store: UnifiedGenericStore, runtime: Arc) -> Self { + Self { store, runtime } + } +} + +/// Convert a Python dict describing a schema into a [`SchemaSpec`]. +/// +/// Accepts the same JSON shape the REST API takes, so the two surfaces cannot +/// drift: `{"id": {"type": "string", "nullable": false}, ...}`, or the +/// shorthand `{"id": "string"}` for a nullable column of that type. +fn schema_spec_from_py(schema: &Bound<'_, PyDict>) -> PyResult { + let mut columns = Vec::with_capacity(schema.len()); + for (key, value) in schema.iter() { + let name: String = key.extract()?; + // Shorthand: a bare type name means a nullable column of that type. + let spec_json = if let Ok(type_name) = value.extract::() { + serde_json::json!({ "type": type_name }) + } else { + let encoded = py_dict_to_json(&value) + .map_err(|e| PyValueError::new_err(format!("column '{name}': {e}")))?; + serde_json::from_str(&encoded) + .map_err(|e| PyValueError::new_err(format!("column '{name}': invalid spec: {e}")))? + }; + let column: ColumnSpec = serde_json::from_value(spec_json) + .map_err(|e| PyValueError::new_err(format!("column '{name}': invalid spec: {e}")))?; + columns.push((name, column)); + } + let spec = SchemaSpec::new(columns); + spec.validate().map_err(PyValueError::new_err)?; + Ok(spec) +} + +/// Serialize an arbitrary Python object to JSON via the `json` module, so +/// nested dicts describing a column type round-trip without a hand-written +/// converter. +fn py_dict_to_json(value: &Bound<'_, PyAny>) -> PyResult { + let json = value.py().import("json")?; + json.call_method1("dumps", (value,))?.extract() +} + +/// Convert a Python mapping into one row. +fn row_from_py(row: &Bound<'_, PyAny>) -> PyResult> { + let encoded = py_dict_to_json(row)?; + let value: Value = serde_json::from_str(&encoded).map_err(to_py_err)?; + match value { + Value::Object(map) => Ok(map), + other => Err(PyValueError::new_err(format!( + "each row must be a mapping, got {other}" + ))), + } +} + +/// Convert a decoded row back into a Python dict. +fn row_to_py(py: Python<'_>, row: &Map) -> PyResult { + let json = py.import("json")?; + let encoded = serde_json::to_string(row).map_err(to_py_err)?; + Ok(json.call_method1("loads", (encoded,))?.unbind()) +} + +#[pymethods] +impl GenericStore { + /// Open (or create) an embedded store at `uri` with the given schema. + /// + /// `schema` maps column name to type, e.g. + /// `{"id": {"type": "string", "nullable": False}, "score": "float32"}`. + /// An `id` column is required. + #[classmethod] + #[pyo3(signature = (uri, schema, storage_options = None, seal_on_add = false))] + fn open( + _cls: &Bound<'_, PyType>, + py: Python<'_>, + uri: &str, + schema: &Bound<'_, PyDict>, + storage_options: Option>, + seal_on_add: bool, + ) -> PyResult { + let spec = schema_spec_from_py(schema)?; + let runtime = Arc::new(Runtime::new().map_err(to_py_err)?); + let store_res = py.allow_threads(|| { + runtime.block_on(UnifiedGenericStore::open_with_options( + uri, + spec, + storage_options, + seal_on_add, + )) + }); + Ok(Self::from_store(store_res.map_err(to_py_err)?, runtime)) + } + + /// Open an existing embedded store, reading its schema from the dataset. + #[classmethod] + #[pyo3(signature = (uri, storage_options = None, seal_on_add = false))] + fn open_existing( + _cls: &Bound<'_, PyType>, + py: Python<'_>, + uri: &str, + storage_options: Option>, + seal_on_add: bool, + ) -> PyResult { + let runtime = Arc::new(Runtime::new().map_err(to_py_err)?); + let store_res = py.allow_threads(|| { + runtime.block_on(UnifiedGenericStore::open_existing( + uri, + storage_options, + seal_on_add, + )) + }); + Ok(Self::from_store(store_res.map_err(to_py_err)?, runtime)) + } + + /// Connect to an existing store on a remote server. + #[classmethod] + fn connect( + _cls: &Bound<'_, PyType>, + py: Python<'_>, + base_url: &str, + name: &str, + ) -> PyResult { + let runtime = Arc::new(Runtime::new().map_err(to_py_err)?); + let store_res = + py.allow_threads(|| runtime.block_on(UnifiedGenericStore::connect(base_url, name))); + Ok(Self::from_store(store_res.map_err(to_py_err)?, runtime)) + } + + /// Connect to a remote store, creating it with `schema` if absent. + #[classmethod] + #[pyo3(signature = (base_url, name, schema, storage_options = None, seal_on_add = false))] + fn connect_or_create( + _cls: &Bound<'_, PyType>, + py: Python<'_>, + base_url: &str, + name: &str, + schema: &Bound<'_, PyDict>, + storage_options: Option>, + seal_on_add: bool, + ) -> PyResult { + let spec = schema_spec_from_py(schema)?; + let req = CreateGenericStoreRequest { + name: name.to_string(), + schema: spec, + storage_options, + seal_on_add, + }; + let runtime = Arc::new(Runtime::new().map_err(to_py_err)?); + let store_res = py.allow_threads(|| { + runtime.block_on(UnifiedGenericStore::connect_or_create(base_url, &req)) + }); + Ok(Self::from_store(store_res.map_err(to_py_err)?, runtime)) + } + + /// The store's schema, as a dict keyed by column name. + fn schema(&self, py: Python<'_>) -> PyResult { + let spec = GenericStoreApi::spec(&self.store); + let json = py.import("json")?; + let encoded = serde_json::to_string(spec).map_err(to_py_err)?; + Ok(json.call_method1("loads", (encoded,))?.unbind()) + } + + /// Base dataset version. + fn version(&self) -> u64 { + GenericStoreApi::version(&self.store) + } + + /// Append rows. Each row is a dict keyed by column name; omitted nullable + /// columns are written as null, and an undeclared column is an error. + fn add(&self, py: Python<'_>, rows: &Bound<'_, PyAny>) -> PyResult { + let mut parsed = Vec::new(); + for row in rows.try_iter()? { + parsed.push(row_from_py(&row?)?); + } + let res = py.allow_threads(|| { + self.runtime + .block_on(GenericStoreApi::add(&self.store, &parsed)) + }); + Ok(res.map_err(to_py_err)?.version) + } + + /// Read rows. Blob columns are excluded; use `get` with `columns` to fetch + /// them one row at a time. + #[pyo3(signature = (limit = None, offset = None, filter = None))] + fn list( + &self, + py: Python<'_>, + limit: Option, + offset: Option, + filter: Option, + ) -> PyResult> { + let rows = py.allow_threads(|| { + self.runtime.block_on(async { + match &filter { + Some(filter) => { + GenericStoreApi::list_filtered(&self.store, filter, limit, offset).await + } + None => GenericStoreApi::list(&self.store, limit, offset).await, + } + }) + }); + rows.map_err(to_py_err)? + .iter() + .map(|row| row_to_py(py, row)) + .collect() + } + + /// Fetch one row by id, or `None`. `columns` selects what to read; omit it + /// to read everything except blob columns. + #[pyo3(signature = (id, columns = None))] + fn get( + &self, + py: Python<'_>, + id: &str, + columns: Option>, + ) -> PyResult> { + let row = py.allow_threads(|| { + self.runtime + .block_on(GenericStoreApi::get(&self.store, id, columns.as_deref())) + }); + match row.map_err(to_py_err)? { + Some(row) => Ok(Some(row_to_py(py, &row)?)), + None => Ok(None), + } + } + + /// Seal the active memtable so added rows become readable. Unnecessary when + /// the store was opened with `seal_on_add=True`. + fn flush(&self, py: Python<'_>) -> PyResult<()> { + py.allow_threads(|| self.runtime.block_on(GenericStoreApi::flush(&self.store))) + .map_err(to_py_err) + } + + /// Merge pending WAL generations into the base table (embedded stores only). + fn cleanup_wal(&mut self, py: Python<'_>) -> PyResult { + py.allow_threads(|| self.runtime.block_on(self.store.cleanup_wal())) + .map_err(to_py_err) + } +} + #[pymodule] fn _internal(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(version, m)?)?; @@ -2964,5 +3210,6 @@ fn _internal(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/python/tests/test_generic_store.py b/python/tests/test_generic_store.py new file mode 100644 index 0000000..ef92d8a --- /dev/null +++ b/python/tests/test_generic_store.py @@ -0,0 +1,192 @@ +"""Stores over a user-declared schema.""" + +from __future__ import annotations + +import base64 +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from pathlib import Path +from lance_context import GenericStore + +SCHEMA = { + "id": {"type": "string", "nullable": False}, + "user": "string", + "score": "float32", + "tags": {"type": "list", "item": {"type": "string"}}, + "embedding": {"type": "vector", "dim": 3, "metric": "cosine"}, + "payload": {"type": "binary", "blob": True}, +} + + +def open_store(tmp_path: Path, name: str = "s", **kwargs) -> GenericStore: + # seal_on_add so each test reads its own writes without an explicit flush; + # the deferred default is exercised separately below. + kwargs.setdefault("seal_on_add", True) + return GenericStore.open(str(tmp_path / f"{name}.lance"), schema=SCHEMA, **kwargs) + + +def test_declared_schema_round_trips(tmp_path: Path) -> None: + store = open_store(tmp_path) + store.add( + [ + { + "id": "r1", + "user": "u1", + "score": 0.5, + "tags": ["a", "b"], + "embedding": [1.0, 2.0, 3.0], + }, + # Every nullable column may be omitted. + {"id": "r2"}, + ] + ) + + rows = store.list() + assert len(rows) == 2 + first = store.get("r1") + assert first is not None + assert first["user"] == "u1" + assert first["tags"] == ["a", "b"] + assert first["embedding"] == [1.0, 2.0, 3.0] + + second = store.get("r2") + assert second is not None + # Nulls are omitted rather than returned as None. + assert set(second) == {"id"} + + assert store.get("missing") is None + + +def test_schema_reports_declared_columns(tmp_path: Path) -> None: + store = open_store(tmp_path) + names = [name for name, _ in store.schema()["columns"]] + assert names == ["id", "user", "score", "tags", "embedding", "payload"] + + +def test_id_column_is_required(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="id"): + GenericStore.open(str(tmp_path / "bad.lance"), schema={"name": "string"}) + + +def test_undeclared_column_is_rejected_not_dropped(tmp_path: Path) -> None: + # A field-name typo must fail the write rather than becoming missing data. + store = open_store(tmp_path) + with pytest.raises(Exception, match="not declared"): + store.add([{"id": "r1", "usr": "typo"}]) + + +def test_type_mismatch_is_rejected(tmp_path: Path) -> None: + store = open_store(tmp_path) + with pytest.raises(Exception): + store.add([{"id": "r1", "score": "not-a-number"}]) + + +def test_vector_length_must_match(tmp_path: Path) -> None: + store = open_store(tmp_path) + with pytest.raises(Exception, match="vector of 3"): + store.add([{"id": "r1", "embedding": [1.0, 2.0]}]) + + +def test_blob_excluded_from_list_but_fetchable_per_row(tmp_path: Path) -> None: + # The cost model for large payloads: bulk reads skip them, point reads ask. + store = open_store(tmp_path) + store.add([{"id": "r1", "payload": [1, 2, 3]}]) + + listed = store.list() + assert "payload" not in listed[0] + + fetched = store.get("r1", columns=["payload"]) + assert fetched is not None + assert base64.b64decode(fetched["payload"]) == bytes([1, 2, 3]) + + +def test_large_blob_round_trips(tmp_path: Path) -> None: + # Multi-megabyte inline payloads are the point of these stores. + store = open_store(tmp_path) + payload = bytes((i % 251) for i in range(2 * 1024 * 1024)) + store.add([{"id": "big", "payload": list(payload)}]) + + fetched = store.get("big", columns=["payload"]) + assert fetched is not None + assert base64.b64decode(fetched["payload"]) == payload + + +def test_filter_and_paging(tmp_path: Path) -> None: + store = open_store(tmp_path) + store.add( + [ + {"id": f"r{i}", "user": "even" if i % 2 == 0 else "odd", "score": float(i)} + for i in range(6) + ] + ) + + evens = store.list(filter="user = 'even'") + assert len(evens) == 3 + assert all(row["user"] == "even" for row in evens) + + page = store.list(limit=2, offset=1) + assert len(page) == 2 + + +def test_duplicate_id_keeps_the_newest_write(tmp_path: Path) -> None: + # `id` is the merge key, so re-adding one supersedes the older row. + store = open_store(tmp_path) + store.add([{"id": "r1", "user": "first"}]) + store.add([{"id": "r1", "user": "second"}]) + + rows = store.list() + assert len(rows) == 1 + assert rows[0]["user"] == "second" + + +def test_deferred_seal_needs_a_flush(tmp_path: Path) -> None: + store = GenericStore.open( + str(tmp_path / "deferred.lance"), schema=SCHEMA, seal_on_add=False + ) + store.add([{"id": "r1"}]) + assert store.list() == [] + + store.flush() + assert len(store.list()) == 1 + + +def test_wal_generations_merge_into_the_base_table(tmp_path: Path) -> None: + store = open_store(tmp_path) + for i in range(3): + store.add([{"id": f"r{i}"}]) + + assert store.cleanup_wal() > 0 + assert len(store.list()) == 3 + + +def test_reopen_reads_the_persisted_schema(tmp_path: Path) -> None: + uri = str(tmp_path / "persist.lance") + store = GenericStore.open(uri, schema=SCHEMA, seal_on_add=True) + store.add([{"id": "r1", "user": "u1"}]) + del store + + reopened = GenericStore.open_existing(uri, seal_on_add=True) + assert [name for name, _ in reopened.schema()["columns"]] == list(SCHEMA) + assert len(reopened.list()) == 1 + + +def test_reopening_with_a_conflicting_schema_is_rejected(tmp_path: Path) -> None: + # Reinterpreting existing data under a new schema is the failure to prevent. + uri = str(tmp_path / "conflict.lance") + GenericStore.open(uri, schema=SCHEMA, seal_on_add=True) + + with pytest.raises(Exception, match="different schema"): + GenericStore.open(uri, schema={**SCHEMA, "extra": "int64"}, seal_on_add=True) + + +def test_shorthand_and_full_type_forms_agree(tmp_path: Path) -> None: + shorthand = GenericStore.open( + str(tmp_path / "short.lance"), + schema={"id": {"type": "string", "nullable": False}, "n": "int64"}, + seal_on_add=True, + ) + shorthand.add([{"id": "a", "n": 7}]) + assert shorthand.list()[0]["n"] == 7