diff --git a/pyproject.toml b/pyproject.toml index 9fbbb7748..3ad62baae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,6 +70,7 @@ redis = [ "redis" ] memorydb = [ "memorydb" ] chromadb = [ "chromadb" ] opensearch = [ "opensearch-py", "boto3", "requests-aws4auth" ] +dynamodb = [ "boto3>=1.40.0" ] aliyun_opensearch = [ "alibabacloud_ha3engine_vector" ] mongodb = [ "pymongo" ] mariadb = [ "mariadb" ] diff --git a/tests/test_dynamodb.py b/tests/test_dynamodb.py new file mode 100644 index 000000000..3ee5af213 --- /dev/null +++ b/tests/test_dynamodb.py @@ -0,0 +1,245 @@ +"""Offline unit tests for the DynamoDB vector client. + +These tests do not require AWS credentials or a running DynamoDB. They freeze +the config contract (metric parsing, credential passthrough, to_dict shape) +and the CLI wiring so a future refactor that breaks them fails CI. +""" + +import ast +from pathlib import Path + +import pytest +from pydantic import SecretStr + +import vectordb_bench.backend.clients.dynamodb.dynamodb as ddb_mod +from vectordb_bench.backend.clients import DB +from vectordb_bench.backend.clients.api import MetricType +from vectordb_bench.backend.clients.dynamodb.config import ( + DynamoDBConfig, + DynamoDBIndexConfig, +) +from vectordb_bench.backend.clients.dynamodb.dynamodb import DynamoDB as DynamoDBClient +from vectordb_bench.backend.filter import LabelFilter, NewIntFilter, NonFilter + +# --------------------------------------------------------------------------- +# Enum registration +# --------------------------------------------------------------------------- + + +def test_dynamodb_enum_resolves_config_and_init(): + assert DB.DynamoDB.value == "DynamoDB" + assert DB.DynamoDB.config_cls is DynamoDBConfig + assert DB.DynamoDB.init_cls is DynamoDBClient + + +# --------------------------------------------------------------------------- +# Metric parsing +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("metric", "expected"), + [ + (MetricType.COSINE, "COSINE"), + (MetricType.L2, "EUCLIDEAN"), + (MetricType.IP, "DOT_PRODUCT"), + ], +) +def test_metric_parsing(metric: MetricType, expected: str): + assert DynamoDBIndexConfig(metric_type=metric).parse_metric() == expected + + +def test_metric_parsing_rejects_unset(): + with pytest.raises(ValueError, match="Unsupported metric"): + DynamoDBIndexConfig(metric_type=None).parse_metric() + + +def test_index_param_carries_partitions_and_search_param_empty(): + cfg = DynamoDBIndexConfig(metric_type=MetricType.COSINE) + assert cfg.index_param() == {"num_partitions": 1} + assert cfg.search_param() == {} + + +# --------------------------------------------------------------------------- +# Partition key (SearchSchema HASH) +# --------------------------------------------------------------------------- + + +def test_use_partition_key_toggles_on_count(): + assert DynamoDBIndexConfig(num_partitions=1).use_partition_key() is False + assert DynamoDBIndexConfig(num_partitions=4).use_partition_key() is True + + +def test_search_schema_elements_add_hash_and_inline_filter(): + client = DynamoDBClient.__new__(DynamoDBClient) + client.use_partition_key = True + client.with_scalar_labels = True + + types = [e["SearchSchemaElementType"] for e in client._search_schema_elements()] + assert types == ["HASH", "INLINE_FILTER"] + + client.use_partition_key = False + client.with_scalar_labels = False + assert client._search_schema_elements() == [] + + +def test_prepare_filter_builds_equality_condition(): + client = DynamoDBClient.__new__(DynamoDBClient) + + client.prepare_filter(NonFilter()) + assert client._condition_expr is None + assert client._expr_values == {} + + lbl = LabelFilter(label_percentage=0.2) + client.prepare_filter(lbl) + assert client._condition_expr == "label = :label" + assert client._expr_values == {":label": {"S": lbl.label_value}} + + +def test_prepare_filter_rejects_numge(): + client = DynamoDBClient.__new__(DynamoDBClient) + with pytest.raises(ValueError, match="Unsupported filter"): + client.prepare_filter(NewIntFilter(int_value=100, filter_rate=0.01)) + + +def test_build_item_sets_partition_and_label(): + client = DynamoDBClient.__new__(DynamoDBClient) + client.use_partition_key = True + client.num_partitions = 3 + client.with_scalar_labels = True + + item = client._build_item(7, [0.1, 0.2], ["lbl"], 0) + assert item["pk"] == {"N": "7"} + assert item["id"] == {"N": "7"} + assert item["vector"] == {"L": [{"N": "0.1"}, {"N": "0.2"}]} + assert item["part"] == {"N": "1"} # 7 % 3 + assert item["label"] == {"S": "lbl"} + + +def test_partition_search_scopes_to_one_random_partition(): + """With a partition key, a search issues ONE scoped SearchVectors call + (not a fan-out), filtered to a single partition value.""" + captured = {} + + class _SearchDDB: + def search_vectors(self, **params): + captured.update(params) + return {"SearchResults": [{"Item": {"id": {"N": "5"}}, "Score": 0.1}]} + + client = DynamoDBClient.__new__(DynamoDBClient) + client.client = _SearchDDB() + client.table_name = "t" + client.index_name = "i" + client.use_partition_key = True + client.num_partitions = 4 + client._condition_expr = None + client._expr_values = {} + + out = client.search_embedding([0.1, 0.2], k=10) + assert out == [5] + # exactly one partition equality in the condition + assert captured["SearchConditionExpression"].startswith("part = :part") + assert "AND" not in captured["SearchConditionExpression"] # no inline filter here + part_val = int(captured["ExpressionAttributeValues"][":part"]["N"]) + assert 0 <= part_val < 4 + + +# --------------------------------------------------------------------------- +# Throttling: UnprocessedItems backoff re-drive +# --------------------------------------------------------------------------- + + +class _FakeDDB: + """Returns the given queue of UnprocessedItems maps, one per call.""" + + def __init__(self, unprocessed_sequence: list[dict]): + self._seq = list(unprocessed_sequence) + self.calls = 0 + + def batch_write_item(self, RequestItems: dict) -> dict: # noqa: N803 (boto3 arg name) + self.calls += 1 + up = self._seq.pop(0) if self._seq else {} + return {"UnprocessedItems": up} + + +def test_write_batch_drains_after_throttled_redrives(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(ddb_mod.time, "sleep", lambda _s: None) # no real waiting + + client = DynamoDBClient.__new__(DynamoDBClient) + one = {"t": [{"PutRequest": {"Item": {"pk": {"N": "1"}}}}]} + # first call -> 1 unprocessed, second -> 1 unprocessed, third -> drained + client.client = _FakeDDB([one, one, {}]) + + client._write_batch_with_retry(one) + assert client.client.calls == 3 # 1 initial + 2 re-drives + + +def test_write_batch_gives_up_after_max_retries(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(ddb_mod.time, "sleep", lambda _s: None) + + client = DynamoDBClient.__new__(DynamoDBClient) + one = {"t": [{"PutRequest": {"Item": {"pk": {"N": "1"}}}}]} + # always throttled -> never drains + client.client = _FakeDDB([one] * 50) + + with pytest.raises(RuntimeError, match="unprocessed items after"): + client._write_batch_with_retry(one) + assert client.client.calls == ddb_mod.UNPROCESSED_MAX_RETRIES + 1 + + +# --------------------------------------------------------------------------- +# Config to_dict() contract +# --------------------------------------------------------------------------- + + +def test_to_dict_without_explicit_credentials(): + """Omitting keys must leave them None so boto3's default chain is used.""" + cfg = DynamoDBConfig(region_name="us-west-2", table_name="t", index_name="i") + d = cfg.to_dict() + assert d["region_name"] == "us-west-2" + assert d["access_key_id"] is None + assert d["secret_access_key"] is None + assert d["session_token"] is None + assert d["table_name"] == "t" + assert d["index_name"] == "i" + + +def test_to_dict_unwraps_secrets(): + cfg = DynamoDBConfig( + access_key_id=SecretStr("AKIA_EXAMPLE"), + secret_access_key=SecretStr("secret_example"), + session_token=SecretStr("token_example"), + ) + d = cfg.to_dict() + assert d["access_key_id"] == "AKIA_EXAMPLE" + assert d["secret_access_key"] == "secret_example" # noqa: S105 + assert d["session_token"] == "token_example" # noqa: S105 + + +# --------------------------------------------------------------------------- +# CLI wiring (static source introspection, no heavy runtime import) +# --------------------------------------------------------------------------- + + +def test_cli_defines_command_and_options(): + cli_src = Path("vectordb_bench/backend/clients/dynamodb/cli.py").read_text() + tree = ast.parse(cli_src) + + class_names = {n.name for n in ast.walk(tree) if isinstance(n, ast.ClassDef)} + func_names = {n.name for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)} + + assert "DynamoDBTypedDict" in class_names + assert "DynamoDBIndexTypedDict" in class_names + assert "DynamoDB" in func_names + + def _fields(cls_name: str) -> set[str]: + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and node.name == cls_name: + return { + i.target.id + for i in node.body + if isinstance(i, ast.AnnAssign) and isinstance(i.target, ast.Name) + } + return set() + + assert "partition_count" in _fields("DynamoDBTypedDict") diff --git a/vectordb_bench/backend/clients/__init__.py b/vectordb_bench/backend/clients/__init__.py index dbd83bb66..7f03410ad 100644 --- a/vectordb_bench/backend/clients/__init__.py +++ b/vectordb_bench/backend/clients/__init__.py @@ -52,6 +52,7 @@ class DB(Enum): LanceDB = "LanceDB" OceanBase = "OceanBase" S3Vectors = "S3Vectors" + DynamoDB = "DynamoDB" Hologres = "Alibaba Cloud Hologres" TencentElasticsearch = "TencentElasticsearch" AliSQL = "AlibabaCloudRDSMySQL" @@ -228,6 +229,11 @@ def init_cls(self) -> type[VectorDB]: # noqa: PLR0911, PLR0912, C901, PLR0915 return S3Vectors + if self == DB.DynamoDB: + from .dynamodb.dynamodb import DynamoDB + + return DynamoDB + if self == DB.Hologres: from .hologres.hologres import Hologres @@ -451,6 +457,11 @@ def config_cls(self) -> type[DBConfig]: # noqa: PLR0911, PLR0912, C901, PLR0915 return S3VectorsConfig + if self == DB.DynamoDB: + from .dynamodb.config import DynamoDBConfig + + return DynamoDBConfig + if self == DB.Hologres: from .hologres.config import HologresConfig @@ -657,6 +668,10 @@ def case_config_cls( # noqa: C901, PLR0911, PLR0912, PLR0915 from .s3_vectors.config import S3VectorsIndexConfig return S3VectorsIndexConfig + if self == DB.DynamoDB: + from .dynamodb.config import DynamoDBIndexConfig + + return DynamoDBIndexConfig if self == DB.Hologres: from .hologres.config import HologresIndexConfig diff --git a/vectordb_bench/backend/clients/dynamodb/__init__.py b/vectordb_bench/backend/clients/dynamodb/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vectordb_bench/backend/clients/dynamodb/cli.py b/vectordb_bench/backend/clients/dynamodb/cli.py new file mode 100644 index 000000000..26ad568c2 --- /dev/null +++ b/vectordb_bench/backend/clients/dynamodb/cli.py @@ -0,0 +1,118 @@ +from typing import Annotated, TypedDict, Unpack + +import click +from pydantic import SecretStr + +from ....cli.cli import ( + CommonTypedDict, + cli, + click_parameter_decorators_from_typed_dict, + get_custom_case_config, + run, +) +from .. import DB +from ..api import MetricType + + +class DynamoDBTypedDict(TypedDict): + region_name: Annotated[ + str, + click.option("--region", type=str, help="AWS region (e.g. us-east-1)", default="us-east-1"), + ] + access_key_id: Annotated[ + str, + click.option( + "--access_key_id", + type=str, + help="AWS access key ID. Omit to use the default boto3 credential chain.", + default=None, + ), + ] + secret_access_key: Annotated[ + str, + click.option( + "--secret_access_key", + type=str, + help="AWS secret access key. Omit to use the default boto3 credential chain.", + default=None, + ), + ] + session_token: Annotated[ + str, + click.option( + "--session_token", + type=str, + help="Optional AWS session token for temporary credentials.", + default=None, + ), + ] + table: Annotated[ + str, + click.option("--table", type=str, help="DynamoDB table name", default="vdbbench_vectors"), + ] + index: Annotated[ + str, + click.option("--index", type=str, help="Vector index name", default="vdbbench-index"), + ] + metric: Annotated[ + str, + click.option( + "--metric", + type=str, + help="Distance metric: cosine, euclidean (l2), or dotproduct (ip).", + default="cosine", + ), + ] + partition_count: Annotated[ + int, + click.option( + "--partition-count", + type=int, + help=( + "Number of vector-index partition-key (SearchSchema HASH) values. " + ">1 defines a partition key; each search is scoped to one randomly " + "chosen partition value (realistic single-partition access, recall " + "~1/N vs whole-dataset ground truth). 1 (default) searches the whole " + "index with no partition key." + ), + default=1, + ), + ] + + +class DynamoDBIndexTypedDict(CommonTypedDict, DynamoDBTypedDict): ... + + +_METRIC_MAP = { + "cosine": MetricType.COSINE, + "euclidean": MetricType.L2, + "l2": MetricType.L2, + "dotproduct": MetricType.IP, + "ip": MetricType.IP, +} + + +@cli.command() +@click_parameter_decorators_from_typed_dict(DynamoDBIndexTypedDict) +def DynamoDB(**parameters: Unpack[DynamoDBIndexTypedDict]): + from .config import DynamoDBConfig, DynamoDBIndexConfig + + parameters["custom_case"] = get_custom_case_config(parameters) + run( + db=DB.DynamoDB, + db_config=DynamoDBConfig( + region_name=parameters["region"], + access_key_id=(SecretStr(parameters["access_key_id"]) if parameters["access_key_id"] else None), + secret_access_key=( + SecretStr(parameters["secret_access_key"]) if parameters["secret_access_key"] else None + ), + session_token=(SecretStr(parameters["session_token"]) if parameters["session_token"] else None), + table_name=parameters["table"], + index_name=parameters["index"] if parameters["index"] else "vdbbench-index", + ), + db_case_config=DynamoDBIndexConfig( + metric_type=_METRIC_MAP.get((parameters["metric"] or "cosine").lower()), + num_partitions=parameters.get("partition_count", 1) or 1, + ), + **parameters, + ) diff --git a/vectordb_bench/backend/clients/dynamodb/config.py b/vectordb_bench/backend/clients/dynamodb/config.py new file mode 100644 index 000000000..b529e3a13 --- /dev/null +++ b/vectordb_bench/backend/clients/dynamodb/config.py @@ -0,0 +1,79 @@ +from pydantic import BaseModel, SecretStr + +from ..api import DBCaseConfig, DBConfig, MetricType + + +class DynamoDBConfig(DBConfig): + """Connection config for Amazon DynamoDB vector search. + + Credentials are optional: when access_key_id / secret_access_key are left + empty the standard boto3 credential chain is used (environment variables, + shared config/credentials files, or an attached IAM role). + """ + + region_name: str = "us-east-1" + access_key_id: SecretStr | None = None + secret_access_key: SecretStr | None = None + session_token: SecretStr | None = None + table_name: str = "vdbbench_vectors" + index_name: str = "vdbbench-index" + + def to_dict(self) -> dict: + return { + "region_name": self.region_name, + "access_key_id": (self.access_key_id.get_secret_value() if self.access_key_id else None), + "secret_access_key": (self.secret_access_key.get_secret_value() if self.secret_access_key else None), + "session_token": (self.session_token.get_secret_value() if self.session_token else None), + "table_name": self.table_name, + "index_name": self.index_name, + } + + +class DynamoDBIndexConfig(DBCaseConfig, BaseModel): + """Case config for a DynamoDB vector index. + + DynamoDB builds and manages the vector index server-side, so there are no + client-tunable ANN build/search params. What is configurable is the + SearchSchema: + + - metric_type: the distance function (set at index creation). + - num_partitions: when > 1, the client defines a SearchSchema HASH key + (vector index partition key) and spreads vectors UNIFORMLY across this + many partition values (``id % num_partitions``). This is a balanced + baseline: every partition holds an equal share, so there are no hot or + empty partitions. Real-world partition keys (Category, tenant, ...) are + usually skewed, so these numbers represent the best-case, evenly-balanced + scenario rather than skewed production behaviour. Each SearchVectors call + is scoped to ONE partition value (chosen at random per query), the + realistic single-partition access pattern the feature is built for. This + searches only ~1/num_partitions of the data, so recall measured against + VectorDBBench's whole-dataset ground truth is expected to be about + 1/num_partitions; QPS/latency reflect the scoped search. Use 1 (default) + for a whole-index recall benchmark with no partition key. + + A vector index partition key must be a low-to-medium cardinality attribute + and, once defined, its value is REQUIRED in the SearchConditionExpression + of every SearchVectors call (per the DynamoDB vector search docs). + """ + + metric_type: MetricType | None = None + num_partitions: int = 1 + + def parse_metric(self) -> str: + if self.metric_type == MetricType.COSINE: + return "COSINE" + if self.metric_type == MetricType.L2: + return "EUCLIDEAN" + if self.metric_type == MetricType.IP: + return "DOT_PRODUCT" + msg = f"Unsupported metric type for DynamoDB: {self.metric_type}" + raise ValueError(msg) + + def use_partition_key(self) -> bool: + return self.num_partitions > 1 + + def index_param(self) -> dict: + return {"num_partitions": self.num_partitions} + + def search_param(self) -> dict: + return {} diff --git a/vectordb_bench/backend/clients/dynamodb/dynamodb.py b/vectordb_bench/backend/clients/dynamodb/dynamodb.py new file mode 100644 index 000000000..ae135a27b --- /dev/null +++ b/vectordb_bench/backend/clients/dynamodb/dynamodb.py @@ -0,0 +1,371 @@ +"""Wrapper around Amazon DynamoDB vector search over VectorDB. + +Uses only the public, GA DynamoDB vector API via the standard boto3 SDK: + - create_table(..., VectorIndexes=[...]) with a SearchSchema to create a + table and vector index + - batch_write_item to load vectors + - search_vectors to run approximate nearest-neighbour queries + +Credentials come from the standard boto3 credential chain unless explicit keys +are given. + +SearchSchema (see the DynamoDB vector search docs). The vector attribute is +named by VectorAttribute, not listed in SearchSchema; the schema holds only: + - HASH: an optional vector index partition key. When present it + scopes each SearchVectors call to vectors that share its + value, which is how search throughput scales over a large + index. Its value is REQUIRED in SearchConditionExpression + on every search. Enabled here via num_partitions > 1. + - INLINE_FILTER: an optional non-vector attribute stored next to the vector + so equality filters are applied during the search. Used for + the StrEqual (label) benchmark filter. Inline filters and + the HASH key support ONLY the equality operator (=). +""" + +import logging +import random +import time +from collections.abc import Iterable +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any + +import boto3 +from botocore.config import Config +from botocore.exceptions import ClientError + +from vectordb_bench.backend.filter import Filter, FilterOp + +from ..api import VectorDB +from .config import DynamoDBIndexConfig + +if TYPE_CHECKING: + from botocore.client import BaseClient + +log = logging.getLogger(__name__) + +WRITE_BATCH_MAX_SIZE = 25 # DynamoDB BatchWriteItem hard limit +UNPROCESSED_MAX_RETRIES = 10 # extra re-drives for throttled UnprocessedItems +UNPROCESSED_BASE_DELAY = 0.05 # seconds; exponential-backoff base +UNPROCESSED_MAX_DELAY = 20.0 # seconds; backoff ceiling +TOP_K_MAX = 100 # SearchVectors TopK valid range is 1..100 + +_ID_FIELD = "id" +_LABEL_FIELD = "label" +_VECTOR_FIELD = "vector" +_PK_FIELD = "pk" +_PARTITION_FIELD = "part" # vector index partition key (SearchSchema HASH) + + +class DynamoDB(VectorDB): + supported_filter_types: list[FilterOp] = [ + FilterOp.NonFilter, + FilterOp.StrEqual, + ] + + def __init__( + self, + dim: int, + db_config: dict, + db_case_config: DynamoDBIndexConfig, + drop_old: bool = False, + with_scalar_labels: bool = False, + **kwargs, + ): + self.dim = dim + self.db_config = db_config + self.case_config = db_case_config + self.with_scalar_labels = with_scalar_labels + + self.table_name = db_config.get("table_name") + self.index_name = db_config.get("index_name") + self.num_partitions = max(1, int(db_case_config.num_partitions)) + self.use_partition_key = db_case_config.use_partition_key() + + # Prepared search condition, set by prepare_filter(). + self._condition_expr: str | None = None + self._expr_values: dict[str, Any] = {} + + client = self._new_client() + try: + if drop_old: + self._drop_table(client) + self._create_table(client, dim) + finally: + client.close() + + # ------------------------------------------------------------------ # + # Connection + # ------------------------------------------------------------------ # + def _new_client(self) -> "BaseClient": + cfg = self.db_config + kwargs: dict[str, Any] = { + "service_name": "dynamodb", + "region_name": cfg.get("region_name"), + "config": Config( + retries={"max_attempts": 8, "mode": "standard"}, + max_pool_connections=100, + ), + } + if cfg.get("access_key_id") and cfg.get("secret_access_key"): + kwargs["aws_access_key_id"] = cfg["access_key_id"] + kwargs["aws_secret_access_key"] = cfg["secret_access_key"] + if cfg.get("session_token"): + kwargs["aws_session_token"] = cfg["session_token"] + return boto3.client(**kwargs) + + # ------------------------------------------------------------------ # + # Table / index lifecycle + # ------------------------------------------------------------------ # + def _drop_table(self, client: "BaseClient") -> None: + try: + log.info(f"DynamoDB dropping old table: {self.table_name}") + client.delete_table(TableName=self.table_name) + client.get_waiter("table_not_exists").wait(TableName=self.table_name) + log.info(f"DynamoDB dropped table: {self.table_name}") + except ClientError as error: + if error.response["Error"]["Code"] == "ResourceNotFoundException": + log.info(f"DynamoDB table does not exist, nothing to drop: {self.table_name}") + else: + raise + + def _search_schema_elements(self) -> list[dict[str, str]]: + """SearchSchema holds only the optional partition key (HASH) and inline + filters. The vector attribute itself is named by VectorAttribute, not + listed here; there is no VECTOR element type in the API.""" + elements: list[dict[str, str]] = [] + if self.use_partition_key: + # At most one HASH (vector index partition key) is allowed. + elements.append({"AttributeName": _PARTITION_FIELD, "SearchSchemaElementType": "HASH"}) + if self.with_scalar_labels: + elements.append({"AttributeName": _LABEL_FIELD, "SearchSchemaElementType": "INLINE_FILTER"}) + return elements + + def _create_table(self, client: "BaseClient", dim: int) -> None: + vector_index: dict[str, Any] = { + "IndexName": self.index_name, + "VectorAttribute": {"AttributeName": _VECTOR_FIELD}, + "Dimensions": dim, + "DistanceFunction": self.case_config.parse_metric(), + "Projection": {"ProjectionType": "ALL"}, + } + search_schema = self._search_schema_elements() + if search_schema: + vector_index["SearchSchema"] = search_schema + + attribute_definitions = [{"AttributeName": _PK_FIELD, "AttributeType": "N"}] + if self.use_partition_key: + # A SearchSchema HASH attribute must also be declared in the table's + # AttributeDefinitions (same rule as secondary-index key attributes). + attribute_definitions.append({"AttributeName": _PARTITION_FIELD, "AttributeType": "N"}) + + try: + log.info( + f"DynamoDB creating table: {self.table_name} " + f"(dim={dim}, partitions={self.num_partitions})" + ) + client.create_table( + TableName=self.table_name, + KeySchema=[{"AttributeName": _PK_FIELD, "KeyType": "HASH"}], + AttributeDefinitions=attribute_definitions, + BillingMode="PAY_PER_REQUEST", # vector indexes require on-demand + VectorIndexes=[vector_index], + ) + client.get_waiter("table_exists").wait(TableName=self.table_name) + log.info(f"DynamoDB table active: {self.table_name}") + except ClientError as error: + if error.response["Error"]["Code"] == "ResourceInUseException": + log.info(f"DynamoDB table already exists: {self.table_name}") + client.get_waiter("table_exists").wait(TableName=self.table_name) + else: + raise + + @contextmanager + def init(self): + """Create and destroy the boto3 client for a worker process. + + Examples: + >>> with self.init(): + >>> self.insert_embeddings() + >>> self.search_embedding() + """ + self.client = self._new_client() + yield + self.client.close() + + def optimize(self, **kwargs): + """No-op: DynamoDB builds the vector index server-side.""" + return + + def need_normalize_cosine(self) -> bool: + return False + + # ------------------------------------------------------------------ # + # Load + # ------------------------------------------------------------------ # + def insert_embeddings( + self, + embeddings: Iterable[list[float]], + metadata: list[int], + labels_data: list[str] | None = None, + **kwargs, + ) -> tuple[int, Exception | None]: + assert self.client is not None + embeddings = list(embeddings) + assert len(embeddings) == len(metadata) + + insert_count = 0 + try: + for start in range(0, len(embeddings), WRITE_BATCH_MAX_SIZE): + end = min(start + WRITE_BATCH_MAX_SIZE, len(embeddings)) + request_items = { + self.table_name: [ + {"PutRequest": {"Item": self._build_item(metadata[i], embeddings[i], labels_data, i)}} + for i in range(start, end) + ] + } + self._write_batch_with_retry(request_items) + insert_count += end - start + except Exception as e: + log.warning(f"DynamoDB failed to insert data: {e}") + return insert_count, e + return insert_count, None + + def _build_item( + self, + row_id: int, + embedding: list[float], + labels_data: list[str] | None, + i: int, + ) -> dict[str, Any]: + item: dict[str, Any] = { + _PK_FIELD: {"N": str(row_id)}, + _ID_FIELD: {"N": str(row_id)}, + _VECTOR_FIELD: {"L": [{"N": str(v)} for v in embedding]}, + } + if self.use_partition_key: + # Uniform (balanced-baseline) distribution: id % N spreads rows + # evenly across partitions. Real-world keys are usually skewed, so + # this is the best-case, no-hot-partition scenario. + item[_PARTITION_FIELD] = {"N": str(row_id % self.num_partitions)} + if self.with_scalar_labels and labels_data is not None: + item[_LABEL_FIELD] = {"S": labels_data[i]} + return item + + def _write_batch_with_retry(self, request_items: dict[str, list[dict]]) -> None: + """Write one batch, re-driving UnprocessedItems with bounded exponential + backoff. + + Two layers of throttling defence: + 1. boto3 ``standard`` retry mode (max_attempts=8) retries throttling + exceptions and 5xx on each BatchWriteItem call with its own + backoff + jitter. + 2. BatchWriteItem can also succeed (HTTP 200) while returning some + rows in UnprocessedItems when the table/index is being throttled. + Those are NOT retried by the SDK, so we re-drive them here. Under + heavy throttling a large fraction comes back unprocessed, so we + back off exponentially (with jitter) between re-drives instead of + hot-looping, and give up after UNPROCESSED_MAX_RETRIES. + """ + response = self.client.batch_write_item(RequestItems=request_items) + unprocessed = response.get("UnprocessedItems") or {} + + attempt = 0 + while unprocessed: + pending = sum(len(v) for v in unprocessed.values()) + if attempt >= UNPROCESSED_MAX_RETRIES: + msg = ( + f"BatchWriteItem still has {pending} unprocessed items after " + f"{UNPROCESSED_MAX_RETRIES} backoff retries; table is throttling " + f"faster than it can drain." + ) + raise RuntimeError(msg) + + delay = min(UNPROCESSED_MAX_DELAY, UNPROCESSED_BASE_DELAY * (2**attempt)) + delay += random.uniform(0, delay) # full jitter + log.debug( + f"DynamoDB BatchWriteItem throttled: {pending} unprocessed, " + f"re-drive {attempt + 1}/{UNPROCESSED_MAX_RETRIES} after {delay:.2f}s" + ) + time.sleep(delay) + + response = self.client.batch_write_item(RequestItems=unprocessed) + unprocessed = response.get("UnprocessedItems") or {} + attempt += 1 + + # ------------------------------------------------------------------ # + # Search + # ------------------------------------------------------------------ # + def prepare_filter(self, filters: Filter): + """Build the SearchConditionExpression for the next searches. + + DynamoDB vector search supports ONLY equality (=) in a + SearchConditionExpression, on HASH and INLINE_FILTER attributes. The + StrEqual benchmark filter maps to an equality on the inline-filter + label attribute; NonFilter clears it. NumGE (id >= N) is intentionally + not advertised in supported_filter_types because the API has no range + operator for the search condition. + """ + if filters.type == FilterOp.NonFilter: + self._condition_expr = None + self._expr_values = {} + elif filters.type == FilterOp.StrEqual: + self._condition_expr = f"{_LABEL_FIELD} = :label" + self._expr_values = {":label": {"S": filters.label_value}} + else: + msg = f"Unsupported filter for DynamoDB vector search - {filters}" + raise ValueError(msg) + + def search_embedding( + self, + query: list[float], + k: int = 100, + timeout: int | None = None, + ) -> list[int]: + assert self.client is not None + top_k = min(k, TOP_K_MAX) + + condition = self._condition_expr + values = dict(self._expr_values) + if self.use_partition_key: + # A SearchVectors call is scoped to ONE partition-key value, which + # is how the feature is used in practice: the caller searches a + # single known partition rather than the whole index. We pick a + # random partition value per query to spread search load evenly + # across partitions (mirroring the id % N write distribution). + # + # NOTE: this searches only ~1/num_partitions of the dataset, while + # VectorDBBench computes recall against WHOLE-dataset ground truth. + # Recall is therefore expected to be ~1/num_partitions here; the + # QPS/latency figures reflect the realistic single-partition access + # pattern. Use num_partitions=1 for a whole-index recall benchmark. + part = random.randrange(self.num_partitions) + part_cond = f"{_PARTITION_FIELD} = :part" + values[":part"] = {"N": str(part)} + condition = f"{part_cond} AND {condition}" if condition else part_cond + + results = self._search_one(query, top_k, condition=condition, values=values) + return [row_id for row_id, _ in results] + + def _search_one( + self, + query: list[float], + top_k: int, + condition: str | None, + values: dict[str, Any], + ) -> list[tuple[int, float]]: + params: dict[str, Any] = { + "TableName": self.table_name, + "IndexName": self.index_name, + "SearchVector": [{"N": str(v)} for v in query], + "TopK": top_k, + "ProjectionExpression": _ID_FIELD, + } + if condition: + params["SearchConditionExpression"] = condition + params["ExpressionAttributeValues"] = values + + response = self.client.search_vectors(**params) + return [ + (int(r["Item"][_ID_FIELD]["N"]), float(r.get("Score", 0.0))) + for r in response.get("SearchResults", []) + ] diff --git a/vectordb_bench/cli/vectordbbench.py b/vectordb_bench/cli/vectordbbench.py index 1bbc462ef..4299358f3 100644 --- a/vectordb_bench/cli/vectordbbench.py +++ b/vectordb_bench/cli/vectordbbench.py @@ -45,6 +45,7 @@ from ..backend.clients.qdrant_local.cli import QdrantLocal from ..backend.clients.redis.cli import Redis from ..backend.clients.s3_vectors.cli import S3Vectors +from ..backend.clients.dynamodb.cli import DynamoDB from ..backend.clients.seekdb.cli import SeekDBHNSW from ..backend.clients.tencent_elasticsearch.cli import TencentElasticsearch from ..backend.clients.test.cli import Test @@ -97,6 +98,7 @@ cli.add_command(ElasticCloudHNSWBBQ) cli.add_command(BatchCli) cli.add_command(S3Vectors) +cli.add_command(DynamoDB) cli.add_command(TencentElasticsearch) cli.add_command(AliSQLHNSW) cli.add_command(Doris)