diff --git a/Cargo.lock b/Cargo.lock index b554db2..9cd03fc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5483,6 +5483,7 @@ dependencies = [ "arrow-ipc 58.3.0", "arrow-json 58.3.0", "arrow-schema 58.3.0", + "base64", "chrono", "datafusion 53.1.0", "futures", diff --git a/crates/lance-context-core/Cargo.toml b/crates/lance-context-core/Cargo.toml index a49e5a7..5789237 100644 --- a/crates/lance-context-core/Cargo.toml +++ b/crates/lance-context-core/Cargo.toml @@ -18,6 +18,7 @@ default = ["metrics"] metrics = ["dep:metrics"] [dependencies] +base64 = "0.22" arrow-array = "58" arrow-ipc = "58" arrow-json = "58" diff --git a/crates/lance-context-core/src/generic_codec.rs b/crates/lance-context-core/src/generic_codec.rs new file mode 100644 index 0000000..dcbbc0d --- /dev/null +++ b/crates/lance-context-core/src/generic_codec.rs @@ -0,0 +1,807 @@ +//! Schema-driven encode/decode: rows as maps, columns from a [`SchemaSpec`]. +//! +//! The built-in stores convert records to Arrow with hand-written builders — 37 +//! of them in the rollout store alone — and decode by naming every column again. +//! That is what makes adding a column a five-to-seven-place edit. Here the +//! schema drives both directions, so an arbitrary user schema needs no code. +//! +//! # Nested values are decoded by name +//! +//! The built-in decoders read struct children *positionally* +//! (`relationships`, `state_metadata`), so reordering fields silently +//! assigns data to the wrong field. Everything here is keyed by name in both +//! directions. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow_array::builder::{ + BooleanBuilder, FixedSizeListBuilder, Float32Builder, Float64Builder, Int32Builder, + Int64Builder, LargeBinaryBuilder, LargeStringBuilder, ListBuilder, StringBuilder, + TimestampMicrosecondBuilder, +}; +use arrow_array::{ + Array, ArrayRef, BooleanArray, FixedSizeListArray, Float32Array, Float64Array, Int32Array, + Int64Array, LargeBinaryArray, LargeStringArray, ListArray, RecordBatch, StringArray, + TimestampMicrosecondArray, +}; +use arrow_schema::{ArrowError, Schema}; +use serde_json::{Map, Number, Value}; + +use crate::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. +pub type Row = Map; + +/// Encode rows into a [`RecordBatch`] matching `schema`. +/// +/// Values are matched to columns **by name**; a key with no matching column is +/// an error rather than being silently dropped, so a typo in a field name +/// surfaces at the write instead of turning into missing data. +/// +/// Binary columns accept either a JSON array of byte values or a base64 string, +/// since JSON has no byte-string type. +/// +/// # Errors +/// +/// - a row is missing a non-nullable column, or supplies null for one +/// - a value's JSON type does not match the declared column type +/// - a vector's length does not match the declared dimension +/// - a row carries a key that the schema does not declare +pub fn rows_to_batch( + spec: &SchemaSpec, + schema: Arc, + rows: &[Row], +) -> Result { + for (index, row) in rows.iter().enumerate() { + for key in row.keys() { + if spec.column(key).is_none() { + return Err(ArrowError::InvalidArgumentError(format!( + "row {index}: column '{key}' is not declared in the store schema" + ))); + } + } + } + + // Build every declared column, then assemble in the dataset's field order. + // Keying by name (rather than pushing positionally) is what lets the + // physical schema reorder or omit columns without breaking the write. + let mut arrays: HashMap<&str, ArrayRef> = HashMap::with_capacity(spec.columns.len()); + for (name, column) in &spec.columns { + arrays.insert(name.as_str(), build_column(name, column, rows)?); + } + + let columns = schema + .fields() + .iter() + .map(|field| { + arrays.get(field.name().as_str()).cloned().ok_or_else(|| { + ArrowError::SchemaError(format!( + "dataset column '{}' is not declared in the store schema", + field.name() + )) + }) + }) + .collect::, _>>()?; + + RecordBatch::try_new(schema, columns) +} + +/// Decode a [`RecordBatch`] back into rows. +/// +/// Only columns present in the batch are decoded, so a projected scan (one that +/// dropped blob columns, say) round-trips to rows without those keys rather +/// than erroring. Null values are omitted from the row entirely. +/// +/// # Errors +/// +/// Returns an error if a column's array type does not match the schema. +pub fn batch_to_rows(spec: &SchemaSpec, batch: &RecordBatch) -> Result, ArrowError> { + let mut rows = vec![Row::new(); batch.num_rows()]; + + for (name, column) in &spec.columns { + let Some(array) = batch.column_by_name(name) else { + // Projected out; not an error. + continue; + }; + decode_column(name, column, array.as_ref(), &mut rows)?; + } + + Ok(rows) +} + +/// Extract the `id` of every row, in order — the merge key the storage layer +/// needs, without decoding the rest of the batch. +/// +/// # Errors +/// +/// Returns an error if the batch has no `id` column or it is not a string. +pub fn ids_from_batch(batch: &RecordBatch) -> Result, ArrowError> { + let array = batch + .column_by_name(ID_COLUMN) + .ok_or_else(|| ArrowError::SchemaError(format!("batch has no '{ID_COLUMN}' column")))?; + let ids = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + ArrowError::SchemaError(format!("'{ID_COLUMN}' column is not a string array")) + })?; + Ok((0..ids.len()) + .map(|row| ids.value(row).to_string()) + .collect()) +} + +// ------------------------------------------------------------------ encoding + +/// Fetch a row's value for `name`, enforcing nullability. +/// +/// `Ok(None)` means "write a null here", which is valid only for a nullable +/// column. An explicit JSON `null` is treated identically to an absent key. +fn value_for<'a>( + name: &str, + column: &ColumnSpec, + row: &'a Row, + index: usize, +) -> Result, ArrowError> { + match row.get(name) { + Some(Value::Null) | None if !column.nullable => Err(ArrowError::InvalidArgumentError( + format!("row {index}: column '{name}' is not nullable and was not supplied"), + )), + Some(Value::Null) | None => Ok(None), + Some(value) => Ok(Some(value)), + } +} + +fn type_error(name: &str, index: usize, expected: &str, got: &Value) -> ArrowError { + ArrowError::InvalidArgumentError(format!( + "row {index}: column '{name}' expects {expected}, got {got}" + )) +} + +/// Error for a bad *element* inside a list, where the row index is not +/// threaded through the append closure. +fn element_error(name: &str, expected: &str, got: &Value) -> ArrowError { + ArrowError::InvalidArgumentError(format!( + "column '{name}': expects {expected}, got element {got}" + )) +} + +fn build_column(name: &str, column: &ColumnSpec, rows: &[Row]) -> Result { + Ok(match &column.column_type { + ColumnType::String { large: false } => { + let mut b = StringBuilder::new(); + for (index, row) in rows.iter().enumerate() { + match value_for(name, column, row, index)? { + Some(Value::String(text)) => b.append_value(text), + Some(other) => return Err(type_error(name, index, "a string", other)), + None => b.append_null(), + } + } + Arc::new(b.finish()) + } + ColumnType::String { large: true } => { + let mut b = LargeStringBuilder::new(); + for (index, row) in rows.iter().enumerate() { + match value_for(name, column, row, index)? { + Some(Value::String(text)) => b.append_value(text), + Some(other) => return Err(type_error(name, index, "a string", other)), + None => b.append_null(), + } + } + Arc::new(b.finish()) + } + ColumnType::Int32 => { + let mut b = Int32Builder::new(); + for (index, row) in rows.iter().enumerate() { + match value_for(name, column, row, index)? { + Some(value) => b.append_value( + as_i64(name, index, value)? + .try_into() + .map_err(|_| type_error(name, index, "a 32-bit integer", value))?, + ), + None => b.append_null(), + } + } + Arc::new(b.finish()) + } + ColumnType::Int64 => { + let mut b = Int64Builder::new(); + for (index, row) in rows.iter().enumerate() { + match value_for(name, column, row, index)? { + Some(value) => b.append_value(as_i64(name, index, value)?), + None => b.append_null(), + } + } + Arc::new(b.finish()) + } + ColumnType::Float32 => { + let mut b = Float32Builder::new(); + for (index, row) in rows.iter().enumerate() { + match value_for(name, column, row, index)? { + Some(value) => b.append_value(as_f64(name, index, value)? as f32), + None => b.append_null(), + } + } + Arc::new(b.finish()) + } + ColumnType::Float64 => { + let mut b = Float64Builder::new(); + for (index, row) in rows.iter().enumerate() { + match value_for(name, column, row, index)? { + Some(value) => b.append_value(as_f64(name, index, value)?), + None => b.append_null(), + } + } + Arc::new(b.finish()) + } + ColumnType::Bool => { + let mut b = BooleanBuilder::new(); + for (index, row) in rows.iter().enumerate() { + match value_for(name, column, row, index)? { + Some(Value::Bool(flag)) => b.append_value(*flag), + Some(other) => return Err(type_error(name, index, "a boolean", other)), + None => b.append_null(), + } + } + Arc::new(b.finish()) + } + ColumnType::Timestamp => { + let mut b = TimestampMicrosecondBuilder::new(); + for (index, row) in rows.iter().enumerate() { + match value_for(name, column, row, index)? { + Some(value) => b.append_value(as_timestamp_micros(name, index, value)?), + None => b.append_null(), + } + } + Arc::new(b.finish()) + } + ColumnType::Binary { .. } => { + let mut b = LargeBinaryBuilder::new(); + for (index, row) in rows.iter().enumerate() { + match value_for(name, column, row, index)? { + Some(value) => b.append_value(as_bytes(name, index, value)?), + None => b.append_null(), + } + } + Arc::new(b.finish()) + } + ColumnType::Vector { dim, .. } => { + let mut b = FixedSizeListBuilder::new(Float32Builder::new(), *dim); + for (index, row) in rows.iter().enumerate() { + match value_for(name, column, row, index)? { + Some(Value::Array(items)) => { + if items.len() != *dim as usize { + return Err(ArrowError::InvalidArgumentError(format!( + "row {index}: column '{name}' expects a vector of {dim} \ + values, got {}", + items.len() + ))); + } + for item in items { + b.values().append_value(as_f64(name, index, item)? as f32); + } + b.append(true); + } + Some(other) => { + return Err(type_error(name, index, "an array of numbers", other)) + } + None => { + // A FixedSizeList null still needs its slots filled. + for _ in 0..*dim { + b.values().append_null(); + } + b.append(false); + } + } + } + Arc::new(b.finish()) + } + ColumnType::List { item } => build_list_column(name, column, item, rows)?, + }) +} + +/// Build a `List` column. Element types are the scalar subset; +/// [`SchemaSpec::validate`] has already rejected nested lists and vector lists. +fn build_list_column( + name: &str, + column: &ColumnSpec, + item: &ColumnType, + rows: &[Row], +) -> Result { + /// Drive a `ListBuilder`, appending each element through `$append`. + macro_rules! build_list { + ($builder:expr, $append:expr) => {{ + let mut b = ListBuilder::new($builder); + for (index, row) in rows.iter().enumerate() { + match value_for(name, column, row, index)? { + Some(Value::Array(items)) => { + for element in items { + let append: &dyn Fn(&mut _, &Value) -> Result<(), ArrowError> = + &$append; + append(b.values(), element)?; + } + b.append(true); + } + Some(other) => return Err(type_error(name, index, "an array", other)), + None => b.append(false), + } + } + Arc::new(b.finish()) as ArrayRef + }}; + } + + Ok(match item { + ColumnType::String { large: false } => build_list!( + StringBuilder::new(), + |values: &mut StringBuilder, element: &Value| match element { + Value::String(text) => { + values.append_value(text); + Ok(()) + } + other => Err(element_error(name, "an array of strings", other)), + } + ), + ColumnType::String { large: true } => build_list!( + LargeStringBuilder::new(), + |values: &mut LargeStringBuilder, element: &Value| match element { + Value::String(text) => { + values.append_value(text); + Ok(()) + } + other => Err(element_error(name, "an array of strings", other)), + } + ), + ColumnType::Int32 => build_list!( + Int32Builder::new(), + |values: &mut Int32Builder, element: &Value| { + let raw = as_i64(name, 0, element)?; + values + .append_value(raw.try_into().map_err(|_| { + element_error(name, "an array of 32-bit integers", element) + })?); + Ok(()) + } + ), + ColumnType::Int64 => build_list!( + Int64Builder::new(), + |values: &mut Int64Builder, element: &Value| { + values.append_value(as_i64(name, 0, element)?); + Ok(()) + } + ), + ColumnType::Float32 => build_list!( + Float32Builder::new(), + |values: &mut Float32Builder, element: &Value| { + values.append_value(as_f64(name, 0, element)? as f32); + Ok(()) + } + ), + ColumnType::Float64 => build_list!( + Float64Builder::new(), + |values: &mut Float64Builder, element: &Value| { + values.append_value(as_f64(name, 0, element)?); + Ok(()) + } + ), + ColumnType::Bool => build_list!( + BooleanBuilder::new(), + |values: &mut BooleanBuilder, element: &Value| match element { + Value::Bool(flag) => { + values.append_value(*flag); + Ok(()) + } + other => Err(element_error(name, "an array of booleans", other)), + } + ), + ColumnType::Timestamp => build_list!( + TimestampMicrosecondBuilder::new(), + |values: &mut TimestampMicrosecondBuilder, element: &Value| { + values.append_value(as_timestamp_micros(name, 0, element)?); + Ok(()) + } + ), + ColumnType::Binary { .. } => build_list!( + LargeBinaryBuilder::new(), + |values: &mut LargeBinaryBuilder, element: &Value| { + values.append_value(as_bytes(name, 0, element)?); + Ok(()) + } + ), + // Rejected by `SchemaSpec::validate`, so unreachable in practice. + ColumnType::List { .. } | ColumnType::Vector { .. } => { + return Err(ArrowError::SchemaError(format!( + "column '{name}': unsupported list element type" + ))) + } + }) +} + +fn as_i64(name: &str, index: usize, value: &Value) -> Result { + value + .as_i64() + .ok_or_else(|| type_error(name, index, "an integer", value)) +} + +fn as_f64(name: &str, index: usize, value: &Value) -> Result { + value + .as_f64() + .ok_or_else(|| type_error(name, index, "a number", value)) +} + +/// Timestamps accept microseconds-since-epoch or an RFC 3339 string. +fn as_timestamp_micros(name: &str, index: usize, value: &Value) -> Result { + match value { + Value::Number(number) => number + .as_i64() + .ok_or_else(|| type_error(name, index, "microseconds since epoch", value)), + Value::String(text) => chrono::DateTime::parse_from_rfc3339(text) + .map(|parsed| parsed.timestamp_micros()) + .map_err(|_| type_error(name, index, "an RFC 3339 timestamp", value)), + other => Err(type_error( + name, + index, + "microseconds since epoch or an RFC 3339 string", + other, + )), + } +} + +/// Binary accepts a base64 string or an array of byte values — JSON has no +/// native byte-string type, and both spellings appear in practice. +fn as_bytes(name: &str, index: usize, value: &Value) -> Result, ArrowError> { + use base64::Engine; + match value { + Value::String(encoded) => base64::engine::general_purpose::STANDARD + .decode(encoded) + .map_err(|_| type_error(name, index, "base64-encoded bytes", value)), + Value::Array(items) => items + .iter() + .map(|item| { + let byte = as_i64(name, index, item)?; + u8::try_from(byte) + .map_err(|_| type_error(name, index, "an array of bytes (0-255)", item)) + }) + .collect(), + other => Err(type_error( + name, + index, + "base64-encoded bytes or an array of byte values", + other, + )), + } +} + +// ------------------------------------------------------------------ decoding + +/// Downcast `array`, or report the mismatch against the declared type. +fn downcast<'a, A: 'static>(name: &str, array: &'a dyn Array) -> Result<&'a A, ArrowError> { + array.as_any().downcast_ref::().ok_or_else(|| { + ArrowError::SchemaError(format!( + "column '{name}' has array type {:?}, which does not match the store schema", + array.data_type() + )) + }) +} + +fn decode_column( + name: &str, + column: &ColumnSpec, + array: &dyn Array, + rows: &mut [Row], +) -> Result<(), ArrowError> { + /// Decode every non-null slot through `$convert`. + macro_rules! decode { + ($ty:ty, $convert:expr) => {{ + let typed = downcast::<$ty>(name, array)?; + for (index, row) in rows.iter_mut().enumerate() { + if typed.is_null(index) { + continue; + } + #[allow(clippy::redundant_closure_call)] + row.insert(name.to_string(), ($convert)(typed.value(index))); + } + }}; + } + + match &column.column_type { + ColumnType::String { large: false } => { + decode!(StringArray, |value: &str| Value::String(value.to_string())) + } + ColumnType::String { large: true } => { + decode!(LargeStringArray, |value: &str| Value::String( + value.to_string() + )) + } + ColumnType::Int32 => decode!(Int32Array, |value: i32| Value::Number(value.into())), + ColumnType::Int64 => decode!(Int64Array, |value: i64| Value::Number(value.into())), + ColumnType::Float32 => decode!(Float32Array, |value: f32| number(f64::from(value))), + ColumnType::Float64 => decode!(Float64Array, number), + ColumnType::Bool => decode!(BooleanArray, Value::Bool), + ColumnType::Timestamp => { + decode!(TimestampMicrosecondArray, |value: i64| Value::Number( + value.into() + )) + } + ColumnType::Binary { .. } => { + // Base64 on the way out, matching the preferred input spelling. + use base64::Engine; + decode!(LargeBinaryArray, |value: &[u8]| Value::String( + base64::engine::general_purpose::STANDARD.encode(value) + )) + } + ColumnType::Vector { .. } => { + let typed = downcast::(name, array)?; + for (index, row) in rows.iter_mut().enumerate() { + if typed.is_null(index) { + continue; + } + let values = typed.value(index); + let floats = downcast::(name, values.as_ref())?; + let vector = (0..floats.len()) + .map(|slot| number(f64::from(floats.value(slot)))) + .collect(); + row.insert(name.to_string(), Value::Array(vector)); + } + } + ColumnType::List { item } => { + let typed = downcast::(name, array)?; + for (index, row) in rows.iter_mut().enumerate() { + if typed.is_null(index) { + continue; + } + let values = typed.value(index); + row.insert( + name.to_string(), + Value::Array(decode_list_values(name, item, values.as_ref())?), + ); + } + } + } + Ok(()) +} + +fn decode_list_values( + name: &str, + item: &ColumnType, + values: &dyn Array, +) -> Result, ArrowError> { + macro_rules! collect { + ($ty:ty, $convert:expr) => {{ + let typed = downcast::<$ty>(name, values)?; + (0..typed.len()) + .map(|slot| { + if typed.is_null(slot) { + Value::Null + } else { + #[allow(clippy::redundant_closure_call)] + ($convert)(typed.value(slot)) + } + }) + .collect() + }}; + } + + Ok(match item { + ColumnType::String { large: false } => { + collect!(StringArray, |v: &str| Value::String(v.to_string())) + } + ColumnType::String { large: true } => { + collect!(LargeStringArray, |v: &str| Value::String(v.to_string())) + } + ColumnType::Int32 => collect!(Int32Array, |v: i32| Value::Number(v.into())), + ColumnType::Int64 => collect!(Int64Array, |v: i64| Value::Number(v.into())), + ColumnType::Float32 => collect!(Float32Array, |v: f32| number(f64::from(v))), + ColumnType::Float64 => collect!(Float64Array, number), + ColumnType::Bool => collect!(BooleanArray, Value::Bool), + ColumnType::Timestamp => { + collect!(TimestampMicrosecondArray, |v: i64| Value::Number(v.into())) + } + ColumnType::Binary { .. } => { + use base64::Engine; + collect!(LargeBinaryArray, |v: &[u8]| Value::String( + base64::engine::general_purpose::STANDARD.encode(v) + )) + } + ColumnType::List { .. } | ColumnType::Vector { .. } => { + return Err(ArrowError::SchemaError(format!( + "column '{name}': unsupported list element type" + ))) + } + }) +} + +/// JSON has no NaN or infinity, so non-finite floats decode to null rather than +/// producing invalid JSON. +fn number(value: f64) -> Value { + Number::from_f64(value).map_or(Value::Null, Value::Number) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::schema_spec::ColumnSpec; + use serde_json::json; + + fn spec() -> SchemaSpec { + SchemaSpec::new(vec![ + ( + ID_COLUMN.to_string(), + ColumnSpec::required(ColumnType::String { large: false }), + ), + ("count".to_string(), ColumnSpec::new(ColumnType::Int64)), + ("score".to_string(), ColumnSpec::new(ColumnType::Float32)), + ("active".to_string(), ColumnSpec::new(ColumnType::Bool)), + ( + "tags".to_string(), + ColumnSpec::new(ColumnType::List { + item: Box::new(ColumnType::String { large: false }), + }), + ), + ( + "embedding".to_string(), + ColumnSpec::new(ColumnType::Vector { + dim: 3, + metric: None, + }), + ), + ( + "blob".to_string(), + ColumnSpec::new(ColumnType::Binary { blob: true }), + ), + ("at".to_string(), ColumnSpec::new(ColumnType::Timestamp)), + ]) + } + + fn row(value: Value) -> Row { + value.as_object().unwrap().clone() + } + + #[test] + fn rows_round_trip_through_arrow() { + let spec = spec(); + let schema = Arc::new(spec.to_arrow().unwrap()); + let rows = vec![ + row(json!({ + "id": "r1", + "count": 42, + "score": 0.5, + "active": true, + "tags": ["a", "b"], + "embedding": [1.0, 2.0, 3.0], + "blob": [0, 1, 255], + "at": 1_700_000_000_000_000i64, + })), + // Every nullable column omitted. + row(json!({"id": "r2"})), + ]; + + let batch = rows_to_batch(&spec, schema, &rows).unwrap(); + assert_eq!(batch.num_rows(), 2); + assert_eq!(ids_from_batch(&batch).unwrap(), vec!["r1", "r2"]); + + let decoded = batch_to_rows(&spec, &batch).unwrap(); + assert_eq!(decoded[0]["id"], json!("r1")); + assert_eq!(decoded[0]["count"], json!(42)); + assert_eq!(decoded[0]["active"], json!(true)); + assert_eq!(decoded[0]["tags"], json!(["a", "b"])); + assert_eq!(decoded[0]["embedding"], json!([1.0, 2.0, 3.0])); + assert_eq!(decoded[0]["at"], json!(1_700_000_000_000_000i64)); + // Binary comes back base64-encoded. + assert_eq!(decoded[0]["blob"], json!("AAH/")); + + // Nulls are omitted rather than emitted as JSON null. + assert_eq!(decoded[1].keys().collect::>(), vec!["id"]); + } + + #[test] + fn binary_accepts_base64_and_byte_arrays_alike() { + let spec = spec(); + let schema = Arc::new(spec.to_arrow().unwrap()); + let rows = vec![ + row(json!({"id": "a", "blob": "AAH/"})), + row(json!({"id": "b", "blob": [0, 1, 255]})), + ]; + let batch = rows_to_batch(&spec, schema, &rows).unwrap(); + let decoded = batch_to_rows(&spec, &batch).unwrap(); + assert_eq!(decoded[0]["blob"], decoded[1]["blob"]); + } + + #[test] + fn missing_required_column_is_rejected() { + let spec = spec(); + let schema = Arc::new(spec.to_arrow().unwrap()); + let err = rows_to_batch(&spec, schema, &[row(json!({"count": 1}))]).unwrap_err(); + assert!(err.to_string().contains("'id' is not nullable"), "{err}"); + } + + #[test] + fn undeclared_column_is_rejected_not_dropped() { + // Silently dropping an unknown key turns a field-name typo into + // missing data that surfaces much later. + let spec = spec(); + let schema = Arc::new(spec.to_arrow().unwrap()); + let err = rows_to_batch(&spec, schema, &[row(json!({"id": "r1", "typo": 1}))]).unwrap_err(); + assert!(err.to_string().contains("not declared"), "{err}"); + } + + #[test] + fn type_mismatches_are_rejected() { + let spec = spec(); + let schema = Arc::new(spec.to_arrow().unwrap()); + for bad in [ + json!({"id": "r1", "count": "not-a-number"}), + json!({"id": "r1", "active": "yes"}), + json!({"id": 7}), + json!({"id": "r1", "tags": "not-a-list"}), + ] { + assert!( + rows_to_batch(&spec, schema.clone(), &[row(bad.clone())]).is_err(), + "{bad} should be rejected" + ); + } + } + + #[test] + fn vector_length_must_match_the_declared_dimension() { + let spec = spec(); + let schema = Arc::new(spec.to_arrow().unwrap()); + let err = rows_to_batch( + &spec, + schema, + &[row(json!({"id": "r1", "embedding": [1.0, 2.0]}))], + ) + .unwrap_err(); + assert!(err.to_string().contains("vector of 3 values"), "{err}"); + } + + #[test] + fn timestamps_accept_rfc3339_strings() { + let spec = spec(); + let schema = Arc::new(spec.to_arrow().unwrap()); + let batch = rows_to_batch( + &spec, + schema, + &[row(json!({"id": "r1", "at": "2023-11-14T22:13:20Z"}))], + ) + .unwrap(); + let decoded = batch_to_rows(&spec, &batch).unwrap(); + assert_eq!(decoded[0]["at"], json!(1_700_000_000_000_000i64)); + } + + #[test] + fn projected_batches_decode_without_the_missing_columns() { + // A list-style scan drops blob columns; decoding must not fail on them. + let spec = spec(); + let schema = Arc::new(spec.to_arrow().unwrap()); + let batch = rows_to_batch( + &spec, + schema, + &[row(json!({"id": "r1", "blob": [1, 2, 3], "count": 5}))], + ) + .unwrap(); + + let projected = batch.project(&[0, 1]).unwrap(); + let decoded = batch_to_rows(&spec, &projected).unwrap(); + assert_eq!(decoded[0]["id"], json!("r1")); + assert_eq!(decoded[0]["count"], json!(5)); + assert!(!decoded[0].contains_key("blob")); + } + + #[test] + fn column_order_follows_the_dataset_not_the_row() { + // Assembly is keyed by name, so a physical schema whose field order + // differs from the declaration still encodes correctly. + let spec = spec(); + let arrow = spec.to_arrow().unwrap(); + let reversed: Vec<_> = arrow.fields().iter().rev().cloned().collect(); + let reversed = Arc::new(Schema::new(reversed)); + + let batch = + rows_to_batch(&spec, reversed, &[row(json!({"id": "r1", "count": 9}))]).unwrap(); + assert_eq!(batch.schema().field(0).name(), "at"); + let decoded = batch_to_rows(&spec, &batch).unwrap(); + assert_eq!(decoded[0]["id"], json!("r1")); + assert_eq!(decoded[0]["count"], json!(9)); + } +} diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs new file mode 100644 index 0000000..5f4d955 --- /dev/null +++ b/crates/lance-context-core/src/generic_store.rs @@ -0,0 +1,701 @@ +//! A store over a user-declared schema. +//! +//! [`GenericStore`] is the fourth store alongside the three fixed-schema ones, +//! and the only one whose columns come from the caller rather than from a +//! hard-coded `fn schema()`. It is deliberately thin: the schema comes from a +//! [`SchemaSpec`], encode/decode from [`crate::generic_codec`], and *all* +//! storage behavior from [`StorageBase`] — the same code path the built-in +//! stores use. +//! +//! That last point is the design constraint, not an implementation detail. +//! `add` and the WAL merge behave identically here to `RolloutStore`, because +//! they are literally the same functions: the resident writer, the fence retry, +//! the surgical generation drain, compaction with `defer_index_remap`, and the +//! LSM read path are all inherited, not reimplemented. +//! +//! # The schema is persisted, not re-declared +//! +//! The spec is written into the dataset's schema metadata at creation and read +//! back on open, so callers pass it once. Reopening with a *different* spec is +//! an error rather than a silent reinterpretation of existing data. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow_array::RecordBatch; +use arrow_schema::{ArrowError, Schema}; +use futures::TryStreamExt; +use lance::dataset::optimize::CompactionMetrics; +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}; + +/// Schema-metadata key holding the serialized [`SchemaSpec`], so a store can be +/// reopened without the caller re-declaring its columns. +const SCHEMA_SPEC_KEY: &str = "lance-context:schema_spec"; + +/// Configuration for opening a [`GenericStore`]. +#[derive(Debug, Clone, Default)] +pub struct GenericStoreOptions { + /// Object-store credentials/config, forwarded to Lance. + pub storage_options: Option>, + /// Stable identity of this writer instance, mapped to a MemWAL shard. Each + /// instance owns exactly one shard, so no two contend. `None` uses a single + /// `"default"` shard — correct for one writer, wrong for several. + pub shard_id: Option, + /// Fold this instance's flushed generations into the base table once it has + /// accumulated this many. `None`/`0` disables the count trigger. + pub merge_after_generations: Option, + /// Shared, capacity-bounded Lance session. + pub session: Option>, + /// Whether [`GenericStore::add`] seals before returning, making the rows it + /// wrote immediately readable. + /// + /// Defaults to `false` — [`crate::RolloutStore`]'s profile: `add` is a + /// durable WAL append, concurrent appends do not serialize behind a + /// per-append seal, and no flushed generation is emitted per call. There is + /// then **no read-your-write guarantee**; drive [`GenericStore::flush`] + /// periodically, or set this to `true` to seal on every `add`. + pub seal_on_add: bool, +} + +/// A Lance-backed store over a user-declared schema. +pub struct GenericStore { + base: StorageBase, + spec: SchemaSpec, + /// Arrow schema as it exists on disk. Encoding assembles against *this*, + /// not against `spec.to_arrow()`, so a dataset whose physical field order + /// differs still writes correctly. + schema: Arc, +} + +impl GenericStore { + /// Create a store with `spec`, or open it if it already exists. + /// + /// # Errors + /// + /// - `spec` violates [`SchemaSpec::validate`] (no `id`, reserved names, …) + /// - the store exists and its persisted spec differs from `spec` + pub async fn open( + uri: &str, + spec: SchemaSpec, + options: GenericStoreOptions, + ) -> LanceResult { + Self::open_inner(uri, Some(spec), options, true).await + } + + /// Open an existing store, reading its schema back from the dataset. + /// + /// # Errors + /// + /// Returns [`LanceError::DatasetNotFound`] when the store does not exist, + /// rather than creating an empty one — a mistyped name must not silently + /// materialize a new store. + pub async fn open_existing(uri: &str, options: GenericStoreOptions) -> LanceResult { + Self::open_inner(uri, None, options, false).await + } + + async fn open_inner( + uri: &str, + spec: Option, + options: GenericStoreOptions, + create_if_missing: bool, + ) -> LanceResult { + if let Some(spec) = &spec { + spec.validate() + .map_err(|error| LanceError::from(ArrowError::SchemaError(error)))?; + } + + // Creating needs a schema up front; opening reads it back from the + // dataset. When no schema was supplied the store must already exist, so + // read its schema first and hand *that* to the base — it validates the + // key column against whatever it is given, and an empty placeholder + // would fail that check. + let create_schema = match &spec { + Some(spec) => Arc::new(schema_with_spec(spec)?), + None if create_if_missing => { + return Err(ArrowError::SchemaError( + "a schema is required to create a store".to_string(), + ) + .into()) + } + None => { + let existing = StorageBase::load_with_options( + uri, + options.storage_options.clone(), + options.session.clone(), + ) + .await?; + Arc::new(Schema::from(existing.schema())) + } + }; + + let base = StorageBase::open( + uri, + StorageBaseOptions { + storage_options: options.storage_options, + shard_id: options.shard_id, + merge_after_generations: options.merge_after_generations, + session: options.session, + schema: create_schema, + // Always `id`: the LSM merge key, which `SchemaSpec::validate` + // guarantees exists as a non-nullable string. + key_column: ID_COLUMN.to_string(), + // User schemas are immutable in v1, so there is nothing to + // evolve an older base table to. + latest_schema: None, + seal_on_put: options.seal_on_add, + }, + create_if_missing, + ) + .await?; + + let schema: Arc = Arc::new(base.dataset.schema().into()); + let persisted = spec_from_schema(&schema)?; + + // Reopening with a different schema would reinterpret existing data. + if let Some(requested) = spec { + if requested != persisted { + return Err(ArrowError::SchemaError(format!( + "store at '{uri}' was created with a different schema; \ + open it without a schema to use the persisted one" + )) + .into()); + } + } + + Ok(Self { + base, + spec: persisted, + schema, + }) + } + + /// The schema this store was created with. + #[must_use] + pub fn spec(&self) -> &SchemaSpec { + &self.spec + } + + /// URI of the underlying Lance dataset. + #[must_use] + pub fn uri(&self) -> &str { + self.base.uri() + } + + /// Current base dataset version. + #[must_use] + pub fn version(&self) -> u64 { + self.base.version() + } + + /// Append rows. + /// + /// Rows are matched to columns by name; an undeclared key is an error, and + /// an omitted nullable column is written as null. Visibility on return is + /// governed by [`GenericStoreOptions::seal_on_add`]. + /// + /// # Errors + /// + /// Propagates encoding failures (missing required column, type mismatch, + /// wrong vector length) before anything is written. + pub async fn add(&self, rows: &[Row]) -> LanceResult { + if rows.is_empty() { + return Ok(self.base.version()); + } + let batch = rows_to_batch(&self.spec, self.schema.clone(), rows)?; + self.base.put(vec![batch]).await?; + Ok(self.base.version()) + } + + /// Read rows, newest-generation-wins by `id`. + /// + /// Blob columns are projected out (see [`SchemaSpec::scan_columns`]), so a + /// list never materializes large payloads. Use [`Self::get`] with explicit + /// columns to fetch them. + pub async fn list(&self, limit: Option, offset: Option) -> LanceResult> { + self.scan(None, limit, offset, &self.spec.scan_columns()) + .await + } + + /// [`Self::list`], filtered by a SQL predicate over the store's columns. + pub async fn list_filtered( + &self, + filter: &str, + limit: Option, + offset: Option, + ) -> LanceResult> { + self.scan(Some(filter), limit, offset, &self.spec.scan_columns()) + .await + } + + /// Fetch one row by `id`, or `None` if absent. + /// + /// `columns` selects what to read: `None` reads everything *except* blob + /// columns, matching [`Self::list`]. Pass an explicit list to fetch blobs — + /// this is the intended way to read a large payload, one row at a time. + pub async fn get(&self, id: &str, columns: Option<&[String]>) -> LanceResult> { + let columns = columns.map_or_else(|| self.spec.scan_columns(), <[String]>::to_vec); + // Always read `id` so a match can be confirmed. + let mut columns = columns; + if !columns.iter().any(|column| column == ID_COLUMN) { + columns.push(ID_COLUMN.to_string()); + } + + let filter = format!("{ID_COLUMN} = '{}'", escape_sql_literal(id)); + let rows = self.scan(Some(&filter), Some(1), None, &columns).await?; + Ok(rows.into_iter().next()) + } + + async fn scan( + &self, + filter: Option<&str>, + limit: Option, + offset: Option, + columns: &[String], + ) -> LanceResult> { + let refs: Vec<&str> = columns.iter().map(String::as_str).collect(); + let mut scanner = self.base.lsm_scanner().await?.project(&refs); + if let Some(filter) = filter { + scanner = scanner.filter(filter)?; + } + if limit.is_some() || offset.is_some() { + scanner = scanner.limit(limit.unwrap_or(usize::MAX), offset); + } + + let mut stream = scanner.try_into_stream().await?; + let mut rows = Vec::new(); + while let Some(batch) = stream.try_next().await? { + rows.extend(batch_to_rows(&self.spec, &batch)?); + } + Ok(rows) + } + + /// Read rows from a chosen [`ListSource`] — base table, pending WAL + /// generations, or their union. + pub async fn list_from(&self, source: ListSource) -> LanceResult> { + let snapshots = match source { + ListSource::Fragments => Vec::new(), + ListSource::Wal | ListSource::All => self.base.wal_shard_snapshots().await?, + }; + let columns = self.spec.scan_columns(); + let refs: Vec<&str> = columns.iter().map(String::as_str).collect(); + let scanner = self + .base + .lsm_scanner_for_source(source, snapshots) + .project(&refs); + + let mut stream = scanner.try_into_stream().await?; + let mut rows = Vec::new(); + while let Some(batch) = stream.try_next().await? { + rows.extend(batch_to_rows(&self.spec, &batch)?); + } + Ok(rows) + } + + /// Seal the active memtable so previously added rows become readable + /// everywhere. A no-op when `seal_on_add` is set, since `add` already seals. + pub async fn flush(&self) -> LanceResult<()> { + self.base.flush().await + } + + /// Close the resident writer, draining its background tasks. Idempotent. + pub async fn close(&mut self) -> LanceResult<()> { + self.base.close().await + } + + /// Merge flushed generations into the base table once the count trigger is + /// met. Returns how many were reclaimed. + pub async fn maybe_merge_wal(&mut self) -> LanceResult { + self.base.maybe_merge_own_shard().await + } + + /// Seal, then merge **every** pending generation into the base table — the + /// time half of the "time OR count" trigger. + pub async fn cleanup_wal(&mut self) -> LanceResult { + self.base.cleanup_own_shard().await + } + + /// Generations pending merge across all shards. Read-only. + pub async fn pending_wal_generations(&self) -> LanceResult { + self.base.pending_wal_generations().await + } + + /// Compact the base table's small fragments. Drive from a single external + /// trigger, not per worker — see [`StorageBase::compact`]. + pub async fn compact( + &mut self, + options: Option, + ) -> LanceResult { + self.base.compact(options).await + } + + /// Whether the base table is fragmented enough to be worth compacting. + #[must_use] + pub fn should_compact(&self, config: &CompactionConfig) -> bool { + self.base.should_compact(config) + } + + /// Compaction statistics for the base table. + #[must_use] + pub fn compaction_stats(&self) -> CompactionStats { + self.base.compaction_stats() + } + + /// Build a ZoneMap scalar index on `id`. Idempotent. + pub async fn create_id_index(&mut self) -> LanceResult<()> { + self.base.create_key_zonemap_index().await + } + + /// Row count of the base table. Excludes rows still in unmerged + /// generations or buffered in the writer. + pub async fn count_base_rows(&self) -> LanceResult { + self.base.dataset.count_rows(None).await + } +} + +/// Attach the serialized spec to the Arrow schema so it round-trips on open. +fn schema_with_spec(spec: &SchemaSpec) -> Result { + let schema = spec.to_arrow()?; + let encoded = serde_json::to_string(spec) + .map_err(|error| ArrowError::SchemaError(format!("could not serialize schema: {error}")))?; + + let mut metadata = schema.metadata().clone(); + metadata.insert(SCHEMA_SPEC_KEY.to_string(), encoded); + Ok(schema.with_metadata(metadata)) +} + +/// Read the spec back out of a dataset's schema metadata. +fn spec_from_schema(schema: &Schema) -> Result { + let encoded = schema.metadata().get(SCHEMA_SPEC_KEY).ok_or_else(|| { + ArrowError::SchemaError( + "dataset has no stored schema spec: it was not created as a generic store".to_string(), + ) + })?; + serde_json::from_str(encoded) + .map_err(|error| ArrowError::SchemaError(format!("stored schema spec is invalid: {error}"))) +} + +/// Escape single quotes for a SQL string literal. +fn escape_sql_literal(value: &str) -> String { + value.replace('\'', "''") +} + +/// Row batches, for callers that already have Arrow data. +impl GenericStore { + /// Append pre-built [`RecordBatch`]es, bypassing row encoding. + /// + /// For bulk ingest where the caller already holds Arrow data. Each batch + /// must match the store's schema exactly. + /// + /// # Errors + /// + /// Returns an error if a batch's schema differs from the store's. + pub async fn add_batches(&self, batches: Vec) -> LanceResult { + if batches.is_empty() { + return Ok(self.base.version()); + } + for batch in &batches { + if batch.schema().fields() != self.schema.fields() { + return Err(ArrowError::SchemaError( + "record batch schema does not match the store schema".to_string(), + ) + .into()); + } + } + self.base.put(batches).await?; + Ok(self.base.version()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::schema_spec::{ColumnSpec, ColumnType}; + use serde_json::json; + use tempfile::TempDir; + + fn spec() -> SchemaSpec { + SchemaSpec::new(vec![ + ( + ID_COLUMN.to_string(), + ColumnSpec::required(ColumnType::String { large: false }), + ), + ( + "user_id".to_string(), + ColumnSpec::new(ColumnType::String { large: false }), + ), + ("score".to_string(), ColumnSpec::new(ColumnType::Float32)), + ( + "payload".to_string(), + ColumnSpec::new(ColumnType::Binary { blob: true }), + ), + ]) + } + + fn row(value: serde_json::Value) -> Row { + value.as_object().unwrap().clone() + } + + fn sealing() -> GenericStoreOptions { + // Seal on add so each test reads its own writes without an explicit + // flush; the deferred default is exercised separately. + GenericStoreOptions { + seal_on_add: true, + ..Default::default() + } + } + + #[test] + fn user_schema_round_trips_add_and_read() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let store = GenericStore::open(&uri, spec(), sealing()).await.unwrap(); + store + .add(&[ + row(json!({"id": "r1", "user_id": "u1", "score": 0.5})), + row(json!({"id": "r2", "user_id": "u2"})), + ]) + .await + .unwrap(); + + let rows = store.list(None, None).await.unwrap(); + assert_eq!(rows.len(), 2); + + let found = store.get("r1", None).await.unwrap().unwrap(); + assert_eq!(found["user_id"], json!("u1")); + assert!(store.get("missing", None).await.unwrap().is_none()); + }); + } + + #[test] + fn schema_is_persisted_and_reopens_without_being_redeclared() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + { + let store = GenericStore::open(&uri, spec(), sealing()).await.unwrap(); + store.add(&[row(json!({"id": "r1"}))]).await.unwrap(); + } + + let reopened = GenericStore::open_existing(&uri, sealing()).await.unwrap(); + assert_eq!(reopened.spec(), &spec()); + assert_eq!(reopened.list(None, None).await.unwrap().len(), 1); + }); + } + + #[test] + fn reopening_with_a_conflicting_schema_is_rejected() { + // Silently reinterpreting existing data under a new schema is the + // failure mode worth preventing here. + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + GenericStore::open(&uri, spec(), sealing()).await.unwrap(); + + let mut other = spec(); + other + .columns + .push(("extra".to_string(), ColumnSpec::new(ColumnType::Int64))); + let err = GenericStore::open(&uri, other, sealing()) + .await + .err() + .expect("conflicting schema must be rejected"); + assert!(err.to_string().contains("different schema"), "{err}"); + }); + } + + #[test] + fn invalid_schema_is_rejected_before_anything_is_created() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + // No `id` column. + let bad = SchemaSpec::new(vec![( + "name".to_string(), + ColumnSpec::new(ColumnType::String { large: false }), + )]); + let err = GenericStore::open(&uri, bad, sealing()) + .await + .err() + .expect("invalid schema must be rejected"); + assert!(err.to_string().contains("must declare an 'id'"), "{err}"); + }); + } + + #[test] + fn blob_columns_are_excluded_from_list_but_fetchable_per_row() { + // The cost model for large payloads: bulk reads never materialize them, + // point reads can ask for them explicitly. + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let store = GenericStore::open(&uri, spec(), sealing()).await.unwrap(); + store + .add(&[row(json!({"id": "r1", "payload": [1, 2, 3]}))]) + .await + .unwrap(); + + let listed = store.list(None, None).await.unwrap(); + assert!( + !listed[0].contains_key("payload"), + "list must not materialize blob columns" + ); + + let fetched = store + .get("r1", Some(&["payload".to_string()])) + .await + .unwrap() + .unwrap(); + assert_eq!(fetched["payload"], json!("AQID")); + }); + } + + #[test] + fn large_blob_round_trips_and_survives_wal_merge() { + // Multi-megabyte inline payloads are the point of this store, so pin + // that they survive both the write path and a WAL merge. + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let mut store = GenericStore::open(&uri, spec(), sealing()).await.unwrap(); + let payload: Vec = (0..4 * 1024 * 1024).map(|i| (i % 251) as u8).collect(); + store + .add(&[row(json!({"id": "big", "payload": payload}))]) + .await + .unwrap(); + + store.cleanup_wal().await.unwrap(); + assert_eq!(store.pending_wal_generations().await.unwrap(), 0); + + let fetched = store + .get("big", Some(&["payload".to_string()])) + .await + .unwrap() + .unwrap(); + use base64::Engine; + let decoded = base64::engine::general_purpose::STANDARD + .decode(fetched["payload"].as_str().unwrap()) + .unwrap(); + assert_eq!(decoded.len(), 4 * 1024 * 1024); + }); + } + + #[test] + fn deferred_seal_defers_visibility_like_the_rollout_store() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let store = GenericStore::open(&uri, spec(), GenericStoreOptions::default()) + .await + .unwrap(); + store.add(&[row(json!({"id": "r1"}))]).await.unwrap(); + assert_eq!(store.list(None, None).await.unwrap().len(), 0); + + store.flush().await.unwrap(); + assert_eq!(store.list(None, None).await.unwrap().len(), 1); + }); + } + + #[test] + fn wal_generations_merge_into_the_base_table() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let mut store = GenericStore::open(&uri, spec(), sealing()).await.unwrap(); + for i in 0..3 { + store + .add(&[row(json!({"id": format!("r{i}")}))]) + .await + .unwrap(); + } + assert!(store.pending_wal_generations().await.unwrap() > 0); + + let reclaimed = store.cleanup_wal().await.unwrap(); + assert!(reclaimed > 0); + assert_eq!(store.pending_wal_generations().await.unwrap(), 0); + assert_eq!(store.count_base_rows().await.unwrap(), 3); + assert_eq!(store.list(None, None).await.unwrap().len(), 3); + }); + } + + #[test] + fn duplicate_ids_dedup_to_the_newest_write() { + // The LSM merge key is `id`, so a re-added id supersedes the old row. + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let store = GenericStore::open(&uri, spec(), sealing()).await.unwrap(); + store + .add(&[row(json!({"id": "r1", "user_id": "first"}))]) + .await + .unwrap(); + store + .add(&[row(json!({"id": "r1", "user_id": "second"}))]) + .await + .unwrap(); + + let rows = store.list(None, None).await.unwrap(); + assert_eq!(rows.len(), 1, "id is the merge key, so rows dedup"); + assert_eq!(rows[0]["user_id"], json!("second")); + }); + } + + #[test] + fn filters_and_paging_work_over_user_columns() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let store = GenericStore::open(&uri, spec(), sealing()).await.unwrap(); + let rows: Vec = (0..6) + .map(|i| { + row(json!({ + "id": format!("r{i}"), + "user_id": if i % 2 == 0 { "even" } else { "odd" }, + "score": i as f32, + })) + }) + .collect(); + store.add(&rows).await.unwrap(); + + let evens = store + .list_filtered("user_id = 'even'", None, None) + .await + .unwrap(); + assert_eq!(evens.len(), 3); + assert!(evens.iter().all(|row| row["user_id"] == json!("even"))); + + let page = store.list(Some(2), Some(1)).await.unwrap(); + assert_eq!(page.len(), 2); + }); + } + + #[test] + fn open_existing_does_not_create() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().join("absent").to_string_lossy().to_string(); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let err = GenericStore::open_existing(&uri, sealing()) + .await + .err() + .expect("open_existing must not create"); + assert!(matches!(err, LanceError::DatasetNotFound { .. }), "{err}"); + }); + } +} diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 13da0f6..40f2d6d 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -7,6 +7,8 @@ mod datagen; mod datagen_store; mod eval; mod export; +pub mod generic_codec; +mod generic_store; mod id; pub mod metrics; mod namespace; @@ -14,6 +16,7 @@ mod record; mod registry; mod rollout; mod rollout_store; +pub mod schema_spec; pub mod serde; mod storage; mod store; @@ -39,6 +42,8 @@ pub use export::{ RolloutExample, RolloutResponse, SftExample, SplitConfig, SplitManifest, TokenStats, EXPORT_SCHEMA_VERSION, }; +pub use generic_codec::{batch_to_rows, ids_from_batch, rows_to_batch, Row}; +pub use generic_store::{GenericStore, GenericStoreOptions}; pub use id::{generate_id, new_uuid}; pub use namespace::{ContextNamespace, PartitionInfo, PartitionSelector, PartitionSpec}; pub use record::{ @@ -53,6 +58,7 @@ 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}; 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-core/src/schema_spec.rs b/crates/lance-context-core/src/schema_spec.rs new file mode 100644 index 0000000..d5da7ef --- /dev/null +++ b/crates/lance-context-core/src/schema_spec.rs @@ -0,0 +1,642 @@ +//! User-defined schemas: declaration, validation, and Arrow mapping. +//! +//! The three built-in stores each hard-code an Arrow schema and then repeat its +//! field list across five to seven places (schema fn, record→Arrow builders, +//! Arrow→record decode, projection, filters, merge key, index name). A +//! [`SchemaSpec`] replaces that with a *value*: declared once at store creation, +//! persisted with the store, and read back on open — so the encode/decode path +//! can be driven by it instead of by hand-written column lists. +//! +//! # What is deliberately constrained +//! +//! A user schema is not an arbitrary Arrow schema. The storage layer +//! ([`crate::store_base::StorageBase`]) makes hard assumptions that a schema +//! must satisfy for `add` and `merge` to behave identically to the built-in +//! stores: +//! +//! - **An `id` column is mandatory** (`Utf8`, non-nullable). It is the LSM merge +//! key, which is what makes a retried append idempotent and a crashed merge +//! safe to redo, and it is the column the scalar index is built on. +//! - **Reserved names are rejected**, because they collide with Lance internals +//! or with this crate's own sentinels. +//! - **Blob columns are declared, not inferred**, and are stored *inline* — see +//! [`ColumnType::Binary`]. +//! - **Vector columns carry their dimension and metric**, because an index +//! cannot be built without them. +//! +//! # Schema is immutable after creation +//! +//! v1 deliberately has no evolution path: the schema is fixed when the store is +//! created. Additive evolution already exists for the built-in rollout schema +//! (`StorageBaseOptions::latest_schema`) and can be extended to user schemas +//! later, but doing it safely needs a versioning story that is out of scope +//! here. + +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"; + +/// Marks the key column as a primary key that Lance does not enforce. Matches +/// the metadata the built-in schemas set on their own key columns. +const UNENFORCED_PRIMARY_KEY: &str = "lance-schema:unenforced-primary-key"; + +/// Schema-metadata key under which a vector column's distance metric is +/// persisted, so it round-trips on open without being re-specified. +const DISTANCE_METRIC_KEY: &str = "lance-context:distance_metric"; + +/// Field-metadata key marking a column as a blob: excluded from scan +/// projections by default. See [`ColumnType::Binary`]. +const BLOB_COLUMN_KEY: &str = "lance-context:blob"; + +/// Column names a user schema may not use. +/// +/// `_rowid` and `_mem_wal` are Lance's; `_rowaddr` likewise. Anything starting +/// with `_` is reserved wholesale so future Lance internals cannot collide with +/// a user column. +const RESERVED_COLUMNS: &[&str] = &["_rowid", "_rowaddr", "_mem_wal", "_distance"]; + +/// Maximum number of columns in a user schema. A guard against pathological +/// declarations, not a Lance limit. +const MAX_COLUMNS: usize = 1024; + +/// Declared type of one user column. +/// +/// Serialized in both a shorthand form (`"string"`) and a tagged form +/// (`{"type": "vector", "dim": 768}`) — see [`ColumnSpec`]. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ColumnType { + /// UTF-8 string. `large: true` selects `LargeUtf8` for values over ~2 GiB + /// cumulative per array. + String { + #[serde(default)] + large: bool, + }, + /// Signed integers. + Int32, + Int64, + /// Floating point. + Float32, + Float64, + Bool, + /// Microsecond-precision UTC timestamp. + Timestamp, + /// Inline binary payload — the encoding for large blobs. + /// + /// # Blobs are stored inline, never blob-v2 offloaded + /// + /// Reads go through the MemWAL LSM scanner, which has **no + /// blob-materialization step**: a blob-v2 (`lance-encoding:blob`) column + /// reads back as `None` through it. Inline `LargeBinary` is therefore the + /// only encoding that round-trips, which is why `offload` is not an option + /// here and why [`SchemaSpec::validate`] rejects any attempt to set the + /// blob-v2 metadata directly. + /// + /// Measured on this storage path, a 120 MB inline value writes in ~1.4 s + /// and reads back in ~0.8 s, and survives a WAL merge intact. + /// + /// `blob: true` does not change *where* bytes live; it changes *who pays + /// for them*. Blob columns are excluded from scan projections by default + /// (see [`SchemaSpec::scan_columns`]), so `list`-style reads never + /// materialize them — fetch them per row instead. + Binary { + /// Exclude from default scan projections. Set this for anything large + /// enough that you would not want it in a list response. + #[serde(default)] + blob: bool, + }, + /// Fixed-size float vector, for similarity search. + Vector { + /// Vector width. Must match on every write. + dim: i32, + /// Ranking metric. Defaults to `l2`. + #[serde(default)] + metric: Option, + }, + /// Variable-length list of a scalar type. + List { + /// Element type. Nested lists and lists of vectors are not supported. + item: Box, + }, +} + +/// One column in a user schema. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ColumnSpec { + /// Column type. + #[serde(flatten)] + pub column_type: ColumnType, + /// Whether the column accepts nulls. Defaults to `true`; a write may then + /// omit the column entirely and it is materialized as null. + #[serde(default = "default_true")] + pub nullable: bool, +} + +fn default_true() -> bool { + true +} + +impl ColumnSpec { + /// A nullable column of the given type. + #[must_use] + pub fn new(column_type: ColumnType) -> Self { + Self { + column_type, + nullable: true, + } + } + + /// A non-nullable column of the given type. + #[must_use] + pub fn required(column_type: ColumnType) -> Self { + Self { + column_type, + nullable: false, + } + } + + /// Whether this column is excluded from default scan projections. + #[must_use] + pub fn is_blob(&self) -> bool { + matches!(self.column_type, ColumnType::Binary { blob: true }) + } +} + +/// A user-declared store schema. +/// +/// Column order is preserved: [`SchemaSpec::to_arrow`] emits fields in +/// declaration order, with `id` first regardless of where it was declared, so +/// the physical layout is predictable. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct SchemaSpec { + /// Columns in declaration order. + pub columns: Vec<(String, ColumnSpec)>, +} + +impl SchemaSpec { + /// Build a spec from an ordered column list. Call [`Self::validate`] before + /// using it to create a store. + #[must_use] + pub fn new(columns: Vec<(String, ColumnSpec)>) -> Self { + Self { columns } + } + + /// Look up a column by name. + #[must_use] + pub fn column(&self, name: &str) -> Option<&ColumnSpec> { + self.columns + .iter() + .find(|(column, _)| column == name) + .map(|(_, spec)| spec) + } + + /// Names of the blob columns — those excluded from default scan + /// projections. + #[must_use] + pub fn blob_columns(&self) -> HashSet { + self.columns + .iter() + .filter(|(_, spec)| spec.is_blob()) + .map(|(name, _)| name.clone()) + .collect() + } + + /// Column names a `list`-style scan should read: everything except blob + /// columns, so large payloads are never materialized by a bulk read. + #[must_use] + pub fn scan_columns(&self) -> Vec { + self.columns + .iter() + .filter(|(_, spec)| !spec.is_blob()) + .map(|(name, _)| name.clone()) + .collect() + } + + /// Names of the vector columns, paired with their dimension. + #[must_use] + pub fn vector_columns(&self) -> Vec<(String, i32)> { + self.columns + .iter() + .filter_map(|(name, spec)| match spec.column_type { + ColumnType::Vector { dim, .. } => Some((name.clone(), dim)), + _ => None, + }) + .collect() + } + + /// Validate the declaration against everything the storage layer assumes. + /// + /// # Errors + /// + /// Returns the first violated rule: + /// - no columns, or more than [`MAX_COLUMNS`] + /// - missing `id`, or `id` that is not a non-nullable string + /// - duplicate, empty, or reserved column names + /// - a vector column with a non-positive dimension or an unknown metric + /// - a nested list, or a list of vectors + pub fn validate(&self) -> Result<(), String> { + if self.columns.is_empty() { + return Err("schema must declare at least one column".to_string()); + } + if self.columns.len() > MAX_COLUMNS { + return Err(format!( + "schema declares {} columns, the maximum is {MAX_COLUMNS}", + self.columns.len() + )); + } + + let mut seen = HashSet::new(); + for (name, spec) in &self.columns { + if name.is_empty() { + return Err("column names must not be empty".to_string()); + } + if !seen.insert(name.as_str()) { + return Err(format!("column '{name}' is declared more than once")); + } + if name.starts_with('_') || RESERVED_COLUMNS.contains(&name.as_str()) { + return Err(format!( + "column '{name}' is reserved: names starting with '_' are used by Lance \ + internals" + )); + } + validate_column_type(name, &spec.column_type)?; + } + + // `id` is what makes the whole storage layer work: it is the LSM merge + // key, so it must be present, non-null and comparable on every row. + let Some(id) = self.column(ID_COLUMN) else { + return Err(format!( + "schema must declare an '{ID_COLUMN}' column: it is the primary key and the \ + LSM merge key" + )); + }; + if !matches!(id.column_type, ColumnType::String { .. }) { + return Err(format!( + "column '{ID_COLUMN}' must be a string, got {:?}", + id.column_type + )); + } + if id.nullable { + return Err(format!( + "column '{ID_COLUMN}' must be non-nullable: it is the primary key" + )); + } + + Ok(()) + } + + /// Convert to an Arrow schema, with `id` first and the declared columns + /// following in order. + /// + /// Validates first, so an invalid spec cannot produce a schema. + /// + /// # Errors + /// + /// Propagates [`Self::validate`] as an [`ArrowError::SchemaError`]. + pub fn to_arrow(&self) -> Result { + self.validate().map_err(ArrowError::SchemaError)?; + + let mut fields = Vec::with_capacity(self.columns.len()); + + // `id` leads, wherever it was declared, so physical layout does not + // depend on declaration order. + let id = self + .column(ID_COLUMN) + .expect("validate() proved id is present"); + fields.push(arrow_field(ID_COLUMN, id)?); + + for (name, spec) in &self.columns { + if name == ID_COLUMN { + continue; + } + fields.push(arrow_field(name, spec)?); + } + + // Persist each vector column's metric so `open` does not have to be + // told again. Keyed per column, since a schema may declare several. + let mut metadata = HashMap::new(); + for (name, spec) in &self.columns { + if let ColumnType::Vector { metric, .. } = &spec.column_type { + let metric = metric.as_deref().unwrap_or("l2"); + metadata.insert(format!("{DISTANCE_METRIC_KEY}:{name}"), metric.to_string()); + } + } + + Ok(Schema::new(fields).with_metadata(metadata)) + } +} + +fn validate_column_type(name: &str, column_type: &ColumnType) -> Result<(), String> { + match column_type { + ColumnType::Vector { dim, metric } => { + if *dim <= 0 { + return Err(format!( + "column '{name}': vector dimension must be positive, got {dim}" + )); + } + if let Some(metric) = metric { + DistanceMetric::parse(metric).map_err(|_| { + format!( + "column '{name}': invalid distance metric '{metric}', expected one of \ + 'l2', 'cosine', 'dot'" + ) + })?; + } + Ok(()) + } + ColumnType::List { item } => match **item { + // A nested list would need a recursive builder and a recursive + // decoder; the built-in schemas never use one, so it stays out + // until something needs it. + ColumnType::List { .. } => { + Err(format!("column '{name}': nested lists are not supported")) + } + ColumnType::Vector { .. } => Err(format!( + "column '{name}': a list of vectors is not supported; declare a vector column" + )), + _ => Ok(()), + }, + _ => Ok(()), + } +} + +fn arrow_field(name: &str, spec: &ColumnSpec) -> Result { + let data_type = arrow_data_type(&spec.column_type)?; + let mut field = Field::new(name, data_type, spec.nullable); + + let mut metadata = HashMap::new(); + if name == ID_COLUMN { + metadata.insert(UNENFORCED_PRIMARY_KEY.to_string(), "true".to_string()); + } + if spec.is_blob() { + // Marks the column for projection exclusion. Deliberately *not* + // `lance-encoding:blob`: that is blob-v2 offload, which the LSM read + // path cannot materialize (see `ColumnType::Binary`). + metadata.insert(BLOB_COLUMN_KEY.to_string(), "true".to_string()); + } + if !metadata.is_empty() { + field = field.with_metadata(metadata); + } + Ok(field) +} + +fn arrow_data_type(column_type: &ColumnType) -> Result { + Ok(match column_type { + ColumnType::String { large: false } => DataType::Utf8, + ColumnType::String { large: true } => DataType::LargeUtf8, + ColumnType::Int32 => DataType::Int32, + ColumnType::Int64 => DataType::Int64, + ColumnType::Float32 => DataType::Float32, + ColumnType::Float64 => DataType::Float64, + ColumnType::Bool => DataType::Boolean, + ColumnType::Timestamp => DataType::Timestamp(TimeUnit::Microsecond, None), + // Always `LargeBinary`, inline. See `ColumnType::Binary`. + ColumnType::Binary { .. } => DataType::LargeBinary, + ColumnType::Vector { dim, .. } => DataType::FixedSizeList( + std::sync::Arc::new(Field::new("item", DataType::Float32, true)), + *dim, + ), + ColumnType::List { item } => DataType::List(std::sync::Arc::new(Field::new( + "item", + arrow_data_type(item)?, + true, + ))), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn id_column() -> (String, ColumnSpec) { + ( + ID_COLUMN.to_string(), + ColumnSpec::required(ColumnType::String { large: false }), + ) + } + + fn spec_with(extra: Vec<(String, ColumnSpec)>) -> SchemaSpec { + let mut columns = vec![id_column()]; + columns.extend(extra); + SchemaSpec::new(columns) + } + + #[test] + fn minimal_schema_is_just_id() { + let spec = spec_with(vec![]); + spec.validate().unwrap(); + let schema = spec.to_arrow().unwrap(); + assert_eq!(schema.fields().len(), 1); + let id = schema.field_with_name(ID_COLUMN).unwrap(); + assert_eq!(id.data_type(), &DataType::Utf8); + assert!(!id.is_nullable()); + assert_eq!( + id.metadata().get(UNENFORCED_PRIMARY_KEY), + Some(&"true".to_string()), + "id must carry the unenforced-primary-key marker like the built-in schemas" + ); + } + + #[test] + fn id_is_mandatory_and_must_be_a_non_null_string() { + let missing = SchemaSpec::new(vec![( + "user_id".to_string(), + ColumnSpec::new(ColumnType::String { large: false }), + )]); + assert!(missing + .validate() + .unwrap_err() + .contains("must declare an 'id'")); + + let wrong_type = SchemaSpec::new(vec![( + ID_COLUMN.to_string(), + ColumnSpec::required(ColumnType::Int64), + )]); + assert!(wrong_type + .validate() + .unwrap_err() + .contains("must be a string")); + + let nullable = SchemaSpec::new(vec![( + ID_COLUMN.to_string(), + ColumnSpec::new(ColumnType::String { large: false }), + )]); + assert!(nullable.validate().unwrap_err().contains("non-nullable")); + } + + #[test] + fn reserved_and_duplicate_names_are_rejected() { + for reserved in ["_rowid", "_mem_wal", "_anything"] { + let spec = spec_with(vec![( + reserved.to_string(), + ColumnSpec::new(ColumnType::Int64), + )]); + assert!( + spec.validate().unwrap_err().contains("reserved"), + "{reserved} must be rejected" + ); + } + + let dup = SchemaSpec::new(vec![ + id_column(), + ("x".to_string(), ColumnSpec::new(ColumnType::Int64)), + ("x".to_string(), ColumnSpec::new(ColumnType::Int64)), + ]); + assert!(dup.validate().unwrap_err().contains("more than once")); + } + + #[test] + fn vector_columns_carry_dimension_and_metric() { + let spec = spec_with(vec![( + "embedding".to_string(), + ColumnSpec::new(ColumnType::Vector { + dim: 768, + metric: Some("cosine".to_string()), + }), + )]); + let schema = spec.to_arrow().unwrap(); + let field = schema.field_with_name("embedding").unwrap(); + match field.data_type() { + DataType::FixedSizeList(_, dim) => assert_eq!(*dim, 768), + other => panic!("expected FixedSizeList, got {other:?}"), + } + assert_eq!( + schema + .metadata() + .get(&format!("{DISTANCE_METRIC_KEY}:embedding")), + Some(&"cosine".to_string()), + "the metric must round-trip so open() need not be told again" + ); + assert_eq!(spec.vector_columns(), vec![("embedding".to_string(), 768)]); + + let bad_dim = spec_with(vec![( + "e".to_string(), + ColumnSpec::new(ColumnType::Vector { + dim: 0, + metric: None, + }), + )]); + assert!(bad_dim.validate().unwrap_err().contains("must be positive")); + + let bad_metric = spec_with(vec![( + "e".to_string(), + ColumnSpec::new(ColumnType::Vector { + dim: 8, + metric: Some("manhattan".to_string()), + }), + )]); + assert!(bad_metric + .validate() + .unwrap_err() + .contains("invalid distance metric")); + } + + #[test] + fn blob_columns_are_inline_and_excluded_from_scans() { + let spec = spec_with(vec![ + ( + "video".to_string(), + ColumnSpec::new(ColumnType::Binary { blob: true }), + ), + ( + "thumbnail".to_string(), + ColumnSpec::new(ColumnType::Binary { blob: false }), + ), + ( + "caption".to_string(), + ColumnSpec::new(ColumnType::String { large: false }), + ), + ]); + let schema = spec.to_arrow().unwrap(); + + // Inline LargeBinary, never blob-v2: the LSM read path has no + // blob-materialization step, so an offloaded column reads back as None. + for column in ["video", "thumbnail"] { + let field = schema.field_with_name(column).unwrap(); + assert_eq!(field.data_type(), &DataType::LargeBinary); + assert!( + field.metadata().get("lance-encoding:blob").is_none(), + "{column} must not be blob-v2 offloaded" + ); + } + + // `blob: true` means "excluded from default scans", nothing else. + assert_eq!(spec.blob_columns(), HashSet::from(["video".to_string()])); + assert_eq!( + spec.scan_columns(), + vec![ + "id".to_string(), + "thumbnail".to_string(), + "caption".to_string() + ], + "a list-style scan must skip blob columns but keep everything else" + ); + } + + #[test] + fn id_leads_the_arrow_schema_regardless_of_declaration_order() { + let spec = SchemaSpec::new(vec![ + ("score".to_string(), ColumnSpec::new(ColumnType::Float32)), + id_column(), + ( + "tags".to_string(), + ColumnSpec::new(ColumnType::List { + item: Box::new(ColumnType::String { large: false }), + }), + ), + ]); + let schema = spec.to_arrow().unwrap(); + let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); + assert_eq!(names, vec!["id", "score", "tags"]); + } + + #[test] + fn nested_lists_and_vector_lists_are_rejected() { + let nested = spec_with(vec![( + "x".to_string(), + ColumnSpec::new(ColumnType::List { + item: Box::new(ColumnType::List { + item: Box::new(ColumnType::Int64), + }), + }), + )]); + assert!(nested.validate().unwrap_err().contains("nested lists")); + + let vec_list = spec_with(vec![( + "x".to_string(), + ColumnSpec::new(ColumnType::List { + item: Box::new(ColumnType::Vector { + dim: 4, + metric: None, + }), + }), + )]); + assert!(vec_list.validate().unwrap_err().contains("list of vectors")); + } + + #[test] + fn json_shorthand_round_trips() { + // The wire form users actually write. + let json = serde_json::json!({ + "columns": [ + ["id", {"type": "string", "nullable": false}], + ["score", {"type": "float32"}], + ["video", {"type": "binary", "blob": true}], + ["embedding", {"type": "vector", "dim": 4, "metric": "cosine"}] + ] + }); + let spec: SchemaSpec = serde_json::from_value(json).unwrap(); + spec.validate().unwrap(); + assert_eq!(spec.blob_columns(), HashSet::from(["video".to_string()])); + + let reparsed: SchemaSpec = + serde_json::from_str(&serde_json::to_string(&spec).unwrap()).unwrap(); + assert_eq!(reparsed, spec); + } +}