"""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 assert_engine_wrote_like_the_gateway(cluster, index): """The engine writes the table's files itself: the filer holds entries the gateway could have written (an ETag, chunks, no inline content), and under VECTOR_ENCRYPT_VOLUMES every chunk carries a cipher key.""" import os import urllib.request table = f"/buckets/{index['vectorBucketName']}/vector/{index['indexName']}" listing = urllib.request.Request(f"http://127.0.0.1:{cluster['filer']}{table}/data/", headers={"Accept": "application/json"}) with urllib.request.urlopen(listing, timeout=5) as response: entries = json.load(response)["Entries"] or [] assert entries, "the engine wrote no data file" for entry in entries: with urllib.request.urlopen(f"http://127.0.0.1:{cluster['filer']}{table}/data/{entry['FullPath'].rsplit('/', 1)[1]}?metadata=true", timeout=5) as response: meta = json.load(response) assert meta["chunks"], meta assert "Seaweed-X-Amz-ETag" in (meta.get("Extended") or {}), meta.get("Extended") if os.environ.get("VECTOR_ENCRYPT_VOLUMES"): assert all(chunk.get("cipher_key") for chunk in meta["chunks"]), meta["chunks"] def run_cli(cluster, *args): """weed-vector-cli's JSON answer. It reaches the cluster through the namespace that names the tables, the masters whose filers hold them and the cluster token's key, and finds security.toml and the read key where the weed processes do.""" import os import subprocess from pathlib import Path cli = Path(os.environ.get("VECTOR_CLI_BINARY", Path(__file__).resolve().parent / "weed-vector-cli")).resolve() command = [str(cli), "--namespace", cluster["catalog"], "--masters", cluster["masters"], "--filer-signing-key", cluster["filer_signing_key"], *args] return json.loads(subprocess.run(command, check=True, capture_output=True, text=True, cwd=cluster["directory"], env=cluster["env"]).stdout) 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, cluster): 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}}] assert_engine_wrote_like_the_gateway(cluster, index) 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_sse_s3_bucket(vectors, cluster): """An index in a bucket that encrypts with SSE-S3 is written encrypted by the engine, under keys the gateway can unwrap: the entries carry the SSE-S3 attributes, the chunks hold ciphertext, and the gateway's own S3 API serves the files decrypted.""" import urllib.request bucket = "sealed-" + uuid.uuid4().hex[:16] vectors.create_vector_bucket(vectorBucketName=bucket, encryptionConfiguration={"sseType": "AES256"}) args = {"vectorBucketName": bucket, "indexName": INDEX} # A small index, whose data file is one chunk, and a wide one whose data # file spans several: the gateway reads the two differently. vectors.create_index(**args, dataType="float32", dimension=2, distanceMetric="euclidean") vectors.put_vectors(**args, vectors=[{"key": "a", "data": {"float32": [1, 0]}}, {"key": "b", "data": {"float32": [0, 1]}}]) hits = vectors.query_vectors(**args, queryVector={"float32": [0, 1]}, topK=1)["vectors"] assert [h["key"] for h in hits] == ["b"] wide = {"vectorBucketName": bucket, "indexName": "wide"} vectors.create_index(**wide, dataType="float32", dimension=4096, distanceMetric="euclidean") vectors.put_vectors(**wide, vectors=[{"key": f"w{i}", "data": {"float32": [float(i)] * 4096}} for i in range(300)]) hits = vectors.query_vectors(**wide, queryVector={"float32": [7.0] * 4096}, topK=1)["vectors"] assert [h["key"] for h in hits] == ["w7"] s3 = boto3.client("s3", endpoint_url=cluster["endpoint"], aws_access_key_id="test", aws_secret_access_key="test", region_name="us-east-1", config=Config(s3={"addressing_style": "path"})) chunk_counts = set() for index in (INDEX, "wide"): table = f"/buckets/{bucket}/vector/{index}" listing = urllib.request.Request(f"http://127.0.0.1:{cluster['filer']}{table}/data/", headers={"Accept": "application/json"}) with urllib.request.urlopen(listing, timeout=5) as response: names = [e["FullPath"].rsplit("/", 1)[1] for e in json.load(response)["Entries"] or []] assert names, f"the engine wrote no data file for {index}" for name in names: with urllib.request.urlopen(f"http://127.0.0.1:{cluster['filer']}{table}/data/{name}?metadata=true", timeout=5) as response: meta = json.load(response) assert "x-seaweedfs-sse-s3-key" in meta["Extended"], meta["Extended"].keys() assert all(chunk.get("sse_type") == 3 for chunk in meta["chunks"]), meta["chunks"] chunk_counts.add(len(meta["chunks"])) # The gateway's own S3 API serves the file decrypted: a Lance # file, its magic in the footer, of the size the filer records. through_gateway = s3.get_object(Bucket=bucket, Key=f"vector/{index}/data/{name}")["Body"].read() assert through_gateway[-4:] == b"LANC", (index, name, through_gateway[-16:]) assert len(through_gateway) == meta["FileSize"], (len(through_gateway), meta["FileSize"]) assert 1 in chunk_counts and max(chunk_counts) > 1, chunk_counts vectors.delete_index(**wide) vectors.delete_index(**args) vectors.delete_vector_bucket(vectorBucketName=bucket) 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 random import urllib.request 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 coverage = run_cli(cluster, "coverage", f"{bucket}$vector${index}") assert len(coverage) == 1 and coverage[0]["segments"] == 1 and coverage[0]["unindexed_rows"] == 0, coverage # A table named by its s3:// location is opened on the cluster as well. uri = run_cli(cluster, "describe", f"{bucket}$vector${index}")["uri"] assert uri.startswith("s3://"), uri direct = run_cli(cluster, "coverage", uri) assert direct == coverage, (direct, 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 run_cli(cluster, "coverage", f"{bucket}$vector${index}")[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): """A gateway with no S3 identities serves an index like one with.""" 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"]