mirror of
https://github.com/seaweedfs/seaweedfs.git
synced 2026-10-26 23:10:44 +01:00
vector: S3 Vectors API gateway (P1) (#12103)
* vector: S3 Vectors API gateway over VectorService weed/s3api/s3vectors serves the AWS S3 Vectors JSON API: bucket CRUD maps to S3 Tables table buckets, index CRUD to Lance table records plus the VectorService CreateIndex/DropIndex RPCs, and the data plane (Put/Get/List/Delete/QueryVectors) calls the shared service. The AWS filter document translates to a Lance SQL predicate over the metadata JSON column; unsupported operators fail loudly rather than overmatch. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: -port.vectors and -vectors.grpc on weed s3 Mirror the Lance server startup: the S3 Vectors API binds its own port and dials the weed-vector-server gRPC address; both disabled by default. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: resolve index ARNs, authorize index ops, fix catalog fields - vector buckets declare FormatLance so a new bucket accepts index tables instead of defaulting to ICEBERG. - ListIndexes decodes the `tables` response field; `tableSummaries` hid every index. - resolveIndexARNs honors indexArn alone; handlers no longer treat indexName as the only carrier. - data-plane RPCs fetch the index record first, so callers authorize against the index resource and a missing index 404s before touching the engine; QueryVectors reports the index's configured metric. - IndexArn added to the index-level request types. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: serve the vectors port over TLS when configured When -key.file/-cert.file are set, the vectors listener mirrors the main API and serves TLS rather than a second plaintext endpoint. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: harden the s3vectors gateway handlers - $ne no longer shadows the operator: exclusion filters emit <> instead of equality. - a null element in a PutVectors list fails the request instead of panicking. - requests are capped at 32 MiB; distances return via *float32 so an exact match still reports 0 when returnDistance is set. - CreateIndex declares the `vector` namespace first and validates dimension (1..=4096) before the catalog write; a failed VectorService.CreateIndex rolls the table record back. - non-filterable metadata keys persist as JSON rather than a comma-joined string, and index records carry their createdAt. - ListVectors honors segmentIndex/segmentCount by slicing the full listing in the gateway. - catalog errors map to AccessDeniedException/ConflictException instead of a blanket ValidationException, and malformed ARNs 400 instead of slicing. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: /{Operation} routes and vectors-port TLS, conn cleanup Path-dispatched clients (POST /QueryVectors) reach the same handlers as X-Amz-Target calls. A failed listener bind now closes the vector gRPC connection instead of leaving it open with the endpoint disabled. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
1 parent
74abcb630e
commit
dcc94f818d
4 files changed
+1279
No files matched your search
@@ -22,17 +22,20 @@ import (
|
||||
"github.com/seaweedfs/seaweedfs/weed/pb/filer_pb"
|
||||
"github.com/seaweedfs/seaweedfs/weed/pb/s3_lifecycle_pb"
|
||||
"github.com/seaweedfs/seaweedfs/weed/pb/s3_pb"
|
||||
"github.com/seaweedfs/seaweedfs/weed/pb/vector_pb"
|
||||
"github.com/seaweedfs/seaweedfs/weed/s3api"
|
||||
"github.com/seaweedfs/seaweedfs/weed/s3api/iceberg"
|
||||
"github.com/seaweedfs/seaweedfs/weed/s3api/lance"
|
||||
"github.com/seaweedfs/seaweedfs/weed/s3api/s3err"
|
||||
"github.com/seaweedfs/seaweedfs/weed/s3api/s3tables"
|
||||
"github.com/seaweedfs/seaweedfs/weed/s3api/s3vectors"
|
||||
"github.com/seaweedfs/seaweedfs/weed/security"
|
||||
stats_collect "github.com/seaweedfs/seaweedfs/weed/stats"
|
||||
"github.com/seaweedfs/seaweedfs/weed/util"
|
||||
"github.com/seaweedfs/seaweedfs/weed/util/grace"
|
||||
util_http "github.com/seaweedfs/seaweedfs/weed/util/http"
|
||||
"github.com/seaweedfs/seaweedfs/weed/util/version"
|
||||
"google.golang.org/grpc"
|
||||
)
|
||||
|
||||
var (
|
||||
@@ -52,6 +55,8 @@ type S3Options struct {
|
||||
portGrpc *int
|
||||
portIceberg *int
|
||||
portLance *int
|
||||
portVectors *int
|
||||
vectorsGrpc *string
|
||||
icebergCredentialRole *string
|
||||
icebergCredentialDuration *int
|
||||
config *string
|
||||
@@ -110,6 +115,8 @@ func init() {
|
||||
s3StandaloneOptions.portGrpc = cmdS3.Flag.Int("port.grpc", 0, "s3 server grpc listen port")
|
||||
s3StandaloneOptions.portIceberg = cmdS3.Flag.Int("port.iceberg", 8181, "Iceberg REST Catalog server listen port (0 to disable)")
|
||||
s3StandaloneOptions.portLance = cmdS3.Flag.Int("port.lance", 9101, "Lance Namespace server listen port (0 to disable); credential vending uses -iceberg.credentialRole")
|
||||
s3StandaloneOptions.portVectors = cmdS3.Flag.Int("port.vectors", 0, "S3 Vectors API server listen port (0 to disable); requires -vectors.grpc")
|
||||
s3StandaloneOptions.vectorsGrpc = cmdS3.Flag.String("vectors.grpc", "", "weed-vector-server gRPC address for the S3 Vectors API")
|
||||
s3StandaloneOptions.icebergCredentialRole = cmdS3.Flag.String("iceberg.credentialRole", "", "IAM role ARN the Iceberg catalog assumes to vend table-scoped credentials (empty disables vending)")
|
||||
s3StandaloneOptions.icebergCredentialDuration = cmdS3.Flag.Int("iceberg.credentialDurationSeconds", 3600, "lifetime of credentials vended by the Iceberg catalog")
|
||||
s3StandaloneOptions.domainName = cmdS3.Flag.String("domainName", "", "suffix of the host name in comma separated list, {bucket}.{domainName}")
|
||||
@@ -433,6 +440,13 @@ func (s3opt *S3Options) startS3Server() bool {
|
||||
}
|
||||
}
|
||||
|
||||
// Start S3 Vectors API server if enabled
|
||||
if s3opt.portVectors != nil && *s3opt.portVectors > 0 {
|
||||
if httpS := s3opt.startVectorsServer(s3ApiServer, grpcDialOption, listenOnLocalhost); httpS != nil {
|
||||
httpServers = append(httpServers, httpS)
|
||||
}
|
||||
}
|
||||
|
||||
if runtime.GOOS != "windows" {
|
||||
localSocket := *s3opt.localSocket
|
||||
explicit := localSocket != ""
|
||||
@@ -692,6 +706,86 @@ func (s3opt *S3Options) startLanceServer(s3ApiServer *s3api.S3ApiServer, listenO
|
||||
return httpS
|
||||
}
|
||||
|
||||
// startVectorsServer starts the S3 Vectors API server on a separate port. It
|
||||
// fans index and data calls out to the shared VectorService over gRPC; with no
|
||||
// -vectors.grpc the API is not served.
|
||||
func (s3opt *S3Options) startVectorsServer(s3ApiServer *s3api.S3ApiServer, grpcDialOption grpc.DialOption, listenOnLocalhost bool) *http.Server {
|
||||
if *s3opt.vectorsGrpc == "" {
|
||||
glog.Warningf("port.vectors set without -vectors.grpc; S3 Vectors API disabled")
|
||||
return nil
|
||||
}
|
||||
conn, err := pb.GrpcDial(context.Background(), *s3opt.vectorsGrpc, false, grpcDialOption)
|
||||
if err != nil {
|
||||
glog.Errorf("S3 Vectors: dial %s: %v", *s3opt.vectorsGrpc, err)
|
||||
return nil
|
||||
}
|
||||
vectorsRouter := mux.NewRouter().SkipClean(true)
|
||||
|
||||
vectorsServer := s3vectors.NewServer(s3ApiServer, s3ApiServer, vector_pb.NewVectorServiceClient(conn))
|
||||
vectorsServer.SetRegion(s3tables.DefaultRegion)
|
||||
vectorsServer.SetAccountID(s3tables.DefaultAccountID)
|
||||
vectorsServer.RegisterRoutes(vectorsRouter)
|
||||
|
||||
listenAddress := fmt.Sprintf("%s:%d", *s3opt.bindIp, *s3opt.portVectors)
|
||||
vectorsListener, vectorsLocalListener, err := util.NewIpAndLocalListeners(
|
||||
*s3opt.bindIp, *s3opt.portVectors, time.Duration(*s3opt.idleTimeout)*time.Second, listenOnLocalhost)
|
||||
if err != nil {
|
||||
glog.Errorf("S3 Vectors listener on %s error: %v", listenAddress, err)
|
||||
conn.Close()
|
||||
return nil
|
||||
}
|
||||
|
||||
// Mirror the main API: when TLS is configured the vectors port serves
|
||||
// HTTPS, not a second plaintext endpoint.
|
||||
if *s3opt.tlsPrivateKey != "" {
|
||||
getCert, _, err := security.NewReloadingServerCertificate(*s3opt.tlsCertificate, *s3opt.tlsPrivateKey)
|
||||
if err != nil {
|
||||
glog.Errorf("S3 Vectors: load HTTPS certificate: %v", err)
|
||||
return nil
|
||||
}
|
||||
clientAuth := tls.NoClientCert
|
||||
if *s3opt.tlsVerifyClientCert {
|
||||
clientAuth = tls.RequireAndVerifyClientCert
|
||||
}
|
||||
tlsConfig := &tls.Config{GetCertificate: getCert, ClientAuth: clientAuth}
|
||||
if err := security.FixTlsConfig(util.GetViper(), tlsConfig); err != nil {
|
||||
glog.Errorf("S3 Vectors: tls config: %v", err)
|
||||
return nil
|
||||
}
|
||||
httpS := newHttpServer(vectorsRouter, tlsConfig)
|
||||
if vectorsLocalListener != nil {
|
||||
go func() {
|
||||
if err := httpS.ServeTLS(vectorsLocalListener, "", ""); err != nil && err != http.ErrServerClosed {
|
||||
glog.V(0).Infof("S3 Vectors localhost listener error: %v", err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
go func() {
|
||||
if err := httpS.ServeTLS(vectorsListener, "", ""); err != nil && err != http.ErrServerClosed {
|
||||
glog.Errorf("S3 Vectors API Server Fail to serve: %v", err)
|
||||
}
|
||||
}()
|
||||
return httpS
|
||||
}
|
||||
|
||||
glog.V(0).Infof("Start S3 Vectors API Server at http://%s", listenAddress)
|
||||
|
||||
httpS := newHttpServer(vectorsRouter, nil)
|
||||
if vectorsLocalListener != nil {
|
||||
go func() {
|
||||
if err := httpS.Serve(vectorsLocalListener); err != nil && err != http.ErrServerClosed {
|
||||
glog.V(0).Infof("S3 Vectors localhost listener error: %v", err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
go func() {
|
||||
if err := httpS.Serve(vectorsListener); err != nil && err != http.ErrServerClosed {
|
||||
glog.Errorf("S3 Vectors API Server Fail to serve: %v", err)
|
||||
}
|
||||
}()
|
||||
return httpS
|
||||
}
|
||||
|
||||
// deriveLanceStorageEndpoint picks the endpoint the Lance namespace puts in
|
||||
// storage_options. It falls back to the advertised -ip where the Iceberg
|
||||
// derivation gives up, because the two clients are not in the same position: a
|
||||
|
||||
@@ -0,0 +1,840 @@
|
||||
package s3vectors
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/seaweedfs/seaweedfs/weed/pb/vector_pb"
|
||||
"github.com/seaweedfs/seaweedfs/weed/s3api/s3tables"
|
||||
"google.golang.org/protobuf/types/known/structpb"
|
||||
)
|
||||
|
||||
var handlers = map[string]func(*Server, http.ResponseWriter, *http.Request){
|
||||
"CreateVectorBucket": (*Server).handleCreateVectorBucket,
|
||||
"GetVectorBucket": (*Server).handleGetVectorBucket,
|
||||
"ListVectorBuckets": (*Server).handleListVectorBuckets,
|
||||
"DeleteVectorBucket": (*Server).handleDeleteVectorBucket,
|
||||
"CreateIndex": (*Server).handleCreateIndex,
|
||||
"GetIndex": (*Server).handleGetIndex,
|
||||
"ListIndexes": (*Server).handleListIndexes,
|
||||
"DeleteIndex": (*Server).handleDeleteIndex,
|
||||
"PutVectors": (*Server).handlePutVectors,
|
||||
"GetVectors": (*Server).handleGetVectors,
|
||||
"ListVectors": (*Server).handleListVectors,
|
||||
"DeleteVectors": (*Server).handleDeleteVectors,
|
||||
"QueryVectors": (*Server).handleQueryVectors,
|
||||
}
|
||||
|
||||
// lance table ids put indexes under a "vector" namespace inside their bucket.
|
||||
const indexNamespace = "vector"
|
||||
|
||||
// Vector bucket records are s3tables table buckets; the S3 Vectors ARN form is
|
||||
// arn:aws:s3vectors:{region}:{account}:bucket/{name}[/index/{name}].
|
||||
func (s *Server) vectorBucketArn(bucket string) string {
|
||||
return fmt.Sprintf("arn:aws:s3vectors:%s:%s:bucket/%s", s.region, s.accountID, bucket)
|
||||
}
|
||||
|
||||
func (s *Server) indexArn(bucket, index string) string {
|
||||
return fmt.Sprintf("%s/index/%s", s.vectorBucketArn(bucket), index)
|
||||
}
|
||||
|
||||
// s3tablesArn renders the same bucket in the table-bucket ARN form the S3
|
||||
// Tables catalog keys records by.
|
||||
func (s *Server) s3tablesArn(bucket string) string {
|
||||
return fmt.Sprintf("arn:aws:s3tables:%s:%s:bucket/%s", s.region, s.accountID, bucket)
|
||||
}
|
||||
|
||||
func arnParts(v *vectorBucketARNs, name, arn string) error {
|
||||
if arn == "" {
|
||||
v.bucket = name
|
||||
return nil
|
||||
}
|
||||
if !strings.HasPrefix(arn, "arn:") {
|
||||
v.bucket = arn
|
||||
return nil
|
||||
}
|
||||
i := strings.Index(arn, "bucket/")
|
||||
if i < 0 {
|
||||
return fmt.Errorf("malformed ARN %q", arn)
|
||||
}
|
||||
rest := arn[i+len("bucket/"):]
|
||||
parts := strings.Split(rest, "/index/")
|
||||
v.bucket = parts[0]
|
||||
if len(parts) == 2 {
|
||||
v.index = parts[1]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// resolveIndexARNs fills bucket and index from any combination of plain names,
|
||||
// a bucket ARN, and a full index ARN — callers may supply only an indexArn.
|
||||
func resolveIndexARNs(bucketName, bucketArn, indexName, indexArn string) (vectorBucketARNs, error) {
|
||||
var v vectorBucketARNs
|
||||
if strings.HasPrefix(indexArn, "arn:") {
|
||||
if err := arnParts(&v, "", indexArn); err != nil {
|
||||
return v, err
|
||||
}
|
||||
} else if indexArn != "" {
|
||||
return v, fmt.Errorf("indexArn must be an ARN, got %q", indexArn)
|
||||
}
|
||||
if strings.HasPrefix(bucketArn, "arn:") {
|
||||
if err := arnParts(&v, "", bucketArn); err != nil {
|
||||
return v, err
|
||||
}
|
||||
}
|
||||
if v.bucket == "" {
|
||||
if bucketArn != "" {
|
||||
v.bucket = bucketArn
|
||||
} else {
|
||||
v.bucket = bucketName
|
||||
}
|
||||
}
|
||||
if v.index == "" {
|
||||
v.index = indexName
|
||||
}
|
||||
if v.bucket == "" || v.index == "" {
|
||||
return v, fmt.Errorf("vectorBucketName and indexName (or indexArn) are required")
|
||||
}
|
||||
return v, nil
|
||||
}
|
||||
|
||||
// indexRecord fetches the index's catalog record; the call also authorizes
|
||||
// the request identity against the index resource.
|
||||
func (s *Server) indexRecord(ctx context.Context, r *http.Request, bucket, index string) (map[string]interface{}, error) {
|
||||
var resp map[string]interface{}
|
||||
err := s.catalog(ctx, "GetTable", &s3tables.GetTableRequest{
|
||||
TableBucketARN: s.s3tablesArn(bucket),
|
||||
Namespace: []string{indexNamespace},
|
||||
Name: index,
|
||||
}, &resp, r)
|
||||
return resp, err
|
||||
}
|
||||
|
||||
func (s *Server) catalog(ctx context.Context, operation string, req, resp interface{}, r *http.Request) error {
|
||||
return s.tables.Execute(ctx, s.filerClient, operation, req, resp, identityOf(r))
|
||||
}
|
||||
|
||||
func isNotFound(err error) bool {
|
||||
return err != nil && (strings.Contains(err.Error(), "NoSuchTableBucket") ||
|
||||
strings.Contains(err.Error(), "NoSuchTable") ||
|
||||
strings.Contains(err.Error(), "NotFoundException") ||
|
||||
strings.Contains(err.Error(), "NoSuchBucket"))
|
||||
}
|
||||
|
||||
func catalogErr(w http.ResponseWriter, err error) {
|
||||
msg := err.Error()
|
||||
switch {
|
||||
case isNotFound(err):
|
||||
writeError(w, http.StatusNotFound, "NotFoundException", msg)
|
||||
case strings.Contains(msg, "AccessDenied") || strings.Contains(msg, "PermissionDenied"):
|
||||
writeError(w, http.StatusForbidden, "AccessDeniedException", "access denied")
|
||||
case strings.Contains(msg, "AlreadyExists") || strings.Contains(msg, "Conflict"):
|
||||
writeError(w, http.StatusConflict, "ConflictException", msg)
|
||||
default:
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", msg)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) handleCreateVectorBucket(w http.ResponseWriter, r *http.Request) {
|
||||
var req CreateVectorBucketRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
if req.VectorBucketName == "" {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", "vectorBucketName is required")
|
||||
return
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
err := s.catalog(r.Context(), "CreateTableBucket",
|
||||
&s3tables.CreateTableBucketRequest{Name: req.VectorBucketName, Format: s3tables.FormatLance}, &resp, r)
|
||||
if err != nil {
|
||||
catalogErr(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, &CreateVectorBucketResponse{})
|
||||
}
|
||||
|
||||
func (s *Server) handleGetVectorBucket(w http.ResponseWriter, r *http.Request) {
|
||||
var req GetVectorBucketRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
var arns vectorBucketARNs
|
||||
if err := arnParts(&arns, req.VectorBucketName, req.VectorBucketArn); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
err := s.catalog(r.Context(), "GetTableBucket",
|
||||
&s3tables.GetTableBucketRequest{TableBucketARN: s.s3tablesArn(arns.bucket)}, &resp, r)
|
||||
if err != nil {
|
||||
catalogErr(w, err)
|
||||
return
|
||||
}
|
||||
created, _ := time.Parse(time.RFC3339, fmt.Sprint(resp["createdAt"]))
|
||||
writeJSON(w, http.StatusOK, &GetVectorBucketResponse{VectorBucket: &VectorBucket{
|
||||
VectorBucketArn: s.vectorBucketArn(arns.bucket),
|
||||
VectorBucketName: arns.bucket,
|
||||
CreationTime: created,
|
||||
}})
|
||||
}
|
||||
|
||||
func (s *Server) handleListVectorBuckets(w http.ResponseWriter, r *http.Request) {
|
||||
var req ListVectorBucketsRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
var resp struct {
|
||||
TableBuckets []struct {
|
||||
Arn string `json:"arn"`
|
||||
Name string `json:"name"`
|
||||
CreatedAt string `json:"createdAt"`
|
||||
} `json:"tableBuckets"`
|
||||
ContinuationToken string `json:"continuationToken"`
|
||||
}
|
||||
err := s.catalog(r.Context(), "ListTableBuckets",
|
||||
&s3tables.ListTableBucketsRequest{
|
||||
MaxBuckets: req.MaxResults, ContinuationToken: req.NextToken, Prefix: req.Prefix,
|
||||
}, &resp, r)
|
||||
if err != nil {
|
||||
catalogErr(w, err)
|
||||
return
|
||||
}
|
||||
out := &ListVectorBucketsResponse{NextToken: resp.ContinuationToken}
|
||||
for _, b := range resp.TableBuckets {
|
||||
created, _ := time.Parse(time.RFC3339, b.CreatedAt)
|
||||
out.VectorBuckets = append(out.VectorBuckets, &VectorBucket{
|
||||
VectorBucketArn: s.vectorBucketArn(b.Name),
|
||||
VectorBucketName: b.Name,
|
||||
CreationTime: created,
|
||||
})
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
func (s *Server) handleDeleteVectorBucket(w http.ResponseWriter, r *http.Request) {
|
||||
var req DeleteVectorBucketRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
var arns vectorBucketARNs
|
||||
if err := arnParts(&arns, req.VectorBucketName, req.VectorBucketArn); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
err := s.catalog(r.Context(), "DeleteTableBucket",
|
||||
&s3tables.DeleteTableBucketRequest{TableBucketARN: s.s3tablesArn(arns.bucket)}, &resp, r)
|
||||
if err != nil {
|
||||
catalogErr(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, &DeleteVectorBucketResponse{})
|
||||
}
|
||||
|
||||
func distanceMetric(metric string) (vector_pb.DistanceMetric, bool) {
|
||||
switch strings.ToLower(metric) {
|
||||
case "", "euclidean":
|
||||
return vector_pb.DistanceMetric_L2, true
|
||||
case "cosine":
|
||||
return vector_pb.DistanceMetric_COSINE, true
|
||||
case "dotproduct":
|
||||
return vector_pb.DistanceMetric_DOT_PRODUCT, true
|
||||
}
|
||||
return vector_pb.DistanceMetric_L2, false
|
||||
}
|
||||
|
||||
func tableID(bucket, index string) *vector_pb.Table {
|
||||
return &vector_pb.Table{Id: []string{bucket, indexNamespace, index}}
|
||||
}
|
||||
|
||||
func (s *Server) handleCreateIndex(w http.ResponseWriter, r *http.Request) {
|
||||
var req CreateIndexRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
if strings.ToLower(req.DataType) != "float32" {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", "only float32 vectors are supported")
|
||||
return
|
||||
}
|
||||
metric, ok := distanceMetric(req.DistanceMetric)
|
||||
if !ok {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", "unknown distanceMetric "+req.DistanceMetric)
|
||||
return
|
||||
}
|
||||
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
if req.Dimension <= 0 || req.Dimension > 4096 {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", "dimension must be 1..=4096")
|
||||
return
|
||||
}
|
||||
metadata := map[string]string{
|
||||
"dimension": fmt.Sprint(req.Dimension),
|
||||
"distanceMetric": req.DistanceMetric,
|
||||
}
|
||||
if keys, err := json.Marshal(req.MetadataConfiguration.NonFilterableMetadataKeys); req.MetadataConfiguration != nil && err == nil {
|
||||
metadata["nonFilterableMetadataKeys"] = string(keys)
|
||||
}
|
||||
metadataJSON, _ := json.Marshal(metadata)
|
||||
var nsResp map[string]interface{}
|
||||
// The index namespace is implicit on AWS; declare it, tolerating an
|
||||
// already-existing namespace.
|
||||
if err := s.catalog(r.Context(), "CreateNamespace", &s3tables.CreateNamespaceRequest{
|
||||
TableBucketARN: s.s3tablesArn(arns.bucket),
|
||||
Namespace: []string{indexNamespace},
|
||||
}, &nsResp, r); err != nil && !strings.Contains(err.Error(), "AlreadyExists") && !strings.Contains(err.Error(), "Conflict") {
|
||||
catalogErr(w, err)
|
||||
return
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
err = s.catalog(r.Context(), "CreateTable", &s3tables.CreateTableRequest{
|
||||
TableBucketARN: s.s3tablesArn(arns.bucket),
|
||||
Namespace: []string{indexNamespace},
|
||||
Name: arns.index,
|
||||
Format: s3tables.FormatLance,
|
||||
Metadata: &s3tables.TableMetadata{FullMetadata: metadataJSON},
|
||||
}, &resp, r)
|
||||
if err != nil {
|
||||
catalogErr(w, err)
|
||||
return
|
||||
}
|
||||
var nonFilterable []string
|
||||
if req.MetadataConfiguration != nil {
|
||||
nonFilterable = req.MetadataConfiguration.NonFilterableMetadataKeys
|
||||
}
|
||||
_, err = s.vectors.CreateIndex(r.Context(), &vector_pb.CreateIndexRequest{
|
||||
Table: tableID(arns.bucket, arns.index),
|
||||
Dimension: uint32(req.Dimension),
|
||||
DistanceMetric: metric,
|
||||
NonFilterableMetadataKeys: nonFilterable,
|
||||
})
|
||||
if err != nil {
|
||||
// Roll the index record back so a half-created index does not linger.
|
||||
var delResp map[string]interface{}
|
||||
_ = s.catalog(r.Context(), "DeleteTable", &s3tables.DeleteTableRequest{
|
||||
TableBucketARN: s.s3tablesArn(arns.bucket),
|
||||
Namespace: []string{indexNamespace},
|
||||
Name: arns.index,
|
||||
}, &delResp, r)
|
||||
writeError(w, http.StatusInternalServerError, "InternalFailure", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, &CreateIndexResponse{})
|
||||
}
|
||||
|
||||
func indexFromRecord(s *Server, bucket, name string, m map[string]interface{}) *Index {
|
||||
idx := &Index{
|
||||
IndexArn: s.indexArn(bucket, name),
|
||||
IndexName: name,
|
||||
VectorBucketName: bucket,
|
||||
DataType: "float32",
|
||||
}
|
||||
if created, ok := m["createdAt"].(string); ok {
|
||||
idx.CreationTime, _ = time.Parse(time.RFC3339, created)
|
||||
}
|
||||
if meta, ok := m["metadata"].(map[string]interface{}); ok {
|
||||
if fm, ok := meta["fullMetadata"].(map[string]interface{}); ok {
|
||||
var dim float64
|
||||
fmt.Sscan(fmt.Sprint(fm["dimension"]), &dim)
|
||||
idx.Dimension = int32(dim)
|
||||
idx.DistanceMetric = fmt.Sprint(fm["distanceMetric"])
|
||||
if keys, ok := fm["nonFilterableMetadataKeys"].([]interface{}); ok {
|
||||
var list []string
|
||||
for _, k := range keys {
|
||||
list = append(list, fmt.Sprint(k))
|
||||
}
|
||||
idx.MetadataConfiguration = &MetadataConfiguration{NonFilterableMetadataKeys: list}
|
||||
}
|
||||
}
|
||||
}
|
||||
if idx.DistanceMetric == "" || idx.DistanceMetric == "<nil>" {
|
||||
idx.DistanceMetric = "euclidean"
|
||||
}
|
||||
return idx
|
||||
}
|
||||
|
||||
func (s *Server) handleGetIndex(w http.ResponseWriter, r *http.Request) {
|
||||
var req GetIndexRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
err = s.catalog(r.Context(), "GetTable", &s3tables.GetTableRequest{
|
||||
TableBucketARN: s.s3tablesArn(arns.bucket),
|
||||
Namespace: []string{indexNamespace},
|
||||
Name: arns.index,
|
||||
}, &resp, r)
|
||||
if err != nil {
|
||||
catalogErr(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, &GetIndexResponse{Index: indexFromRecord(s, arns.bucket, req.IndexName, resp)})
|
||||
}
|
||||
|
||||
func (s *Server) handleListIndexes(w http.ResponseWriter, r *http.Request) {
|
||||
var req ListIndexesRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
var arns vectorBucketARNs
|
||||
if err := arnParts(&arns, req.VectorBucketName, req.VectorBucketArn); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
var resp struct {
|
||||
TableSummaries []struct {
|
||||
Name string `json:"name"`
|
||||
CreatedAt string `json:"createdAt"`
|
||||
} `json:"tables"`
|
||||
ContinuationToken string `json:"continuationToken"`
|
||||
}
|
||||
err := s.catalog(r.Context(), "ListTables", &s3tables.ListTablesRequest{
|
||||
TableBucketARN: s.s3tablesArn(arns.bucket),
|
||||
Namespace: []string{indexNamespace},
|
||||
MaxTables: req.MaxResults,
|
||||
ContinuationToken: req.NextToken,
|
||||
Prefix: req.Prefix,
|
||||
}, &resp, r)
|
||||
if err != nil {
|
||||
catalogErr(w, err)
|
||||
return
|
||||
}
|
||||
out := &ListIndexesResponse{NextToken: resp.ContinuationToken}
|
||||
for _, t := range resp.TableSummaries {
|
||||
out.Indexes = append(out.Indexes, indexFromRecord(s, arns.bucket, t.Name, nil))
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
func (s *Server) handleDeleteIndex(w http.ResponseWriter, r *http.Request) {
|
||||
var req DeleteIndexRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
_, _ = s.vectors.DropIndex(r.Context(), &vector_pb.DropIndexRequest{
|
||||
Table: tableID(arns.bucket, arns.index),
|
||||
})
|
||||
var resp map[string]interface{}
|
||||
err = s.catalog(r.Context(), "DeleteTable", &s3tables.DeleteTableRequest{
|
||||
TableBucketARN: s.s3tablesArn(arns.bucket),
|
||||
Namespace: []string{indexNamespace},
|
||||
Name: arns.index,
|
||||
}, &resp, r)
|
||||
if err != nil {
|
||||
catalogErr(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, &DeleteIndexResponse{})
|
||||
}
|
||||
|
||||
func metadataStruct(m map[string]interface{}) (*structpb.Struct, error) {
|
||||
if m == nil {
|
||||
return nil, nil
|
||||
}
|
||||
return structpb.NewStruct(m)
|
||||
}
|
||||
|
||||
func structToMap(s *structpb.Struct) map[string]interface{} {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
return s.AsMap()
|
||||
}
|
||||
|
||||
func (s *Server) handlePutVectors(w http.ResponseWriter, r *http.Request) {
|
||||
var req PutVectorsRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
if _, err := s.indexRecord(r.Context(), r, arns.bucket, arns.index); err != nil {
|
||||
catalogErr(w, err)
|
||||
return
|
||||
}
|
||||
vectors := make([]*vector_pb.Vector, 0, len(req.Vectors))
|
||||
for _, v := range req.Vectors {
|
||||
if v == nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", "vectors contains a null element")
|
||||
return
|
||||
}
|
||||
metadata, err := metadataStruct(v.Metadata)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
var data []float32
|
||||
if v.Data != nil {
|
||||
data = v.Data.Float32
|
||||
}
|
||||
vectors = append(vectors, &vector_pb.Vector{Key: v.Key, Data: data, Metadata: metadata})
|
||||
}
|
||||
_, err = s.vectors.PutVectors(r.Context(), &vector_pb.PutVectorsRequest{
|
||||
Table: tableID(arns.bucket, arns.index),
|
||||
Vectors: vectors,
|
||||
})
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "InternalFailure", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, &PutVectorsResponse{})
|
||||
}
|
||||
|
||||
func toOutputVectors(vectors []*vector_pb.Vector, returnData, returnMetadata bool) []*OutputVector {
|
||||
out := make([]*OutputVector, 0, len(vectors))
|
||||
for _, v := range vectors {
|
||||
o := &OutputVector{Key: v.Key}
|
||||
if returnData {
|
||||
o.Data = &VectorData{Float32: v.Data}
|
||||
}
|
||||
if returnMetadata {
|
||||
o.Metadata = structToMap(v.Metadata)
|
||||
}
|
||||
out = append(out, o)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (s *Server) handleGetVectors(w http.ResponseWriter, r *http.Request) {
|
||||
var req GetVectorsRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
if _, err := s.indexRecord(r.Context(), r, arns.bucket, arns.index); err != nil {
|
||||
catalogErr(w, err)
|
||||
return
|
||||
}
|
||||
resp, err := s.vectors.GetVectors(r.Context(), &vector_pb.GetVectorsRequest{
|
||||
Table: tableID(arns.bucket, arns.index),
|
||||
Keys: req.Keys,
|
||||
ReturnData: req.ReturnData,
|
||||
ReturnMetadata: req.ReturnMetadata,
|
||||
})
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "InternalFailure", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, &GetVectorsResponse{
|
||||
Vectors: toOutputVectors(resp.Vectors, req.ReturnData, req.ReturnMetadata),
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) handleListVectors(w http.ResponseWriter, r *http.Request) {
|
||||
var req ListVectorsRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
if _, err := s.indexRecord(r.Context(), r, arns.bucket, arns.index); err != nil {
|
||||
catalogErr(w, err)
|
||||
return
|
||||
}
|
||||
// AWS exposes a static split: segmentIndex/segmentCount divides the full
|
||||
// listing. The engine service does not know AWS segments, so page through
|
||||
// the whole index here and slice the caller's share — O(index) per call,
|
||||
// acceptable for the MVP.
|
||||
if req.SegmentCount > 0 {
|
||||
var all []*vector_pb.Vector
|
||||
token := ""
|
||||
for {
|
||||
page, err := s.vectors.ListVectors(r.Context(), &vector_pb.ListVectorsRequest{
|
||||
Table: tableID(arns.bucket, arns.index),
|
||||
MaxResults: 1000,
|
||||
NextToken: token,
|
||||
ReturnData: req.ReturnData,
|
||||
ReturnMetadata: req.ReturnMetadata,
|
||||
})
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "InternalFailure", err.Error())
|
||||
return
|
||||
}
|
||||
all = append(all, page.Vectors...)
|
||||
if page.NextToken == "" {
|
||||
break
|
||||
}
|
||||
token = page.NextToken
|
||||
}
|
||||
var out []*vector_pb.Vector
|
||||
for i, v := range all {
|
||||
if i%int(req.SegmentCount) == int(req.SegmentIndex) {
|
||||
out = append(out, v)
|
||||
}
|
||||
}
|
||||
writeJSON(w, http.StatusOK, &ListVectorsResponse{
|
||||
Vectors: toOutputVectors(out, req.ReturnData, req.ReturnMetadata),
|
||||
})
|
||||
return
|
||||
}
|
||||
if req.MaxResults < 0 {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", "maxResults must be non-negative")
|
||||
return
|
||||
}
|
||||
resp, err := s.vectors.ListVectors(r.Context(), &vector_pb.ListVectorsRequest{
|
||||
Table: tableID(arns.bucket, arns.index),
|
||||
MaxResults: uint32(req.MaxResults),
|
||||
NextToken: req.NextToken,
|
||||
ReturnData: req.ReturnData,
|
||||
ReturnMetadata: req.ReturnMetadata,
|
||||
})
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "InternalFailure", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, &ListVectorsResponse{
|
||||
NextToken: resp.NextToken,
|
||||
Vectors: toOutputVectors(resp.Vectors, req.ReturnData, req.ReturnMetadata),
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) handleDeleteVectors(w http.ResponseWriter, r *http.Request) {
|
||||
var req DeleteVectorsRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
if _, err := s.indexRecord(r.Context(), r, arns.bucket, arns.index); err != nil {
|
||||
catalogErr(w, err)
|
||||
return
|
||||
}
|
||||
_, err = s.vectors.DeleteVectors(r.Context(), &vector_pb.DeleteVectorsRequest{
|
||||
Table: tableID(arns.bucket, arns.index),
|
||||
Keys: req.Keys,
|
||||
})
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "InternalFailure", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, &DeleteVectorsResponse{})
|
||||
}
|
||||
|
||||
// translateFilter renders an AWS S3 Vectors filter document as a Lance SQL
|
||||
// predicate over the metadata JSON column. Only the documented operators are
|
||||
// translated; anything else fails loudly instead of silently overmatching.
|
||||
func translateFilter(filter map[string]interface{}) (string, error) {
|
||||
if len(filter) == 0 {
|
||||
return "", nil
|
||||
}
|
||||
return filterExpr(filter)
|
||||
}
|
||||
|
||||
func filterExpr(filter map[string]interface{}) (string, error) {
|
||||
parts := make([]string, 0, len(filter))
|
||||
for key, cond := range filter {
|
||||
if strings.HasPrefix(key, "$") {
|
||||
sub, ok := cond.([]interface{})
|
||||
if !ok {
|
||||
return "", fmt.Errorf("%s expects a list", key)
|
||||
}
|
||||
exprs := make([]string, 0, len(sub))
|
||||
for _, item := range sub {
|
||||
m, ok := item.(map[string]interface{})
|
||||
if !ok {
|
||||
return "", fmt.Errorf("%s expects objects", key)
|
||||
}
|
||||
e, err := filterExpr(m)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
exprs = append(exprs, e)
|
||||
}
|
||||
joiner := " AND "
|
||||
if key == "$or" {
|
||||
joiner = " OR "
|
||||
}
|
||||
if key != "$and" && key != "$or" {
|
||||
return "", fmt.Errorf("unsupported operator %s", key)
|
||||
}
|
||||
parts = append(parts, "("+strings.Join(exprs, joiner)+")")
|
||||
continue
|
||||
}
|
||||
clause, err := fieldClause(key, cond)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
parts = append(parts, clause)
|
||||
}
|
||||
return strings.Join(parts, " AND "), nil
|
||||
}
|
||||
|
||||
func fieldClause(field string, cond interface{}) (string, error) {
|
||||
col := fmt.Sprintf("json_get_str(`metadata`, '%s')", strings.ReplaceAll(field, "'", "''"))
|
||||
switch c := cond.(type) {
|
||||
case map[string]interface{}:
|
||||
parts := []string{}
|
||||
for op, operand := range c {
|
||||
switch op {
|
||||
case "$exists":
|
||||
if operand == true {
|
||||
parts = append(parts, fmt.Sprintf("json_exists(`metadata`, '%s')", strings.ReplaceAll(field, "'", "''")))
|
||||
} else {
|
||||
parts = append(parts, fmt.Sprintf("NOT json_exists(`metadata`, '%s')", strings.ReplaceAll(field, "'", "''")))
|
||||
}
|
||||
case "$eq", "$ne":
|
||||
lit, err := sqlLiteral(operand)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
sym := "="
|
||||
if op == "$ne" {
|
||||
sym = "<>"
|
||||
}
|
||||
parts = append(parts, fmt.Sprintf("%s %s %s", col, sym, lit))
|
||||
case "$in", "$nin":
|
||||
list, ok := operand.([]interface{})
|
||||
if !ok {
|
||||
return "", fmt.Errorf("%s expects a list", op)
|
||||
}
|
||||
lits := make([]string, 0, len(list))
|
||||
for _, item := range list {
|
||||
lit, err := sqlLiteral(item)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
lits = append(lits, lit)
|
||||
}
|
||||
neg := ""
|
||||
if op == "$nin" {
|
||||
neg = "NOT "
|
||||
}
|
||||
parts = append(parts, fmt.Sprintf("%s%s in (%s)", neg, col, strings.Join(lits, ", ")))
|
||||
case "$lt", "$lte", "$gt", "$gte":
|
||||
lit, err := sqlLiteral(operand)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
sym := map[string]string{"$lt": "<", "$lte": "<=", "$gt": ">", "$gte": ">="}[op]
|
||||
numCol := fmt.Sprintf("json_get_float(`metadata`, '%s')", strings.ReplaceAll(field, "'", "''"))
|
||||
parts = append(parts, fmt.Sprintf("%s %s %s", numCol, sym, lit))
|
||||
default:
|
||||
return "", fmt.Errorf("unsupported operator %s", op)
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, " AND "), nil
|
||||
default:
|
||||
lit, err := sqlLiteral(cond)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return fmt.Sprintf("%s = %s", col, lit), nil
|
||||
}
|
||||
}
|
||||
|
||||
func sqlLiteral(v interface{}) (string, error) {
|
||||
switch t := v.(type) {
|
||||
case string:
|
||||
return "'" + strings.ReplaceAll(t, "'", "''") + "'", nil
|
||||
case bool:
|
||||
if t {
|
||||
return "true", nil
|
||||
}
|
||||
return "false", nil
|
||||
case float64:
|
||||
return fmt.Sprintf("%v", t), nil
|
||||
}
|
||||
return "", fmt.Errorf("unsupported filter operand %v", v)
|
||||
}
|
||||
|
||||
func (s *Server) handleQueryVectors(w http.ResponseWriter, r *http.Request) {
|
||||
var req QueryVectorsRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
var queryVector []float32
|
||||
if req.QueryVector != nil {
|
||||
queryVector = req.QueryVector.Float32
|
||||
}
|
||||
if len(queryVector) == 0 {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", "queryVector is required")
|
||||
return
|
||||
}
|
||||
record, err := s.indexRecord(r.Context(), r, arns.bucket, arns.index)
|
||||
if err != nil {
|
||||
catalogErr(w, err)
|
||||
return
|
||||
}
|
||||
metric := indexFromRecord(s, arns.bucket, arns.index, record).DistanceMetric
|
||||
filter, err := translateFilter(req.Filter)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
||||
return
|
||||
}
|
||||
if req.TopK < 0 {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", "topK must be non-negative")
|
||||
return
|
||||
}
|
||||
resp, err := s.vectors.QueryVectors(r.Context(), &vector_pb.QueryVectorsRequest{
|
||||
Table: tableID(arns.bucket, arns.index),
|
||||
QueryVector: queryVector,
|
||||
TopK: uint32(req.TopK),
|
||||
Filter: filter,
|
||||
ReturnMetadata: req.ReturnMetadata,
|
||||
ReturnDistance: req.ReturnDistance,
|
||||
})
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "InternalFailure", err.Error())
|
||||
return
|
||||
}
|
||||
out := &QueryVectorsResponse{DistanceMetric: metric}
|
||||
for _, v := range resp.Vectors {
|
||||
o := &QueryOutputVector{Key: v.Key}
|
||||
if req.ReturnDistance {
|
||||
o.Distance = &v.Distance
|
||||
}
|
||||
if req.ReturnMetadata {
|
||||
o.Metadata = structToMap(v.Metadata)
|
||||
}
|
||||
out.Vectors = append(out.Vectors, o)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
// Package s3vectors serves the AWS S3 Vectors API on top of the shared
|
||||
// VectorService gRPC boundary and the S3 Tables catalog.
|
||||
package s3vectors
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/gorilla/mux"
|
||||
"github.com/seaweedfs/seaweedfs/weed/pb/filer_pb"
|
||||
"github.com/seaweedfs/seaweedfs/weed/pb/vector_pb"
|
||||
"github.com/seaweedfs/seaweedfs/weed/s3api/s3err"
|
||||
"github.com/seaweedfs/seaweedfs/weed/s3api/s3tables"
|
||||
)
|
||||
|
||||
// FilerClient provides access to the filer for storage operations.
|
||||
type FilerClient interface {
|
||||
WithFilerClient(streamingMode bool, fn func(client filer_pb.SeaweedFilerClient) error) error
|
||||
}
|
||||
|
||||
// S3Authenticator authenticates a request, returning the caller identity.
|
||||
type S3Authenticator interface {
|
||||
AuthenticateRequest(r *http.Request) (string, interface{}, s3err.ErrorCode)
|
||||
DefaultAllow() bool
|
||||
}
|
||||
|
||||
// Server implements the AWS S3 Vectors JSON API.
|
||||
type Server struct {
|
||||
filerClient FilerClient
|
||||
tables *s3tables.Manager
|
||||
vectors vector_pb.VectorServiceClient
|
||||
authenticator S3Authenticator
|
||||
region string
|
||||
accountID string
|
||||
}
|
||||
|
||||
// NewServer creates an S3 Vectors server over the given filer and catalog.
|
||||
func NewServer(filerClient FilerClient, authenticator S3Authenticator, vectors vector_pb.VectorServiceClient) *Server {
|
||||
manager := s3tables.NewManager()
|
||||
if authenticator != nil {
|
||||
manager.SetDefaultAllow(authenticator.DefaultAllow())
|
||||
}
|
||||
return &Server{
|
||||
filerClient: filerClient,
|
||||
tables: manager,
|
||||
vectors: vectors,
|
||||
authenticator: authenticator,
|
||||
}
|
||||
}
|
||||
|
||||
// SetRegion configures the region embedded in the ARNs this server mints.
|
||||
func (s *Server) SetRegion(region string) { s.region = region }
|
||||
|
||||
// SetAccountID configures the account id embedded in the ARNs this server mints.
|
||||
func (s *Server) SetAccountID(id string) { s.accountID = id }
|
||||
|
||||
// RegisterRoutes mounts POST / dispatching on the X-Amz-Target header and
|
||||
// POST /{Operation} for clients that route per-path.
|
||||
func (s *Server) RegisterRoutes(router *mux.Router) {
|
||||
router.HandleFunc("/", s.Auth(s.dispatch)).Methods(http.MethodPost)
|
||||
router.HandleFunc("/{op}", s.Auth(s.dispatch)).Methods(http.MethodPost)
|
||||
}
|
||||
|
||||
func (s *Server) dispatch(w http.ResponseWriter, r *http.Request) {
|
||||
// Clients carry the operation in the X-Amz-Target header or in the path —
|
||||
// POST /QueryVectors and POST / with a target mean the same call here.
|
||||
op := strings.Trim(r.URL.Path, "/")
|
||||
if i := strings.LastIndex(op, "/"); i >= 0 {
|
||||
op = op[i+1:]
|
||||
}
|
||||
if op == "" {
|
||||
target := r.Header.Get("X-Amz-Target")
|
||||
op = target[strings.LastIndex(target, ".")+1:]
|
||||
}
|
||||
handler, ok := handlers[op]
|
||||
if !ok {
|
||||
writeError(w, http.StatusBadRequest, "ValidationException", "unknown operation "+op)
|
||||
return
|
||||
}
|
||||
handler(s, w, r)
|
||||
}
|
||||
|
||||
// Auth authenticates the caller and puts the identity in the request context.
|
||||
func (s *Server) Auth(handler http.HandlerFunc) http.HandlerFunc {
|
||||
return func(w http.ResponseWriter, r *http.Request) {
|
||||
identity, _, errCode := s.authenticator.AuthenticateRequest(r)
|
||||
if errCode != s3err.ErrNone {
|
||||
writeError(w, http.StatusForbidden, "AccessDenied", "authentication failed")
|
||||
return
|
||||
}
|
||||
r = r.WithContext(contextWithIdentity(r.Context(), identity))
|
||||
handler(w, r)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,251 @@
|
||||
package s3vectors
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
type contextKey string
|
||||
|
||||
const identityKey contextKey = "s3vectors-identity"
|
||||
|
||||
func contextWithIdentity(ctx context.Context, identity string) context.Context {
|
||||
return context.WithValue(ctx, identityKey, identity)
|
||||
}
|
||||
|
||||
func identityOf(r *http.Request) string {
|
||||
v, _ := r.Context().Value(identityKey).(string)
|
||||
return v
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, status int, payload interface{}) {
|
||||
w.Header().Set("Content-Type", "application/x-amz-json-1.1")
|
||||
w.WriteHeader(status)
|
||||
if payload != nil {
|
||||
_ = json.NewEncoder(w).Encode(payload)
|
||||
}
|
||||
}
|
||||
|
||||
func writeError(w http.ResponseWriter, status int, code, message string) {
|
||||
writeJSON(w, status, map[string]string{
|
||||
"__type": code,
|
||||
"code": code,
|
||||
"message": message,
|
||||
"requestId": "local",
|
||||
})
|
||||
}
|
||||
|
||||
const maxRequestBodyBytes = 32 << 20
|
||||
|
||||
func readJSON(r *http.Request, req interface{}) error {
|
||||
return json.NewDecoder(io.LimitReader(r.Body, maxRequestBodyBytes)).Decode(req)
|
||||
}
|
||||
|
||||
// Request/response shapes for the AWS S3 Vectors API (awsJson1_1).
|
||||
|
||||
type vectorBucketARNs struct {
|
||||
bucket string
|
||||
index string
|
||||
}
|
||||
|
||||
type CreateVectorBucketRequest struct {
|
||||
VectorBucketName string `json:"vectorBucketName"`
|
||||
EncryptionConfiguration map[string]interface{} `json:"encryptionConfiguration,omitempty"`
|
||||
MetadataResourceTags map[string]string `json:"-"`
|
||||
}
|
||||
|
||||
type CreateVectorBucketResponse struct{}
|
||||
|
||||
type GetVectorBucketRequest struct {
|
||||
VectorBucketName string `json:"vectorBucketName"`
|
||||
VectorBucketArn string `json:"vectorBucketArn"`
|
||||
}
|
||||
|
||||
type GetVectorBucketResponse struct {
|
||||
VectorBucket *VectorBucket `json:"vectorBucket"`
|
||||
}
|
||||
|
||||
type VectorBucket struct {
|
||||
VectorBucketArn string `json:"vectorBucketArn"`
|
||||
VectorBucketName string `json:"vectorBucketName"`
|
||||
CreationTime time.Time `json:"creationTime"`
|
||||
}
|
||||
|
||||
type ListVectorBucketsRequest struct {
|
||||
MaxResults int `json:"maxResults"`
|
||||
NextToken string `json:"nextToken"`
|
||||
Prefix string `json:"prefix"`
|
||||
}
|
||||
|
||||
type ListVectorBucketsResponse struct {
|
||||
NextToken string `json:"nextToken,omitempty"`
|
||||
VectorBuckets []*VectorBucket `json:"vectorBuckets"`
|
||||
}
|
||||
|
||||
type DeleteVectorBucketRequest struct {
|
||||
VectorBucketName string `json:"vectorBucketName"`
|
||||
VectorBucketArn string `json:"vectorBucketArn"`
|
||||
}
|
||||
|
||||
type DeleteVectorBucketResponse struct{}
|
||||
|
||||
type CreateIndexRequest struct {
|
||||
DataType string `json:"dataType"`
|
||||
Dimension int32 `json:"dimension"`
|
||||
DistanceMetric string `json:"distanceMetric"`
|
||||
IndexArn string `json:"indexArn"`
|
||||
IndexName string `json:"indexName"`
|
||||
MetadataConfiguration *MetadataConfiguration `json:"metadataConfiguration,omitempty"`
|
||||
VectorBucketName string `json:"vectorBucketName"`
|
||||
VectorBucketArn string `json:"vectorBucketArn"`
|
||||
}
|
||||
|
||||
type MetadataConfiguration struct {
|
||||
NonFilterableMetadataKeys []string `json:"nonFilterableMetadataKeys"`
|
||||
}
|
||||
|
||||
type CreateIndexResponse struct{}
|
||||
|
||||
type GetIndexRequest struct {
|
||||
IndexArn string `json:"indexArn"`
|
||||
IndexName string `json:"indexName"`
|
||||
VectorBucketName string `json:"vectorBucketName"`
|
||||
VectorBucketArn string `json:"vectorBucketArn"`
|
||||
}
|
||||
|
||||
type Index struct {
|
||||
CreationTime time.Time `json:"creationTime"`
|
||||
DataType string `json:"dataType"`
|
||||
Dimension int32 `json:"dimension"`
|
||||
DistanceMetric string `json:"distanceMetric"`
|
||||
IndexArn string `json:"indexArn"`
|
||||
IndexName string `json:"indexName"`
|
||||
MetadataConfiguration *MetadataConfiguration `json:"metadataConfiguration,omitempty"`
|
||||
VectorBucketName string `json:"vectorBucketName"`
|
||||
}
|
||||
|
||||
type GetIndexResponse struct {
|
||||
Index *Index `json:"index"`
|
||||
}
|
||||
|
||||
type ListIndexesRequest struct {
|
||||
MaxResults int `json:"maxResults"`
|
||||
NextToken string `json:"nextToken"`
|
||||
Prefix string `json:"prefix"`
|
||||
VectorBucketName string `json:"vectorBucketName"`
|
||||
VectorBucketArn string `json:"vectorBucketArn"`
|
||||
}
|
||||
|
||||
type ListIndexesResponse struct {
|
||||
Indexes []*Index `json:"indexes"`
|
||||
NextToken string `json:"nextToken,omitempty"`
|
||||
}
|
||||
|
||||
type DeleteIndexRequest struct {
|
||||
IndexArn string `json:"indexArn"`
|
||||
IndexName string `json:"indexName"`
|
||||
VectorBucketName string `json:"vectorBucketName"`
|
||||
VectorBucketArn string `json:"vectorBucketArn"`
|
||||
}
|
||||
|
||||
type DeleteIndexResponse struct{}
|
||||
|
||||
type PutInputVector struct {
|
||||
Data *VectorData `json:"data"`
|
||||
Key string `json:"key"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
type VectorData struct {
|
||||
Float32 []float32 `json:"float32"`
|
||||
}
|
||||
|
||||
type PutVectorsRequest struct {
|
||||
VectorBucketName string `json:"vectorBucketName"`
|
||||
VectorBucketArn string `json:"vectorBucketArn"`
|
||||
IndexArn string `json:"indexArn"`
|
||||
IndexName string `json:"indexName"`
|
||||
Vectors []*PutInputVector `json:"vectors"`
|
||||
}
|
||||
|
||||
type PutVectorsResponse struct{}
|
||||
|
||||
type GetVectorsRequest struct {
|
||||
IndexArn string `json:"indexArn"`
|
||||
IndexName string `json:"indexName"`
|
||||
VectorBucketName string `json:"vectorBucketName"`
|
||||
VectorBucketArn string `json:"vectorBucketArn"`
|
||||
Keys []string `json:"keys"`
|
||||
ReturnData bool `json:"returnData"`
|
||||
ReturnMetadata bool `json:"returnMetadata"`
|
||||
}
|
||||
|
||||
type OutputVector struct {
|
||||
Data *VectorData `json:"data,omitempty"`
|
||||
Key string `json:"key"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
type GetVectorsResponse struct {
|
||||
Vectors []*OutputVector `json:"vectors"`
|
||||
}
|
||||
|
||||
type ListVectorsRequest struct {
|
||||
IndexArn string `json:"indexArn"`
|
||||
IndexName string `json:"indexName"`
|
||||
VectorBucketName string `json:"vectorBucketName"`
|
||||
VectorBucketArn string `json:"vectorBucketArn"`
|
||||
MaxResults int `json:"maxResults"`
|
||||
NextToken string `json:"nextToken"`
|
||||
ReturnData bool `json:"returnData"`
|
||||
ReturnMetadata bool `json:"returnMetadata"`
|
||||
SegmentCount int `json:"segmentCount"`
|
||||
SegmentIndex int `json:"segmentIndex"`
|
||||
}
|
||||
|
||||
type ListVectorsResponse struct {
|
||||
NextToken string `json:"nextToken,omitempty"`
|
||||
Vectors []*OutputVector `json:"vectors"`
|
||||
}
|
||||
|
||||
type DeleteVectorsRequest struct {
|
||||
IndexArn string `json:"indexArn"`
|
||||
IndexName string `json:"indexName"`
|
||||
VectorBucketName string `json:"vectorBucketName"`
|
||||
VectorBucketArn string `json:"vectorBucketArn"`
|
||||
Keys []string `json:"keys"`
|
||||
}
|
||||
|
||||
type DeleteVectorsResponse struct{}
|
||||
|
||||
type QueryVector struct {
|
||||
Float32 []float32 `json:"float32"`
|
||||
}
|
||||
|
||||
type QueryVectorsRequest struct {
|
||||
DistanceMetric string `json:"distanceMetric"`
|
||||
Filter map[string]interface{} `json:"filter"`
|
||||
IndexArn string `json:"indexArn"`
|
||||
IndexName string `json:"indexName"`
|
||||
QueryVector *QueryVector `json:"queryVector"`
|
||||
ReturnDistance bool `json:"returnDistance"`
|
||||
ReturnMetadata bool `json:"returnMetadata"`
|
||||
TopK int32 `json:"topK"`
|
||||
VectorBucketName string `json:"vectorBucketName"`
|
||||
VectorBucketArn string `json:"vectorBucketArn"`
|
||||
}
|
||||
|
||||
type QueryOutputVector struct {
|
||||
Data *VectorData `json:"data,omitempty"`
|
||||
Distance *float32 `json:"distance,omitempty"`
|
||||
Key string `json:"key"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
type QueryVectorsResponse struct {
|
||||
DistanceMetric string `json:"distanceMetric"`
|
||||
Vectors []*QueryOutputVector `json:"vectors"`
|
||||
}
|
||||
Reference in new issue
Block a user