mirror of
https://github.com/seaweedfs/seaweedfs.git
synced 2026-11-05 11:37:44 +01:00
Chris LuandClaude Opus 5.5
b74111ea2f
A query's lease tag is deleted by a background task after the query answers, so the encryption helper could list it and then get a 404 on HEAD. Skip only that case; any other missing object still fails. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
669 lines
34 KiB
Python
669 lines
34 KiB
Python
"""Run the S3 Vectors API against Go/Rust volumes and single/distributed queries."""
|
|
|
|
import contextlib
|
|
import json
|
|
import time
|
|
import uuid
|
|
|
|
import boto3
|
|
from botocore.config import Config
|
|
from botocore.exceptions import ClientError
|
|
import pytest
|
|
|
|
BUCKET = "vectors-bucket"
|
|
INDEX = "idx"
|
|
|
|
|
|
def test_vector_lifecycle(vectors):
|
|
vectors.create_vector_bucket(vectorBucketName=BUCKET)
|
|
buckets = [b["vectorBucketName"] for b in vectors.list_vector_buckets()["vectorBuckets"]]
|
|
assert BUCKET in buckets
|
|
|
|
vectors.create_index(
|
|
vectorBucketName=BUCKET,
|
|
indexName=INDEX,
|
|
dataType="float32",
|
|
dimension=4,
|
|
distanceMetric="euclidean",
|
|
)
|
|
idx = vectors.get_index(vectorBucketName=BUCKET, indexName=INDEX)["index"]
|
|
assert idx["dimension"] == 4
|
|
|
|
vectors.put_vectors(
|
|
vectorBucketName=BUCKET,
|
|
indexName=INDEX,
|
|
vectors=[
|
|
{"key": "a", "data": {"float32": [1, 0, 0, 0]}, "metadata": {"color": "red"}},
|
|
{"key": "b", "data": {"float32": [0.95, 0.05, 0, 0]}, "metadata": {"color": "blue"}},
|
|
{"key": "c", "data": {"float32": [0.9, 0.1, 0, 0]}, "metadata": {"color": "red"}},
|
|
],
|
|
)
|
|
|
|
got = vectors.get_vectors(
|
|
vectorBucketName=BUCKET, indexName=INDEX, keys=["a"], returnData=True, returnMetadata=True
|
|
)["vectors"]
|
|
assert got[0]["key"] == "a" and got[0]["metadata"]["color"] == "red"
|
|
|
|
listed = vectors.list_vectors(vectorBucketName=BUCKET, indexName=INDEX)["vectors"]
|
|
assert sorted(v["key"] for v in listed) == ["a", "b", "c"]
|
|
|
|
hits = vectors.query_vectors(
|
|
vectorBucketName=BUCKET,
|
|
indexName=INDEX,
|
|
queryVector={"float32": [1, 0, 0, 0]},
|
|
topK=2,
|
|
returnDistance=True,
|
|
returnMetadata=True,
|
|
filter={"color": {"$eq": "red"}},
|
|
)["vectors"]
|
|
# b is closer than c but excluded by the filter.
|
|
assert [h["key"] for h in hits] == ["a", "c"]
|
|
distances = {h["key"]: h["distance"] for h in hits}
|
|
assert distances["a"] == pytest.approx(0.0, abs=1e-6)
|
|
assert distances["c"] == pytest.approx(0.02 ** 0.5, abs=1e-4)
|
|
assert all(h["metadata"]["color"] == "red" for h in hits)
|
|
|
|
vectors.delete_vectors(vectorBucketName=BUCKET, indexName=INDEX, keys=["b"])
|
|
listed = vectors.list_vectors(vectorBucketName=BUCKET, indexName=INDEX)["vectors"]
|
|
assert sorted(v["key"] for v in listed) == ["a", "c"]
|
|
|
|
vectors.delete_index(vectorBucketName=BUCKET, indexName=INDEX)
|
|
vectors.delete_vector_bucket(vectorBucketName=BUCKET)
|
|
|
|
|
|
def test_vector_bucket_policy(vectors):
|
|
vectors.create_vector_bucket(vectorBucketName=BUCKET)
|
|
try:
|
|
policy = json.dumps(
|
|
{
|
|
"Version": "2012-10-17",
|
|
"Statement": [
|
|
{
|
|
"Effect": "Allow",
|
|
"Principal": {"AWS": "*"},
|
|
"Action": "s3vectors:*",
|
|
"Resource": "*",
|
|
}
|
|
],
|
|
}
|
|
)
|
|
invalid = json.loads(policy)
|
|
invalid["Statement"][0]["Action"] = "s3tables:*"
|
|
expect_error("ValidationException", vectors.put_vector_bucket_policy,
|
|
vectorBucketName=BUCKET, policy=json.dumps(invalid))
|
|
vectors.put_vector_bucket_policy(vectorBucketName=BUCKET, policy=policy)
|
|
got = vectors.get_vector_bucket_policy(vectorBucketName=BUCKET)["policy"]
|
|
assert json.loads(got) == json.loads(policy)
|
|
|
|
vectors.delete_vector_bucket_policy(vectorBucketName=BUCKET)
|
|
with pytest.raises(ClientError) as exc:
|
|
vectors.get_vector_bucket_policy(vectorBucketName=BUCKET)
|
|
assert exc.value.response["Error"]["Code"] == "NotFoundException"
|
|
finally:
|
|
vectors.delete_vector_bucket(vectorBucketName=BUCKET)
|
|
|
|
|
|
@pytest.fixture
|
|
def index(vectors):
|
|
bucket = "vectors-" + uuid.uuid4().hex[:16]
|
|
vectors.create_vector_bucket(vectorBucketName=bucket)
|
|
vectors.create_index(vectorBucketName=bucket, indexName=INDEX, dataType="float32",
|
|
dimension=4, distanceMetric="euclidean",
|
|
metadataConfiguration={"nonFilterableMetadataKeys": ["private"]})
|
|
args = {"vectorBucketName": bucket, "indexName": INDEX}
|
|
yield args
|
|
vectors.delete_index(**args)
|
|
vectors.delete_vector_bucket(vectorBucketName=bucket)
|
|
|
|
|
|
def expect_error(code, operation, **kwargs):
|
|
with pytest.raises(ClientError) as exc:
|
|
operation(**kwargs)
|
|
assert exc.value.response["Error"]["Code"] == code
|
|
|
|
|
|
def test_empty_catalog(vectors):
|
|
bucket = "empty-" + uuid.uuid4().hex[:16]
|
|
assert vectors.list_vector_buckets(prefix=bucket)["vectorBuckets"] == []
|
|
vectors.create_vector_bucket(vectorBucketName=bucket)
|
|
try:
|
|
assert vectors.list_indexes(vectorBucketName=bucket)["indexes"] == []
|
|
finally:
|
|
vectors.delete_vector_bucket(vectorBucketName=bucket)
|
|
|
|
|
|
def test_empty_index(vectors, index):
|
|
assert vectors.get_vectors(**index, keys=["absent"])["vectors"] == []
|
|
assert vectors.list_vectors(**index)["vectors"] == []
|
|
assert vectors.query_vectors(**index, queryVector={"float32": [1, 0, 0, 0]}, topK=3)["vectors"] == []
|
|
expect_error("ValidationException", vectors.query_vectors, **index,
|
|
queryVector={"float32": [1, 0, 0]}, topK=3)
|
|
|
|
|
|
def test_upsert_and_invalid_batch(vectors, index):
|
|
vectors.put_vectors(**index, vectors=[{"key": "a", "data": {"float32": [1, 0, 0, 0]}, "metadata": {"revision": 1}}])
|
|
vectors.put_vectors(**index, vectors=[{"key": "a", "data": {"float32": [0, 1, 0, 0]}, "metadata": {"revision": 2}}])
|
|
expected = [{"key": "a", "data": {"float32": [0, 1, 0, 0]}, "metadata": {"revision": 2}}]
|
|
for invalid in (
|
|
[{"key": "a", "data": {"float32": [1, 1, 1, 1]}}, {"key": "bad", "data": {"float32": [1, 2]}}],
|
|
[{"key": "a", "data": {"float32": [1, 1, 1, 1]}}, {"key": "a", "data": {"float32": [2, 2, 2, 2]}}],
|
|
):
|
|
expect_error("ValidationException", vectors.put_vectors, **index, vectors=invalid)
|
|
assert vectors.list_vectors(**index, returnData=True, returnMetadata=True)["vectors"] == expected
|
|
|
|
|
|
def test_metadata_and_arn(vectors, index):
|
|
description = vectors.get_index(**index)["index"]
|
|
assert description["metadataConfiguration"] == {"nonFilterableMetadataKeys": ["private"]}
|
|
arn = {"indexArn": description["indexArn"]}
|
|
metadata = {"color": "red", "private": "secret", "nested": {"x": [1, True, "text"]}}
|
|
vectors.put_vectors(**arn, vectors=[{"key": "quote'key", "data": {"float32": [1, 0, 0, 0]}, "metadata": metadata}])
|
|
assert vectors.get_vectors(**arn, keys=["quote'key"], returnMetadata=True)["vectors"][0]["metadata"] == metadata
|
|
hits = vectors.query_vectors(**arn, queryVector={"float32": [1, 0, 0, 0]}, topK=1, returnMetadata=True)["vectors"]
|
|
assert hits[0]["metadata"] == metadata
|
|
assert "distance" not in hits[0]
|
|
expect_error("ValidationException", vectors.query_vectors, **arn, queryVector={"float32": [1, 0, 0, 0]},
|
|
topK=1, filter={"$and": [{"color": "red"}, {"private": "secret"}]})
|
|
vectors.delete_vectors(**arn, keys=["quote'key"])
|
|
assert vectors.list_vectors(**arn)["vectors"] == []
|
|
|
|
|
|
def test_list_pagination_and_segments(vectors, index):
|
|
rows = [{"key": f"key-{i:02}", "data": {"float32": [i, 1, 0, 0]}, "metadata": {"row": i}} for i in range(17)]
|
|
vectors.put_vectors(**index, vectors=rows)
|
|
|
|
def collect(**extra):
|
|
result, tokens = [], set()
|
|
while True:
|
|
page = vectors.list_vectors(**index, maxResults=2, returnData=True, returnMetadata=True, **extra)
|
|
result.extend(page["vectors"])
|
|
token = page.get("nextToken")
|
|
if not token:
|
|
return result
|
|
assert token not in tokens
|
|
tokens.add(token)
|
|
extra["nextToken"] = token
|
|
|
|
assert sorted(collect(), key=lambda r: r["key"]) == rows
|
|
segmented = [row for segment in range(4) for row in collect(segmentCount=4, segmentIndex=segment)]
|
|
assert sorted(segmented, key=lambda r: r["key"]) == rows
|
|
expect_error("ValidationException", vectors.list_vectors, **index, nextToken="invalid")
|
|
for bad in ({"maxResults": 1001}, {"segmentCount": 17, "segmentIndex": 0}, {"segmentIndex": 0}):
|
|
expect_error("ValidationException", vectors.list_vectors, **index, **bad)
|
|
|
|
|
|
def test_filter_types(vectors, index):
|
|
values = [1, "1", True, "true"]
|
|
vectors.put_vectors(**index, vectors=[{"key": str(i), "data": {"float32": [1, 0, 0, 0]}, "metadata": {"value": value}}
|
|
for i, value in enumerate(values)])
|
|
for i, value in enumerate(values):
|
|
hits = vectors.query_vectors(**index, queryVector={"float32": [1, 0, 0, 0]}, topK=4,
|
|
filter={"value": {"$eq": value}})["vectors"]
|
|
assert [hit["key"] for hit in hits] == [str(i)]
|
|
for op, members in (("$in", [1, "true"]), ("$nin", ["1", True])):
|
|
hits = vectors.query_vectors(**index, queryVector={"float32": [1, 0, 0, 0]}, topK=4,
|
|
filter={"value": {op: members}})["vectors"]
|
|
assert sorted(hit["key"] for hit in hits) == ["0", "3"]
|
|
|
|
|
|
def test_filter_arrays(vectors, index):
|
|
rows = {"space": ["space", "classic"], "nums": [1, 2.5], "strs": ["1", "2.5"], "flag": [True], "scalar": "space"}
|
|
vectors.put_vectors(**index, vectors=[{"key": key, "data": {"float32": [1, 0, 0, 0]}, "metadata": {"tags": tags}}
|
|
for key, tags in rows.items()])
|
|
|
|
def matches(condition):
|
|
hits = vectors.query_vectors(**index, queryVector={"float32": [1, 0, 0, 0]}, topK=10,
|
|
filter={"tags": condition})["vectors"]
|
|
return sorted(hit["key"] for hit in hits)
|
|
|
|
assert matches("space") == ["scalar", "space"]
|
|
assert matches({"$eq": 1}) == ["nums"]
|
|
assert matches({"$eq": "1"}) == ["strs"]
|
|
assert matches({"$eq": True}) == ["flag"]
|
|
assert matches({"$in": ["classic", 2.5]}) == ["nums", "space"]
|
|
assert matches({"$ne": "space"}) == ["flag", "nums", "strs"]
|
|
assert matches({"$nin": [1, "2.5"]}) == ["flag", "scalar", "space"]
|
|
|
|
|
|
def test_euclidean_distance(vectors, index):
|
|
vectors.put_vectors(**index, vectors=[{"key": "p", "data": {"float32": [3, 4, 0, 0]}}])
|
|
result = vectors.query_vectors(**index, queryVector={"float32": [0, 0, 0, 0]}, topK=1, returnDistance=True)
|
|
assert result["distanceMetric"] == "euclidean"
|
|
assert result["vectors"][0]["distance"] == pytest.approx(5.0)
|
|
|
|
|
|
def test_cosine_metric(vectors):
|
|
bucket = "cosine-" + uuid.uuid4().hex[:16]
|
|
vectors.create_vector_bucket(vectorBucketName=bucket)
|
|
args = {"vectorBucketName": bucket, "indexName": INDEX}
|
|
vectors.create_index(**args, dataType="float32", dimension=2, distanceMetric="cosine")
|
|
try:
|
|
vectors.put_vectors(**args, vectors=[{"key": "far", "data": {"float32": [10, 0]}},
|
|
{"key": "near", "data": {"float32": [1, 1]}}])
|
|
result = vectors.query_vectors(**args, queryVector={"float32": [1, 0]}, topK=2, returnDistance=True)
|
|
assert result["distanceMetric"] == "cosine"
|
|
assert [v["key"] for v in result["vectors"]] == ["far", "near"]
|
|
assert result["vectors"][0]["distance"] == pytest.approx(0, abs=1e-6)
|
|
expect_error("ValidationException", vectors.query_vectors, **args, queryVector={"float32": [1, 0, 0]}, topK=1)
|
|
finally:
|
|
vectors.delete_index(**args)
|
|
vectors.delete_vector_bucket(vectorBucketName=bucket)
|
|
|
|
|
|
def test_delete_bucket_conflict(vectors, index):
|
|
expect_error("ConflictException", vectors.delete_vector_bucket, vectorBucketName=index["vectorBucketName"])
|
|
assert vectors.get_index(**index)["index"]["dimension"] == 4
|
|
|
|
|
|
def test_vector_policy_permissions(vectors, vector_clients, index):
|
|
reader, writer = vector_clients["reader"], vector_clients["writer"]
|
|
arn = vectors.get_index(**index)["index"]["indexArn"]
|
|
row = {"key": "policy-row", "data": {"float32": [1, 0, 0, 0]}, "metadata": {"owner": "writer"}}
|
|
expect_error("AccessDeniedException", writer.put_vectors, **index, vectors=[row])
|
|
expect_error("AccessDeniedException", reader.get_vectors, **index, keys=[row["key"]])
|
|
statements = [
|
|
{"Effect": "Allow", "Principal": {"AWS": "writer"},
|
|
"Action": "s3vectors:PutVectors", "Resource": arn},
|
|
{"Effect": "Allow", "Principal": {"AWS": "reader"},
|
|
"Action": ["s3vectors:GetVectors", "s3vectors:QueryVectors"], "Resource": arn},
|
|
]
|
|
|
|
def put_policy():
|
|
vectors.put_vector_bucket_policy(vectorBucketName=index["vectorBucketName"],
|
|
policy=json.dumps({"Version": "2012-10-17", "Statement": statements}))
|
|
|
|
put_policy()
|
|
writer.put_vectors(**index, vectors=[row])
|
|
assert reader.get_vectors(**index, keys=[row["key"]], returnData=True, returnMetadata=True)["vectors"] == [row]
|
|
assert reader.query_vectors(**index, queryVector=row["data"], topK=1)["vectors"][0]["key"] == row["key"]
|
|
expect_error("AccessDeniedException", reader.put_vectors, **index, vectors=[row])
|
|
expect_error("AccessDeniedException", writer.get_vectors, **index, keys=[row["key"]])
|
|
expect_error("AccessDeniedException", reader.delete_index, **index)
|
|
expect_error("AccessDeniedException", reader.list_vectors, **index)
|
|
other = dict(index, indexName="other-index")
|
|
vectors.create_index(**other, dataType="float32", dimension=4, distanceMetric="euclidean")
|
|
try:
|
|
expect_error("AccessDeniedException", reader.get_vectors, **other, keys=[row["key"]])
|
|
expect_error("AccessDeniedException", writer.put_vectors, **other, vectors=[row])
|
|
statements.append({"Effect": "Deny", "Principal": {"AWS": "writer"},
|
|
"Action": "s3vectors:PutVectors", "Resource": arn})
|
|
put_policy()
|
|
expect_error("AccessDeniedException", writer.put_vectors, **index, vectors=[row])
|
|
assert reader.get_vectors(**index, keys=[row["key"]])["vectors"][0]["key"] == row["key"]
|
|
vectors.delete_vector_bucket_policy(vectorBucketName=index["vectorBucketName"])
|
|
expect_error("AccessDeniedException", reader.get_vectors, **index, keys=[row["key"]])
|
|
expect_error("AccessDeniedException", writer.put_vectors, **index, vectors=[row])
|
|
vectors.put_vectors(**index, vectors=[row])
|
|
finally:
|
|
vectors.delete_index(**other)
|
|
|
|
|
|
def index_object_encryption(vectors, bucket):
|
|
s3 = boto3.client("s3", endpoint_url=vectors.meta.endpoint_url, aws_access_key_id="test",
|
|
aws_secret_access_key="test", region_name="us-east-1")
|
|
keys = [o["Key"] for page in s3.get_paginator("list_objects_v2").paginate(Bucket=bucket, Prefix=f"vector/{INDEX}/")
|
|
for o in page.get("Contents", []) if o["Size"] > 0]
|
|
assert any("/data/" in key for key in keys) and any("/_versions/" in key for key in keys)
|
|
encryption = {}
|
|
for key in keys:
|
|
try:
|
|
encryption[key] = s3.head_object(Bucket=bucket, Key=key).get("ServerSideEncryption")
|
|
except ClientError as exc:
|
|
# a query's lease tag is released in the background after it answers
|
|
if "/_refs/tags/weed-vector-query-" in key and exc.response["Error"]["Code"] in ("404", "NoSuchKey"):
|
|
continue
|
|
raise
|
|
return encryption
|
|
|
|
|
|
def test_vector_bucket_encryption(vectors):
|
|
bucket = "encryption-" + uuid.uuid4().hex[:16]
|
|
plain = bucket + "-plain"
|
|
config = {"sseType": "AES256"}
|
|
vectors.create_vector_bucket(vectorBucketName=bucket, encryptionConfiguration=config)
|
|
vectors.create_vector_bucket(vectorBucketName=plain)
|
|
try:
|
|
assert vectors.get_vector_bucket(vectorBucketName=bucket)["vectorBucket"]["encryptionConfiguration"] == config
|
|
expect_error("ValidationException", vectors.create_vector_bucket, vectorBucketName=bucket + "-kms",
|
|
encryptionConfiguration={"sseType": "aws:kms", "kmsKeyArn": "arn:aws:kms:us-east-1:123456789012:key/test"})
|
|
assert vectors.list_vector_buckets(prefix=bucket + "-kms")["vectorBuckets"] == []
|
|
|
|
for name in (bucket, plain):
|
|
vectors.create_index(vectorBucketName=name, indexName=INDEX, dataType="float32", dimension=4,
|
|
distanceMetric="euclidean")
|
|
vectors.put_vectors(vectorBucketName=name, indexName=INDEX, vectors=[
|
|
{"key": "a", "data": {"float32": [1, 0, 0, 0]}, "metadata": {"color": "red"}},
|
|
{"key": "b", "data": {"float32": [0, 1, 0, 0]}, "metadata": {"color": "blue"}},
|
|
])
|
|
hits = vectors.query_vectors(vectorBucketName=name, indexName=INDEX, queryVector={"float32": [1, 0, 0, 0]},
|
|
topK=1, returnMetadata=True)["vectors"]
|
|
assert [(h["key"], h["metadata"]) for h in hits] == [("a", {"color": "red"})]
|
|
assert set(index_object_encryption(vectors, bucket).values()) == {"AES256"}
|
|
assert set(index_object_encryption(vectors, plain).values()) == {None}
|
|
finally:
|
|
for name in (bucket, plain):
|
|
with contextlib.suppress(ClientError):
|
|
vectors.delete_index(vectorBucketName=name, indexName=INDEX)
|
|
vectors.delete_vector_bucket(vectorBucketName=name)
|
|
|
|
|
|
def test_catalog_pagination(vectors, index):
|
|
prefix = "pages-" + uuid.uuid4().hex[:16]
|
|
buckets, indexes = [], []
|
|
try:
|
|
for i in range(3):
|
|
bucket = f"{prefix}-{i}"
|
|
vectors.create_vector_bucket(vectorBucketName=bucket)
|
|
buckets.append(bucket)
|
|
for name in ("idx-a", "idx-b"):
|
|
args = dict(index, indexName=name)
|
|
vectors.create_index(**args, dataType="float32", dimension=4, distanceMetric="euclidean")
|
|
indexes.append(args)
|
|
pages = list(vectors.get_paginator("list_vector_buckets").paginate(
|
|
prefix=prefix, PaginationConfig={"PageSize": 1}))
|
|
assert all(len(page["vectorBuckets"]) <= 1 for page in pages)
|
|
assert sorted(b["vectorBucketName"] for page in pages for b in page["vectorBuckets"]) == buckets
|
|
pages = list(vectors.get_paginator("list_indexes").paginate(
|
|
vectorBucketName=index["vectorBucketName"], prefix="idx", PaginationConfig={"PageSize": 1}))
|
|
assert all(len(page["indexes"]) <= 1 for page in pages)
|
|
assert sorted(i["indexName"] for page in pages for i in page["indexes"]) == [INDEX, "idx-a", "idx-b"]
|
|
expect_error("ValidationException", vectors.list_vector_buckets, maxResults=501)
|
|
expect_error("ValidationException", vectors.list_indexes, vectorBucketName=index["vectorBucketName"], maxResults=501)
|
|
finally:
|
|
for args in indexes:
|
|
vectors.delete_index(**args)
|
|
for bucket in buckets:
|
|
vectors.delete_vector_bucket(vectorBucketName=bucket)
|
|
|
|
|
|
def test_delete_and_recreate_index(vectors, index):
|
|
vectors.put_vectors(**index, vectors=[{"key": "old", "data": {"float32": [1, 0, 0, 0]}}])
|
|
vectors.delete_index(**index)
|
|
expect_error("NotFoundException", vectors.get_index, **index)
|
|
vectors.create_index(**index, dataType="float32", dimension=2, distanceMetric="cosine")
|
|
assert vectors.list_vectors(**index)["vectors"] == []
|
|
row = {"key": "new", "data": {"float32": [1, 0]}}
|
|
vectors.put_vectors(**index, vectors=[row])
|
|
assert vectors.get_vectors(**index, keys=["old", "new"], returnData=True)["vectors"] == [row]
|
|
result = vectors.query_vectors(**index, queryVector=row["data"], topK=1, returnDistance=True)
|
|
assert result["distanceMetric"] == "cosine"
|
|
assert result["vectors"][0]["key"] == "new"
|
|
assert result["vectors"][0]["distance"] == pytest.approx(0, abs=1e-6)
|
|
|
|
|
|
def test_index_arn_validation(vectors, index):
|
|
arn = vectors.get_index(**index)["index"]["indexArn"]
|
|
foreign = []
|
|
for field, value in ((1, "aws-cn"), (3, "eu-west-1"), (4, "999999999999")):
|
|
parts = arn.split(":", 5)
|
|
assert parts[field] != value
|
|
parts[field] = value
|
|
foreign.append(":".join(parts))
|
|
for invalid in (
|
|
*foreign,
|
|
arn.replace(":s3vectors:", ":s3tables:"),
|
|
arn.replace(":s3vectors:", ":s3:"),
|
|
arn + "/index/extra",
|
|
arn + "/",
|
|
):
|
|
expect_error("ValidationException", vectors.get_index, indexArn=invalid)
|
|
expect_error("ValidationException", vectors.delete_index, indexArn=invalid)
|
|
expect_error("ValidationException", vectors.get_vector_bucket, vectorBucketArn=arn)
|
|
expect_error("ValidationException", vectors.list_indexes, vectorBucketArn=arn)
|
|
bucket_arn = arn.rsplit("/index/", 1)[0]
|
|
for request in (
|
|
{"vectorBucketName": "other-" + index["vectorBucketName"], "indexArn": arn},
|
|
{"indexName": "other-" + index["indexName"], "indexArn": arn},
|
|
):
|
|
expect_error("ValidationException", vectors.get_index, **request)
|
|
expect_error("ValidationException", vectors.delete_index, **request)
|
|
for operation in (vectors.get_vector_bucket, vectors.list_indexes, vectors.delete_vector_bucket):
|
|
expect_error("ValidationException", operation, vectorBucketName="other-" + index["vectorBucketName"],
|
|
vectorBucketArn=bucket_arn)
|
|
assert vectors.get_index(**index, indexArn=arn)["index"]["indexArn"] == arn
|
|
assert vectors.get_vector_bucket(vectorBucketName=index["vectorBucketName"],
|
|
vectorBucketArn=bucket_arn)["vectorBucket"]["vectorBucketArn"] == bucket_arn
|
|
assert vectors.get_index(**index)["index"]["indexArn"] == arn
|
|
|
|
|
|
def test_resource_tags(vectors, vector_clients, index):
|
|
reader = vector_clients["reader"]
|
|
bucket_arn = vectors.get_vector_bucket(vectorBucketName=index["vectorBucketName"])["vectorBucket"]["vectorBucketArn"]
|
|
index_arn = vectors.get_index(**index)["index"]["indexArn"]
|
|
assert vectors.list_tags_for_resource(resourceArn=bucket_arn)["tags"] == {}
|
|
vectors.tag_resource(resourceArn=bucket_arn, tags={"team": "search", "env": "dev"})
|
|
vectors.tag_resource(resourceArn=index_arn, tags={"team": "ranking"})
|
|
assert vectors.list_tags_for_resource(resourceArn=bucket_arn)["tags"] == {"team": "search", "env": "dev"}
|
|
assert vectors.list_tags_for_resource(resourceArn=index_arn)["tags"] == {"team": "ranking"}
|
|
vectors.untag_resource(resourceArn=bucket_arn, tagKeys=["env", "absent"])
|
|
assert vectors.list_tags_for_resource(resourceArn=bucket_arn)["tags"] == {"team": "search"}
|
|
expect_error("ValidationException", vectors.tag_resource, resourceArn=bucket_arn,
|
|
tags={f"k{i}": "v" for i in range(50)})
|
|
expect_error("NotFoundException", vectors.list_tags_for_resource, resourceArn=bucket_arn + "/index/missing")
|
|
expect_error("AccessDeniedException", reader.list_tags_for_resource, resourceArn=bucket_arn)
|
|
expect_error("AccessDeniedException", reader.tag_resource, resourceArn=index_arn, tags={"team": "other"})
|
|
assert vectors.list_tags_for_resource(resourceArn=index_arn)["tags"] == {"team": "ranking"}
|
|
|
|
|
|
def test_create_with_tags(vectors):
|
|
bucket = "tagged-" + uuid.uuid4().hex[:16]
|
|
bucket_arn = vectors.create_vector_bucket(vectorBucketName=bucket, tags={"team": "search"})["vectorBucketArn"]
|
|
try:
|
|
assert vectors.list_tags_for_resource(resourceArn=bucket_arn)["tags"] == {"team": "search"}
|
|
index_arn = vectors.create_index(vectorBucketName=bucket, indexName=INDEX, dataType="float32", dimension=4,
|
|
distanceMetric="cosine", tags={"tier": "hot"})["indexArn"]
|
|
try:
|
|
assert vectors.list_tags_for_resource(resourceArn=index_arn)["tags"] == {"tier": "hot"}
|
|
finally:
|
|
vectors.delete_index(vectorBucketName=bucket, indexName=INDEX)
|
|
finally:
|
|
vectors.delete_vector_bucket(vectorBucketName=bucket)
|
|
|
|
|
|
def test_tag_conditions(vectors, vector_clients, index):
|
|
reader, writer = vector_clients["reader"], vector_clients["writer"]
|
|
bucket = index["vectorBucketName"]
|
|
bucket_arn = vectors.get_vector_bucket(vectorBucketName=bucket)["vectorBucket"]["vectorBucketArn"]
|
|
index_arn = vectors.get_index(**index)["index"]["indexArn"]
|
|
other = dict(index, indexName="other-index")
|
|
vectors.create_index(**other, dataType="float32", dimension=4, distanceMetric="euclidean")
|
|
vectors.tag_resource(resourceArn=bucket_arn, tags={"env": "prod"})
|
|
vectors.tag_resource(resourceArn=index_arn, tags={"stage": "live"})
|
|
indexes = bucket_arn + "/index/*"
|
|
statements = [
|
|
{"Effect": "Allow", "Principal": {"AWS": "reader"}, "Action": "s3vectors:GetIndex", "Resource": indexes,
|
|
"Condition": {"StringEquals": {"aws:ResourceTag/stage": "live"}}},
|
|
{"Effect": "Allow", "Principal": {"AWS": "reader"}, "Action": "s3vectors:QueryVectors", "Resource": indexes,
|
|
"Condition": {"StringEquals": {"s3vectors:VectorBucketTag/env": "prod"}}},
|
|
{"Effect": "Allow", "Principal": {"AWS": "writer"}, "Action": "s3vectors:TagResource", "Resource": indexes,
|
|
"Condition": {"StringEquals": {"aws:RequestTag/team": "search"},
|
|
"ForAllValues:StringEquals": {"aws:TagKeys": ["team"]}}},
|
|
]
|
|
try:
|
|
vectors.put_vector_bucket_policy(vectorBucketName=bucket,
|
|
policy=json.dumps({"Version": "2012-10-17", "Statement": statements}))
|
|
assert reader.get_index(**index)["index"]["indexArn"] == index_arn
|
|
expect_error("AccessDeniedException", reader.get_index, **other)
|
|
assert reader.query_vectors(**other, queryVector={"float32": [1, 0, 0, 0]}, topK=1)["vectors"] == []
|
|
writer.tag_resource(resourceArn=index_arn, tags={"team": "search"})
|
|
expect_error("AccessDeniedException", writer.tag_resource, resourceArn=index_arn, tags={"team": "other"})
|
|
expect_error("AccessDeniedException", writer.tag_resource, resourceArn=index_arn,
|
|
tags={"team": "search", "env": "dev"})
|
|
assert vectors.list_tags_for_resource(resourceArn=index_arn)["tags"] == {"stage": "live", "team": "search"}
|
|
vectors.tag_resource(resourceArn=bucket_arn, tags={"env": "dev"})
|
|
expect_error("AccessDeniedException", reader.query_vectors, **other,
|
|
queryVector={"float32": [1, 0, 0, 0]}, topK=1)
|
|
unsupported = dict(statements[0], Condition={"IpAddress": {"aws:SourceIp": "10.0.0.0/8"}})
|
|
expect_error("ValidationException", vectors.put_vector_bucket_policy, vectorBucketName=bucket,
|
|
policy=json.dumps({"Version": "2012-10-17", "Statement": [unsupported]}))
|
|
finally:
|
|
vectors.delete_vector_bucket_policy(vectorBucketName=bucket)
|
|
vectors.delete_index(**other)
|
|
|
|
|
|
def exact_neighbours(rows, query, k):
|
|
import math
|
|
qn = math.sqrt(sum(x * x for x in query))
|
|
scored = sorted(((sum(a * b for a, b in zip(v, query)) / (math.sqrt(sum(a * a for a in v)) * qn), key)
|
|
for key, v in rows.items()), reverse=True)
|
|
return {key for _, key in scored[:k]}
|
|
|
|
|
|
@pytest.mark.worker
|
|
def test_vector_index_job(vectors, cluster):
|
|
"""The maintenance worker indexes a table past the row threshold through the
|
|
admin server, queries then run through the index, and rows written after
|
|
the build are still found."""
|
|
import os
|
|
import random
|
|
import subprocess
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
bucket, index, dim = "vectors-indexed", "idx", 16
|
|
vectors.create_vector_bucket(vectorBucketName=bucket)
|
|
vectors.create_index(vectorBucketName=bucket, indexName=index, dataType="float32",
|
|
dimension=dim, distanceMetric="cosine")
|
|
# Embedding-like data: points around shared centres. Independent Gaussian
|
|
# vectors have no locality for an IVF index to use, so recall through one
|
|
# would only track the fraction of partitions probed.
|
|
rng = random.Random(7)
|
|
centres = [[rng.gauss(0, 1) for _ in range(dim)] for _ in range(30)]
|
|
|
|
def point():
|
|
centre = rng.choice(centres)
|
|
return [c + 0.3 * rng.gauss(0, 1) for c in centre]
|
|
|
|
rows = {f"v{i}": point() for i in range(1200)}
|
|
keys = list(rows)
|
|
for offset in range(0, len(keys), 400):
|
|
vectors.put_vectors(vectorBucketName=bucket, indexName=index,
|
|
vectors=[{"key": key, "data": {"float32": rows[key]}} for key in keys[offset:offset + 400]])
|
|
|
|
admin = cluster["admin"]
|
|
job = "lance_vector_index"
|
|
|
|
def call(path, method="GET", body=None):
|
|
data = json.dumps(body).encode() if body is not None else None
|
|
request = urllib.request.Request(f"{admin}/api/plugin/job-types/{job}{path}", data=data, method=method,
|
|
headers={"Content-Type": "application/json"} if data else {})
|
|
with urllib.request.urlopen(request, timeout=120) as response:
|
|
return json.load(response)
|
|
|
|
call("/config", "PUT", {"admin_runtime": {"enabled": True},
|
|
"worker_config_values": {"min_rows": {"int64_value": "1000"}}})
|
|
run = call("/run", "POST", {"timeout_seconds": 110})
|
|
assert run["success_count"] == 1 and run["error_count"] == 0, run
|
|
|
|
cli = Path(os.environ.get("VECTOR_CLI_BINARY", Path(__file__).resolve().parent / "weed-vector-cli"))
|
|
coverage = json.loads(subprocess.run(
|
|
[str(cli), "--namespace", cluster["catalog"], "--access-key", "test", "--secret-key", "test",
|
|
"coverage", f"{bucket}$vector${index}"], check=True, capture_output=True, text=True).stdout)
|
|
assert len(coverage) == 1 and coverage[0]["segments"] == 1 and coverage[0]["unindexed_rows"] == 0, coverage
|
|
|
|
found = 0
|
|
for _ in range(20):
|
|
query = point()
|
|
hits = vectors.query_vectors(vectorBucketName=bucket, indexName=index,
|
|
queryVector={"float32": query}, topK=10)["vectors"]
|
|
assert len(hits) == 10
|
|
found += len({h["key"] for h in hits} & exact_neighbours(rows, query, 10))
|
|
assert found >= 180, f"recall {found}/200"
|
|
|
|
late = point()
|
|
vectors.put_vectors(vectorBucketName=bucket, indexName=index, vectors=[{"key": "late", "data": {"float32": late}}])
|
|
hits = vectors.query_vectors(vectorBucketName=bucket, indexName=index, queryVector={"float32": late}, topK=1)
|
|
assert [h["key"] for h in hits["vectors"]] == ["late"]
|
|
|
|
again = call("/run", "POST", {"timeout_seconds": 60})
|
|
assert again["detected_count"] == 0, again
|
|
|
|
def coverage_now():
|
|
return json.loads(subprocess.run(
|
|
[str(cli), "--namespace", cluster["catalog"], "--access-key", "test", "--secret-key", "test",
|
|
"coverage", f"{bucket}$vector${index}"], check=True, capture_output=True, text=True).stdout)[0]
|
|
|
|
def optimize(expect_detected):
|
|
request = urllib.request.Request(
|
|
f"{admin}/api/plugin/job-types/lance_optimize_indices/run", method="POST",
|
|
data=json.dumps({"timeout_seconds": 110}).encode(), headers={"Content-Type": "application/json"})
|
|
with urllib.request.urlopen(request, timeout=120) as response:
|
|
result = json.load(response)
|
|
assert result["detected_count"] == expect_detected and result["error_count"] == 0, result
|
|
|
|
# A tenth of the index written after the build brings the optimize job
|
|
# in, even though it is far below its absolute budget.
|
|
optimize(0)
|
|
more = {f"m{i}": point() for i in range(200)}
|
|
vectors.put_vectors(vectorBucketName=bucket, indexName=index,
|
|
vectors=[{"key": key, "data": {"float32": v}} for key, v in more.items()])
|
|
before = coverage_now()
|
|
assert before["unindexed_rows"] == 201, before
|
|
optimize(1)
|
|
after = coverage_now()
|
|
assert after["unindexed_rows"] == 0 and after["segments"] == 2, after
|
|
|
|
# Five times the rows the index was sized for: the job trains it again
|
|
# instead of extending it, and the partitions follow the new row count.
|
|
rows.update(more)
|
|
for offset in range(0, 4600, 400):
|
|
batch = {f"g{offset + i}": point() for i in range(400)}
|
|
rows.update(batch)
|
|
vectors.put_vectors(vectorBucketName=bucket, indexName=index,
|
|
vectors=[{"key": key, "data": {"float32": v}} for key, v in batch.items()])
|
|
optimize(1)
|
|
retrained = coverage_now()
|
|
assert retrained["unindexed_rows"] == 0 and retrained["segments"] == 1, retrained
|
|
assert retrained["partitions"] > after["partitions"], (before, after, retrained)
|
|
found = 0
|
|
for _ in range(20):
|
|
query = point()
|
|
hits = vectors.query_vectors(vectorBucketName=bucket, indexName=index,
|
|
queryVector={"float32": query}, topK=10)["vectors"]
|
|
found += len({h["key"] for h in hits} & exact_neighbours(rows, query, 10))
|
|
assert found >= 180, f"recall after retrain {found}/200"
|
|
|
|
|
|
def test_invalid_credentials(cluster):
|
|
client = boto3.client("s3vectors", endpoint_url=cluster["endpoint"], aws_access_key_id="unknown",
|
|
aws_secret_access_key="unknown", region_name="us-east-1")
|
|
with pytest.raises(ClientError) as exc:
|
|
client.list_vector_buckets()
|
|
assert exc.value.response["ResponseMetadata"]["HTTPStatusCode"] == 403
|
|
|
|
|
|
def test_query_without_writer(vectors, index, cluster):
|
|
from conftest import stop
|
|
|
|
vectors.put_vectors(**index, vectors=[{"key": "survivor", "data": {"float32": [1, 0, 0, 0]}}])
|
|
stop(cluster["processes"].pop("vector"))
|
|
args = dict(index, queryVector={"float32": [1, 0, 0, 0]}, topK=1)
|
|
if cluster["mode"] == "single":
|
|
expect_error("ServiceUnavailableException", vectors.query_vectors, **args)
|
|
return
|
|
deadline = time.monotonic() + 15
|
|
while True:
|
|
try:
|
|
result = vectors.query_vectors(**args)
|
|
break
|
|
except ClientError as exc:
|
|
if exc.response["Error"]["Code"] != "ServiceUnavailableException" or time.monotonic() >= deadline:
|
|
raise
|
|
time.sleep(0.2)
|
|
assert [v["key"] for v in result["vectors"]] == ["survivor"]
|
|
|
|
|
|
def test_engine_without_credentials(open_cluster):
|
|
"""With no S3 identities, an engine given no credentials still serves an index."""
|
|
client = boto3.client("s3vectors", endpoint_url=open_cluster["endpoint"], aws_access_key_id="any",
|
|
aws_secret_access_key="any", region_name="us-east-1",
|
|
config=Config(connect_timeout=3, read_timeout=20, retries={"max_attempts": 0}))
|
|
index = {"vectorBucketName": BUCKET, "indexName": INDEX}
|
|
client.create_vector_bucket(vectorBucketName=BUCKET)
|
|
client.create_index(**index, dataType="float32", dimension=2, distanceMetric="cosine")
|
|
client.put_vectors(**index, vectors=[{"key": "a", "data": {"float32": [1, 0]}},
|
|
{"key": "b", "data": {"float32": [0, 1]}}])
|
|
hits = client.query_vectors(**index, queryVector={"float32": [1, 0]}, topK=1)["vectors"]
|
|
assert [h["key"] for h in hits] == ["a"]
|
|
|