diff --git a/seaweed-vector/Cargo.lock b/seaweed-vector/Cargo.lock index 879a1c4fbb..8b9da51181 100644 --- a/seaweed-vector/Cargo.lock +++ b/seaweed-vector/Cargo.lock @@ -6475,6 +6475,34 @@ dependencies = [ "weed-lance-worker", ] +[[package]] +name = "weed-vector-server" +version = "0.1.0" +dependencies = [ + "anyhow", + "arrow-array", + "arrow-cast", + "arrow-schema", + "clap", + "futures", + "lance", + "lance-linalg", + "prost", + "prost-types", + "protoc-bin-vendored", + "reqwest 0.12.28", + "serde", + "serde_json", + "tokio", + "tonic", + "tonic-prost", + "tonic-prost-build", + "tracing", + "tracing-subscriber", + "weed-lance-worker", + "weed-vector-core", +] + [[package]] name = "winapi-util" version = "0.1.11" diff --git a/seaweed-vector/Cargo.toml b/seaweed-vector/Cargo.toml index 93fe4ff858..71648c1621 100644 --- a/seaweed-vector/Cargo.toml +++ b/seaweed-vector/Cargo.toml @@ -6,7 +6,7 @@ # milestone. [workspace] resolver = "2" -members = ["crates/core"] +members = ["crates/core", "crates/server"] [workspace.package] version = "0.1.0" diff --git a/seaweed-vector/crates/core/src/coverage.rs b/seaweed-vector/crates/core/src/coverage.rs index 4dd68620a6..ad7d8e650a 100644 --- a/seaweed-vector/crates/core/src/coverage.rs +++ b/seaweed-vector/crates/core/src/coverage.rs @@ -37,11 +37,7 @@ pub async fn coverage(dataset: &Dataset, vector_column: &str) -> Result Result<(), Box> { + if std::env::var_os("PROTOC").is_none() { + // SAFETY: a build script's main runs single-threaded before anything + // else in this process, so no other thread can be reading the + // environment concurrently. + unsafe { + std::env::set_var("PROTOC", protoc_bin_vendored::protoc_bin_path()?); + } + } + tonic_prost_build::configure() + .build_server(true) + .build_client(true) + .compile_protos(&["../../../weed/pb/vector.proto"], &["../../../weed/pb/"])?; + println!("cargo:rerun-if-changed=../../../weed/pb/vector.proto"); + Ok(()) +} diff --git a/seaweed-vector/crates/server/src/lib.rs b/seaweed-vector/crates/server/src/lib.rs new file mode 100644 index 0000000000..b46863fe45 --- /dev/null +++ b/seaweed-vector/crates/server/src/lib.rs @@ -0,0 +1,9 @@ +//! The VectorService gRPC server the S3 gateway talks to. + +pub mod service; + +pub mod pb { + pub mod vector { + tonic::include_proto!("vector_pb"); + } +} diff --git a/seaweed-vector/crates/server/src/main.rs b/seaweed-vector/crates/server/src/main.rs new file mode 100644 index 0000000000..7c0aa7a623 --- /dev/null +++ b/seaweed-vector/crates/server/src/main.rs @@ -0,0 +1,47 @@ +use anyhow::Result; +use clap::Parser; +use tonic::transport::Server; +use weed_lance_worker::catalog::NamespaceClient; +use weed_lance_worker::dataset::FallbackOptions; +use weed_vector_server::pb::vector::vector_service_server::VectorServiceServer; +use weed_vector_server::service::VectorServer; + +/// The VectorService the S3 gateway fans vector calls out to. +#[derive(Parser)] +struct Args { + /// Lance Namespace base URL, e.g. http://localhost:8888/lance. + #[arg( + long, + env = "WEED_VECTOR_NAMESPACE", + default_value = "http://localhost:8888/lance" + )] + namespace: String, + /// gRPC listen address. + #[arg(long, env = "WEED_VECTOR_LISTEN", default_value = "127.0.0.1:8080")] + listen: String, + /// object_store option fallback for deployments that vend no credentials: + /// key=value, repeatable. + #[arg(long = "storage-option", value_parser = parse_option)] + storage_options: Vec<(String, String)>, +} + +fn parse_option(s: &str) -> Result<(String, String), String> { + s.split_once('=') + .map(|(k, v)| (k.to_string(), v.to_string())) + .ok_or_else(|| "expected key=value".to_string()) +} + +#[tokio::main] +async fn main() -> Result<()> { + tracing_subscriber::fmt::init(); + let args = Args::parse(); + let fallback: FallbackOptions = args.storage_options.into_iter().collect(); + let client = NamespaceClient::new(args.namespace.clone()).with_fallback(&fallback); + let service = VectorServer::new(client, fallback); + tracing::info!(namespace = %args.namespace, listen = %args.listen, "weed-vector-server"); + Server::builder() + .add_service(VectorServiceServer::new(service)) + .serve(args.listen.parse()?) + .await?; + Ok(()) +} diff --git a/seaweed-vector/crates/server/src/service.rs b/seaweed-vector/crates/server/src/service.rs new file mode 100644 index 0000000000..e49a357ff2 --- /dev/null +++ b/seaweed-vector/crates/server/src/service.rs @@ -0,0 +1,615 @@ +use anyhow::{Result, bail}; +use arrow_array::{Array, ArrayRef, FixedSizeListArray, Float32Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema}; +use futures::TryStreamExt; +use lance::Dataset; +use lance::dataset::ProjectionRequest; +use lance::dataset::WriteParams; +use lance::dataset::{MergeInsertBuilder, WhenMatched, WhenNotMatched}; +use lance::io::ObjectStoreParams; +use lance::io::StorageOptionsAccessor; +use lance_linalg::distance::DistanceType; +use prost_types::{ListValue, Struct, Value, value::Kind}; +use std::collections::BTreeMap; +use std::sync::Arc; +use tonic::{Request, Response, Status}; +use weed_lance_worker::catalog::NamespaceClient; +use weed_lance_worker::dataset::{self, FallbackOptions}; +use weed_vector_core::query::{self, Query}; +use weed_vector_core::table::VectorTable; + +use crate::pb::vector::{self as pb, vector_service_server::VectorService}; + +fn internal(err: impl Into) -> Status { + Status::internal(format!("{:#}", err.into())) +} + +fn invalid(err: impl Into) -> Status { + Status::invalid_argument(format!("{:#}", err.into())) +} + +fn open_status(err: anyhow::Error) -> Status { + let missing = err.chain().any(|cause| { + cause + .downcast_ref::() + .is_some_and(|e| e.status() == Some(reqwest::StatusCode::NOT_FOUND)) + || cause.downcast_ref::().is_some_and(|e| { + matches!( + e, + lance::Error::DatasetNotFound { .. } + | lance::Error::NotFound { .. } + | lance::Error::VersionNotFound { .. } + ) + }) + }); + if missing { + Status::not_found(format!("{:#}", err)) + } else { + Status::unavailable(format!("{:#}", err)) + } +} + +const KEY_COLUMN: &str = "key"; +const METADATA_COLUMN: &str = "metadata"; +const VECTOR_COLUMN: &str = "vector"; + +fn schema_for( + dimension: usize, + metric: pb::DistanceMetric, + non_filterable: &[String], +) -> Arc { + let vector_field = Field::new( + VECTOR_COLUMN, + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, false)), + dimension as i32, + ), + false, + ) + .with_metadata( + [( + "distance_metric".to_string(), + metric.as_str_name().to_string(), + )] + .into_iter() + .collect(), + ); + let mut metadata_field = Field::new(METADATA_COLUMN, DataType::Utf8, true); + if !non_filterable.is_empty() { + metadata_field = metadata_field.with_metadata( + [( + "non_filterable_metadata_keys".to_string(), + serde_json::to_string(non_filterable).unwrap_or_default(), + )] + .into_iter() + .collect(), + ); + } + Arc::new(Schema::new(vec![ + Field::new(KEY_COLUMN, DataType::Utf8, false), + vector_field, + metadata_field, + ])) +} + +fn batch_of(vectors: &[pb::Vector], dimension: usize, schema: Arc) -> Result { + let mut seen = std::collections::HashSet::new(); + for vector in vectors { + if !seen.insert(vector.key.as_str()) { + bail!("duplicate key in request: {}", vector.key); + } + } + let keys = StringArray::from(vectors.iter().map(|v| v.key.as_str()).collect::>()); + // Check every vector's length before flattening: equal totals would + // still let a short vector borrow from the next one's values. + if vectors.iter().any(|v| v.data.len() != dimension) { + bail!("a vector's data length does not match the index dimension"); + } + let flat = vectors + .iter() + .flat_map(|v| v.data.iter().copied()) + .collect::>(); + let values: ArrayRef = Arc::new(Float32Array::from(flat)); + let vectors_array = FixedSizeListArray::try_new( + Field::new("item", DataType::Float32, false).into(), + dimension as i32, + values, + None, + )?; + let metas = vectors + .iter() + .map(|v| { + v.metadata + .as_ref() + .map(|m| serde_json::to_string(&struct_to_json(m))) + .transpose() + }) + .collect::, _>>()?; + RecordBatch::try_new( + schema, + vec![ + Arc::new(keys), + Arc::new(vectors_array), + Arc::new(StringArray::from(metas)), + ], + ) + .map_err(Into::into) +} + +fn struct_to_json(s: &Struct) -> serde_json::Value { + serde_json::Value::Object( + s.fields + .iter() + .map(|(k, v)| (k.clone(), value_to_json(v))) + .collect(), + ) +} + +fn value_to_json(v: &Value) -> serde_json::Value { + match &v.kind { + None | Some(Kind::NullValue(_)) => serde_json::Value::Null, + Some(Kind::NumberValue(n)) => serde_json::json!(n), + Some(Kind::StringValue(s)) => serde_json::json!(s), + Some(Kind::BoolValue(b)) => serde_json::json!(b), + Some(Kind::StructValue(s)) => struct_to_json(s), + Some(Kind::ListValue(l)) => { + serde_json::Value::Array(l.values.iter().map(value_to_json).collect()) + } + } +} + +fn json_to_value(v: serde_json::Value) -> Value { + let kind = match v { + serde_json::Value::Null => Kind::NullValue(0), + serde_json::Value::Bool(b) => Kind::BoolValue(b), + serde_json::Value::Number(n) => Kind::NumberValue(n.as_f64().unwrap_or(0.0)), + serde_json::Value::String(s) => Kind::StringValue(s), + serde_json::Value::Array(items) => Kind::ListValue(ListValue { + values: items.into_iter().map(json_to_value).collect(), + }), + serde_json::Value::Object(map) => Kind::StructValue(json_to_struct(map)), + }; + Value { kind: Some(kind) } +} + +fn json_to_struct(map: serde_json::Map) -> Struct { + Struct { + fields: map + .into_iter() + .map(|(k, v)| (k, json_to_value(v))) + .collect::>(), + } +} + +fn parse_struct(s: &str) -> Result { + let v: serde_json::Value = serde_json::from_str(s)?; + match v { + serde_json::Value::Object(map) => Ok(json_to_struct(map)), + _ => bail!("metadata is not a JSON object"), + } +} + +fn key_list(keys: &[String]) -> String { + format!( + "`{KEY_COLUMN}` in ({})", + keys.iter() + .map(|k| format!("'{}'", k.replace('\'', "''"))) + .collect::>() + .join(", ") + ) +} + +fn rows_as_vectors( + batches: &[RecordBatch], + return_data: bool, + return_metadata: bool, +) -> Vec { + let mut out = Vec::new(); + for batch in batches { + let Some(keys) = batch + .column_by_name(KEY_COLUMN) + .and_then(|c| c.as_any().downcast_ref::()) + else { + continue; + }; + let vectors = batch + .column_by_name(VECTOR_COLUMN) + .and_then(|c| c.as_any().downcast_ref::()); + let metas = batch + .column_by_name(METADATA_COLUMN) + .and_then(|c| c.as_any().downcast_ref::()); + for i in 0..batch.num_rows() { + let data = match (return_data, vectors) { + (true, Some(list)) => list + .value(i) + .as_any() + .downcast_ref::() + .map(|a| a.values().to_vec()) + .unwrap_or_default(), + _ => Vec::new(), + }; + let metadata = match (return_metadata, metas) { + (true, Some(m)) if !m.is_null(i) => parse_struct(m.value(i)).ok(), + _ => None, + }; + out.push(pb::Vector { + key: keys.value(i).to_string(), + data, + metadata, + }); + } + } + out +} + +pub struct VectorServer { + client: NamespaceClient, + fallback: FallbackOptions, +} + +impl VectorServer { + pub fn new(client: NamespaceClient, fallback: FallbackOptions) -> Self { + Self { client, fallback } + } + + async fn open(&self, table: &pb::Table) -> Result { + let opened = dataset::open(&self.client, &table.id, &self.fallback).await?; + let mut dataset = opened.dataset; + if table.version != 0 { + dataset = dataset.checkout_version(table.version).await?; + } + Ok(dataset) + } +} + +#[tonic::async_trait] +impl VectorService for VectorServer { + async fn create_index( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let metric = req.distance_metric(); + let table = req + .table + .ok_or_else(|| Status::invalid_argument("missing table"))?; + let dimension = req.dimension as usize; + // S3 Vectors caps at 4096; reject beyond it instead of letting an + // oversized request build a huge Arrow schema. + if dimension == 0 || dimension > 4096 { + return Err(Status::invalid_argument("dimension must be 1..=4096")); + } + let description = self + .client + .declare_table(&table.id, None) + .await + .map_err(internal)?; + let mut options = self.fallback.clone(); + options.extend(description.storage_options.clone()); + let schema = schema_for(dimension, metric, &req.non_filterable_metadata_keys); + let empty = RecordBatch::new_empty(schema.clone()); + let reader = arrow_array::RecordBatchIterator::new(vec![Ok(empty)].into_iter(), schema); + let dataset = Dataset::write( + reader, + &description.location, + Some(WriteParams { + store_params: Some(ObjectStoreParams { + storage_options_accessor: Some(Arc::new( + StorageOptionsAccessor::with_static_options(options), + )), + ..Default::default() + }), + ..Default::default() + }), + ) + .await + .map_err(internal)?; + Ok(Response::new(pb::CreateIndexResponse { + version: dataset.version().version, + })) + } + + async fn drop_index( + &self, + _request: Request, + ) -> Result, Status> { + // Dropping is a catalog decision; the gateway deletes the table record. + Err(Status::unimplemented("drop through the catalog")) + } + + async fn put_vectors( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let table = req + .table + .ok_or_else(|| Status::invalid_argument("missing table"))?; + let dataset = self.open(&table).await.map_err(open_status)?; + if req.vectors.is_empty() { + return Ok(Response::new(pb::PutVectorsResponse { + version: dataset.version().version, + })); + } + let dimension = VectorTable::resolve(&dataset, None, Some(KEY_COLUMN)) + .map_err(invalid)? + .dimension; + let batch = batch_of( + &req.vectors, + dimension, + Arc::new(Into::::into(dataset.schema())), + ) + .map_err(invalid)?; + // Upsert by key in one merge commit: no window where the old rows are + // already deleted, and a losing writer fails the request instead of + // leaving duplicate keys. + let mut builder = + MergeInsertBuilder::try_new(Arc::new(dataset), vec![KEY_COLUMN.to_string()]) + .map_err(internal)?; + builder.when_matched(WhenMatched::UpdateAll); + builder.when_not_matched(WhenNotMatched::InsertAll); + let (dataset, _stats) = builder + .try_build() + .map_err(internal)? + .execute_batches(vec![batch]) + .await + .map_err(internal)?; + Ok(Response::new(pb::PutVectorsResponse { + version: dataset.version().version, + })) + } + + async fn get_vectors( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let table = req + .table + .ok_or_else(|| Status::invalid_argument("missing table"))?; + let dataset = self.open(&table).await.map_err(open_status)?; + let resolved = VectorTable::resolve(&dataset, None, Some(KEY_COLUMN)).map_err(invalid)?; + let batches = query::get(&dataset, &resolved, &req.keys) + .await + .map_err(internal)?; + Ok(Response::new(pb::GetVectorsResponse { + vectors: rows_as_vectors(&batches, req.return_data, req.return_metadata), + })) + } + + async fn delete_vectors( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let table = req + .table + .ok_or_else(|| Status::invalid_argument("missing table"))?; + let mut dataset = self.open(&table).await.map_err(open_status)?; + if !req.keys.is_empty() { + dataset + .delete(&key_list(&req.keys)) + .await + .map_err(internal)?; + } + Ok(Response::new(pb::DeleteVectorsResponse { + version: dataset.version().version, + })) + } + + async fn list_vectors( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let table = req + .table + .ok_or_else(|| Status::invalid_argument("missing table"))?; + let dataset = self.open(&table).await.map_err(open_status)?; + let offset: u64 = if req.next_token.is_empty() { + 0 + } else { + req.next_token + .parse() + .map_err(|_| Status::invalid_argument("invalid next_token"))? + }; + let mut columns = vec![KEY_COLUMN.to_string()]; + if req.return_data { + columns.push(VECTOR_COLUMN.to_string()); + } + if req.return_metadata { + columns.push(METADATA_COLUMN.to_string()); + } + // An unset max_results arrives as 0; page it like S3 Vectors does. + const DEFAULT_PAGE: u32 = 500; + const MAX_PAGE: u32 = 1000; + let page = match req.max_results { + 0 => DEFAULT_PAGE, + n => n.min(MAX_PAGE), + }; + let mut scanner = dataset.scan(); + scanner + .project(&columns) + .map_err(internal)? + .limit(Some(page as i64), Some(offset as i64)) + .map_err(internal)?; + let batches: Vec = scanner + .try_into_stream() + .await + .map_err(internal)? + .try_collect() + .await + .map_err(internal)?; + let rows = batches.iter().map(|b| b.num_rows()).sum::(); + let next = offset + rows as u64; + let next_token = if rows == page as usize && rows > 0 { + next.to_string() + } else { + String::new() + }; + Ok(Response::new(pb::ListVectorsResponse { + vectors: rows_as_vectors(&batches, req.return_data, req.return_metadata), + next_token, + })) + } + + async fn query_vectors( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let table = req + .table + .ok_or_else(|| Status::invalid_argument("missing table"))?; + let dataset = self.open(&table).await.map_err(open_status)?; + let resolved = VectorTable::resolve(&dataset, None, Some(KEY_COLUMN)).map_err(invalid)?; + if req.query_vector.len() != resolved.dimension { + return Err(Status::invalid_argument(format!( + "query vector carries {} floats, the index dimension is {}", + req.query_vector.len(), + resolved.dimension + ))); + } + if req.top_k == 0 { + return Err(Status::invalid_argument("top_k must be positive")); + } + // The index's configured metric lives in the vector field's metadata. + let distance = match dataset + .schema() + .field(VECTOR_COLUMN) + .and_then(|f| f.metadata.get("distance_metric").map(String::as_str)) + { + Some("COSINE") => DistanceType::Cosine, + Some("DOT_PRODUCT") => DistanceType::Dot, + _ => DistanceType::L2, + }; + let query = Query { + vector: req.query_vector, + k: req.top_k as usize, + filter: if req.filter.is_empty() { + None + } else { + Some(req.filter) + }, + postfilter: req.postfilter, + distance, + nprobes: None, + }; + let hits = query::ann(&dataset, &resolved, &query) + .await + .map_err(internal)?; + // One batched read for every hit's columns; keyed by the hit key so + // row order does not matter. + let mut metas: std::collections::HashMap> = + std::collections::HashMap::new(); + let mut data_map: std::collections::HashMap> = + std::collections::HashMap::new(); + if (req.return_metadata || req.return_data) && !hits.is_empty() { + let row_ids: Vec = hits.iter().map(|h| h.row_id).collect(); + let mut columns = vec![KEY_COLUMN]; + if req.return_metadata { + columns.push(METADATA_COLUMN); + } + if req.return_data { + columns.push(VECTOR_COLUMN); + } + let batch = dataset + .take_rows( + &row_ids, + ProjectionRequest::from_columns(columns, dataset.schema()), + ) + .await + .map_err(internal)?; + let keys = batch + .column_by_name(KEY_COLUMN) + .and_then(|c| c.as_any().downcast_ref::()); + let meta_col = batch + .column_by_name(METADATA_COLUMN) + .and_then(|c| c.as_any().downcast_ref::()); + let data_col = batch + .column_by_name(VECTOR_COLUMN) + .and_then(|c| c.as_any().downcast_ref::()); + if let Some(keys) = keys { + for i in 0..batch.num_rows() { + let key = keys.value(i).to_string(); + if let Some(meta_col) = meta_col { + metas.insert( + key.clone(), + if meta_col.is_null(i) { + None + } else { + parse_struct(meta_col.value(i)).ok() + }, + ); + } + if let Some(data_col) = data_col { + data_map.insert( + key, + data_col + .value(i) + .as_any() + .downcast_ref::() + .map(|a| a.values().to_vec()) + .unwrap_or_default(), + ); + } + } + } + } + let mut vectors = Vec::with_capacity(hits.len()); + for hit in &hits { + vectors.push(pb::QueryResult { + key: hit.key.clone(), + distance: if req.return_distance { + hit.distance + } else { + 0.0 + }, + metadata: metas.remove(&hit.key).flatten(), + data: data_map.remove(&hit.key).unwrap_or_default(), + }); + } + Ok(Response::new(pb::QueryVectorsResponse { vectors })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn batch_of_rejects_a_ragged_vector() { + let vectors = vec![ + pb::Vector { + key: "a".to_string(), + data: vec![1.0, 2.0], + metadata: None, + }, + pb::Vector { + key: "b".to_string(), + data: vec![1.0; 4], + metadata: None, + }, + ]; + let schema = schema_for(4, pb::DistanceMetric::L2, &[]); + assert!(batch_of(&vectors, 4, schema).is_err()); + } + + #[test] + fn batch_of_rejects_a_repeated_key() { + let schema = schema_for(4, pb::DistanceMetric::L2, &[]); + let vectors = vec![ + pb::Vector { + key: "a".to_string(), + data: vec![1.0; 4], + metadata: None, + }, + pb::Vector { + key: "a".to_string(), + data: vec![2.0; 4], + metadata: None, + }, + ]; + assert!(batch_of(&vectors, 4, schema).is_err()); + } +} diff --git a/seaweed-worker/crates/lance/src/catalog/namespace.rs b/seaweed-worker/crates/lance/src/catalog/namespace.rs index 61a0e099f3..6aca480495 100644 --- a/seaweed-worker/crates/lance/src/catalog/namespace.rs +++ b/seaweed-worker/crates/lance/src/catalog/namespace.rs @@ -127,6 +127,30 @@ impl NamespaceClient { Ok(response.namespaces) } + /// Declare a table that does not exist yet; the catalog allocates the + /// location when none is given and returns the table description. + pub async fn declare_table( + &self, + id: &[String], + location: Option<&str>, + ) -> Result { + let url = format!("{}/v1/table/{}/declare", self.base_url, encode_id(id)); + let mut body = serde_json::json!({ "id": id, "vend_credentials": true }); + if let Some(location) = location { + body["location"] = location.into(); + } + let description: TableDescription = self + .authed(self.http.post(&url)) + .json(&body) + .send() + .await + .context("declare table")? + .error_for_status()? + .json() + .await?; + Ok(description) + } + /// Resolve a table to a location and the credentials to reach it. The /// credentials expire, so a long compaction re-describes rather than /// carrying one set for the whole job.