"""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"]