mirror of
https://github.com/seaweedfs/seaweedfs.git
synced 2026-10-27 23:40:41 +01:00
vector: weed-vector-server, the VectorService gRPC (P1) (#12102)
* vector: VectorService protobuf contract The shared RPC boundary between the Go S3 gateway and the Rust vector engine from the vector design doc: CreateIndex/DropIndex, Put/Get/ List/Delete/Query vectors. A Table reference carries the Lance table id and an optional pinned version, matching the snapshot discipline queries are required to keep. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: generated Go stubs for vector.proto protoc output via weed/pb/Makefile. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: carry return_data through QueryVectors S3 Vectors QueryVectors honours returnData, but the contract had no way to ask for or carry vector data in a query response, so the gateway would need a second GetVectors round-trip per query. Add return_data to the request and data to QueryResult now, while the contract is still new. * vector: pin metadata numbers to double precision in the contract Metadata arrives as JSON, where numbers carry double precision; say so on the field so neither side promises integer fidelity past 2^53. * vector: declare_table on the lance namespace client CreateIndex needs a way to ask the catalog for a new table's location; expose POST /v1/table/{id}/declare. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: cargo fmt coverage.rs Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: weed-vector-server, the VectorService gRPC A tonic server over weed-vector-core: CreateIndex declares the table through the namespace and writes an empty {key, vector, metadata} dataset; PutVectors upserts by key (delete + append until merge_insert is exercised); Get/List/Delete/Query route through the engine's resolved table, honoring pinned versions and the prefilter default. Metadata travels as a Utf8 JSON column, proto Struct on the wire. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: stamp the metric in the schema, tighten put validation CreateIndex now bounds dimension at the S3 Vectors 4096 maximum and records the distance metric in the vector field's Arrow metadata, so a later reader recovers the index's configured metric from the table itself. batch_of checks every vector's data length individually before flattening — equal totals could still let a short vector borrow floats from its neighbor — and writes against the dataset's own schema. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: upsert vectors in a single merge commit PutVectors ran delete-then-append: a failed append left the old keys deleted, and two concurrent puts could both append a replacement. Use lance's merge_insert on the key column — one commit, no delete window, conflicts fail the request instead of duplicating keys. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: read the index's metric; fetch metadata in one take QueryVectors hardcoded L2 — it now reads the metric the index was created with from the vector field's metadata. return_metadata fetched each hit's row one at a time; collect the row ids into a single take_rows call and join metadata back on the hit key. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: honor return_data in query results QueryResult grew a data field but query_vectors never populated it. Fetch the vector column in the same batched take_rows as metadata and return each hit's stored embedding. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: reject a malformed query as invalid_argument A zero top_k silently became one result and a dim-mismatched query vector surfaced as an internal scan error; reject both before querying. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: reject a malformed list next_token A token that failed to parse silently restarted the listing from offset zero; surface it as invalid_argument instead. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: answer an empty put without a commit and cover ragged input An empty vectors list would still run a merge commit; return the current version instead. Pin the ragged-vector check with a test. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: keep non_filterable_metadata_keys on the table schema create_index dropped the declared filter policy. Persist it as JSON on the metadata field so the index carries which keys a filter may touch; enforcing the list belongs to the gateway that translates the request filter into Lance SQL. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: reject duplicate keys in one put A repeated key for a new row committed both copies, and for an existing key the merge failure surfaced as internal. Reject duplicates up front as invalid_argument. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: report open failures as not_found or unavailable A catalog 404 or a missing dataset/version is not a client mistake, and a catalog outage is retryable: mapping every open error to invalid_argument made both indistinguishable from a bad request. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: page list_vectors when max_results is unset A proto3 client that omits max_results sends 0, which used to limit the scan to zero rows and return no token. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: drop stray empty service.rs.tmp * vector: rename datas map to data_map codespell flags datas as a misspelling of data; use a distinct map name. --------- Co-authored-by: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
1 parent
d771a91571
commit
87f15e4d6c
9 files changed
+783
-6
No files matched your search
Generated
+28
@@ -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"
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
# milestone.
|
||||
[workspace]
|
||||
resolver = "2"
|
||||
members = ["crates/core"]
|
||||
members = ["crates/core", "crates/server"]
|
||||
|
||||
[workspace.package]
|
||||
version = "0.1.0"
|
||||
|
||||
@@ -37,11 +37,7 @@ pub async fn coverage(dataset: &Dataset, vector_column: &str) -> Result<Vec<Inde
|
||||
.or_default()
|
||||
.push(index.fragment_bitmap.clone());
|
||||
}
|
||||
let all: RoaringBitmap = dataset
|
||||
.fragments()
|
||||
.iter()
|
||||
.map(|f| f.id as u32)
|
||||
.collect();
|
||||
let all: RoaringBitmap = dataset.fragments().iter().map(|f| f.id as u32).collect();
|
||||
let mut out = Vec::new();
|
||||
for (name, bitmaps) in by_name {
|
||||
let coverage_known = bitmaps.iter().all(|b| b.is_some());
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
[package]
|
||||
name = "weed-vector-server"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
rust-version.workspace = true
|
||||
description = "VectorService gRPC server over SeaweedFS Lance table buckets"
|
||||
|
||||
[lib]
|
||||
name = "weed_vector_server"
|
||||
|
||||
[[bin]]
|
||||
name = "weed-vector-server"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
weed-vector-core = { path = "../core" }
|
||||
weed-lance-worker = { path = "../../../seaweed-worker/crates/lance" }
|
||||
lance = { version = "10", default-features = false, features = ["aws"] }
|
||||
lance-linalg = "10"
|
||||
arrow-array = "58"
|
||||
arrow-cast = "58"
|
||||
arrow-schema = "58"
|
||||
prost = "0.14"
|
||||
prost-types = "0.14"
|
||||
reqwest = { version = "0.12", default-features = false }
|
||||
tonic = { version = "0.14", features = ["tls-aws-lc"] }
|
||||
tonic-prost = "0.14"
|
||||
anyhow.workspace = true
|
||||
clap = { version = "4", features = ["derive", "env"] }
|
||||
futures = "0.3"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
tokio.workspace = true
|
||||
tracing.workspace = true
|
||||
tracing-subscriber.workspace = true
|
||||
|
||||
[build-dependencies]
|
||||
tonic-prost-build = "0.14"
|
||||
protoc-bin-vendored = "3"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,16 @@
|
||||
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
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(())
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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<anyhow::Error>) -> Status {
|
||||
Status::internal(format!("{:#}", err.into()))
|
||||
}
|
||||
|
||||
fn invalid(err: impl Into<anyhow::Error>) -> Status {
|
||||
Status::invalid_argument(format!("{:#}", err.into()))
|
||||
}
|
||||
|
||||
fn open_status(err: anyhow::Error) -> Status {
|
||||
let missing = err.chain().any(|cause| {
|
||||
cause
|
||||
.downcast_ref::<reqwest::Error>()
|
||||
.is_some_and(|e| e.status() == Some(reqwest::StatusCode::NOT_FOUND))
|
||||
|| cause.downcast_ref::<lance::Error>().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<Schema> {
|
||||
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<Schema>) -> Result<RecordBatch> {
|
||||
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::<Vec<_>>());
|
||||
// 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::<Vec<_>>();
|
||||
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::<Result<Vec<_>, _>>()?;
|
||||
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<String, serde_json::Value>) -> Struct {
|
||||
Struct {
|
||||
fields: map
|
||||
.into_iter()
|
||||
.map(|(k, v)| (k, json_to_value(v)))
|
||||
.collect::<BTreeMap<_, _>>(),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_struct(s: &str) -> Result<Struct> {
|
||||
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::<Vec<_>>()
|
||||
.join(", ")
|
||||
)
|
||||
}
|
||||
|
||||
fn rows_as_vectors(
|
||||
batches: &[RecordBatch],
|
||||
return_data: bool,
|
||||
return_metadata: bool,
|
||||
) -> Vec<pb::Vector> {
|
||||
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::<StringArray>())
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let vectors = batch
|
||||
.column_by_name(VECTOR_COLUMN)
|
||||
.and_then(|c| c.as_any().downcast_ref::<FixedSizeListArray>());
|
||||
let metas = batch
|
||||
.column_by_name(METADATA_COLUMN)
|
||||
.and_then(|c| c.as_any().downcast_ref::<StringArray>());
|
||||
for i in 0..batch.num_rows() {
|
||||
let data = match (return_data, vectors) {
|
||||
(true, Some(list)) => list
|
||||
.value(i)
|
||||
.as_any()
|
||||
.downcast_ref::<Float32Array>()
|
||||
.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<Dataset> {
|
||||
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<pb::CreateIndexRequest>,
|
||||
) -> Result<Response<pb::CreateIndexResponse>, 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<pb::DropIndexRequest>,
|
||||
) -> Result<Response<pb::DropIndexResponse>, 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<pb::PutVectorsRequest>,
|
||||
) -> Result<Response<pb::PutVectorsResponse>, 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::<arrow_schema::Schema>::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<pb::GetVectorsRequest>,
|
||||
) -> Result<Response<pb::GetVectorsResponse>, 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<pb::DeleteVectorsRequest>,
|
||||
) -> Result<Response<pb::DeleteVectorsResponse>, 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<pb::ListVectorsRequest>,
|
||||
) -> Result<Response<pb::ListVectorsResponse>, 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<RecordBatch> = scanner
|
||||
.try_into_stream()
|
||||
.await
|
||||
.map_err(internal)?
|
||||
.try_collect()
|
||||
.await
|
||||
.map_err(internal)?;
|
||||
let rows = batches.iter().map(|b| b.num_rows()).sum::<usize>();
|
||||
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<pb::QueryVectorsRequest>,
|
||||
) -> Result<Response<pb::QueryVectorsResponse>, 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<String, Option<Struct>> =
|
||||
std::collections::HashMap::new();
|
||||
let mut data_map: std::collections::HashMap<String, Vec<f32>> =
|
||||
std::collections::HashMap::new();
|
||||
if (req.return_metadata || req.return_data) && !hits.is_empty() {
|
||||
let row_ids: Vec<u64> = 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::<StringArray>());
|
||||
let meta_col = batch
|
||||
.column_by_name(METADATA_COLUMN)
|
||||
.and_then(|c| c.as_any().downcast_ref::<StringArray>());
|
||||
let data_col = batch
|
||||
.column_by_name(VECTOR_COLUMN)
|
||||
.and_then(|c| c.as_any().downcast_ref::<FixedSizeListArray>());
|
||||
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::<Float32Array>()
|
||||
.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());
|
||||
}
|
||||
}
|
||||
@@ -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<TableDescription> {
|
||||
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.
|
||||
|
||||
Reference in new issue
Block a user