diff --git a/README.md b/README.md index 85d34d76d..b957a966f 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,7 @@ All the database client supported | pgvector, pgvectorscale, pgdiskann, alloydb, vectorchord | `pip install vectordb-bench[pgvector]` | | pgvecto.rs | `pip install vectordb-bench[pgvecto_rs]` | | redis | `pip install vectordb-bench[redis]` | +| valkey | `pip install vectordb-bench[valkey]` | | memorydb | `pip install vectordb-bench[memorydb]` | | chromadb | `pip install vectordb-bench[chromadb]` | | cockroachdb | `pip install vectordb-bench[cockroachdb]` | diff --git a/install/requirements_py3.11.txt b/install/requirements_py3.11.txt index a6f3f32cf..0a4fe571d 100644 --- a/install/requirements_py3.11.txt +++ b/install/requirements_py3.11.txt @@ -11,6 +11,7 @@ pgvector pgvecto_rs[psycopg3]>=0.2.1 sqlalchemy redis +valkey-glide-sync>=2.5.1 chromadb pytz streamlit-autorefresh diff --git a/pyproject.toml b/pyproject.toml index 223291721..33c89c087 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,7 @@ pgvector = [ "psycopg", "psycopg-binary", "pgvector" ] pgvecto_rs = [ "pgvecto_rs[psycopg3]>=0.2.2" ] redis = [ "redis" ] +valkey = [ "valkey-glide-sync>=2.5.1" ] memorydb = [ "memorydb" ] chromadb = [ "chromadb" ] opensearch = [ "opensearch-py", "boto3", "requests-aws4auth" ] diff --git a/tests/test_valkey.py b/tests/test_valkey.py new file mode 100644 index 000000000..2433060dd --- /dev/null +++ b/tests/test_valkey.py @@ -0,0 +1,298 @@ +from typing import Any +from unittest.mock import patch + +import numpy as np +import pytest +from glide_sync import Batch +from pydantic import ValidationError + +from vectordb_bench.backend.clients import DB +from vectordb_bench.backend.clients.api import IndexType, MetricType +from vectordb_bench.backend.clients.valkey.config import ValkeyConfig, ValkeyHNSWConfig +from vectordb_bench.backend.clients.valkey.valkey import Valkey +from vectordb_bench.backend.filter import IntFilter, LabelFilter, non_filter + + +def command_name(command: Any) -> str: + command = getattr(command, "value", command) + return command.decode() if isinstance(command, bytes) else str(command) + + +class FakeClusterCursor: + def __init__(self): + self.finished = False + + def is_finished(self): + return self.finished + + +class FakeClient: + def __init__(self): + self.created = False + self.closed = False + self.custom_commands = [] + self.exec_batches = [] + self.scan_keys = [] + + def custom_command(self, args: list[Any]): + self.custom_commands.append(args) + command = command_name(args[0]) + if command == "FT._LIST": + return [b"vdbbench_valkey"] if self.created else [] + if command == "FT.CREATE": + self.created = True + return b"OK" + if command == "FT.DROPINDEX": + self.created = False + return b"OK" + if command == "FT.SEARCH": + return [1, {b"vdbbench_valkey:7": {}}] + msg = f"Unexpected command: {args}" + raise AssertionError(msg) + + def exec(self, batch: Any, raise_on_error: bool): + assert raise_on_error is True + self.exec_batches.append(batch.commands) + return [1] * len(batch.commands) + + def scan(self, cursor: Any, match: str, count: int): + assert match == "vdbbench_valkey:*" + assert count > 0 + keys, self.scan_keys = self.scan_keys, [] + if isinstance(cursor, FakeClusterCursor): + cursor.finished = True + return [cursor, keys] + return [b"0", keys] + + def close(self): + self.closed = True + + +def make_adapter(): + setup_client = FakeClient() + runtime_client = FakeClient() + config = ValkeyHNSWConfig(M=24, efConstruction=300, ef=40) + db_config = { + "host": "localhost", + "port": 6379, + "password": None, + "ssl": False, + "insecure_tls": False, + "request_timeout_ms": 12_000, + "connection_timeout_ms": 13_000, + "cmd": True, + } + with patch( + "vectordb_bench.backend.clients.valkey.valkey.GlideClient.create", + side_effect=[setup_client, runtime_client], + ) as create: + adapter = Valkey(dim=3, db_config=db_config, db_case_config=config) + context = adapter.init() + context.__enter__() + return adapter, context, setup_client, runtime_client, create + + +def test_valkey_registration_and_config(): + assert DB.Valkey.value == "Valkey" + assert DB.Valkey.init_cls is Valkey + assert DB.Valkey.config_cls is ValkeyConfig + assert DB.Valkey.case_config_cls(IndexType.HNSW) is ValkeyHNSWConfig + assert DB.Valkey.case_config_cls(IndexType.AUTOINDEX) is ValkeyHNSWConfig + assert Valkey.supports_full_text_search() is False + + config = ValkeyConfig(host="localhost") + assert config.to_dict()["password"] is None + assert config.to_dict()["port"] == 6379 + assert config.to_dict()["ssl"] is True + assert config.to_dict()["insecure_tls"] is False + assert config.to_dict()["cmd"] is False + assert config.to_dict()["request_timeout_ms"] == 600_000 + assert config.to_dict()["connection_timeout_ms"] == 10_000 + assert config.to_dict()["collection_name"] == "vdbbench_valkey" + + assert ValkeyHNSWConfig().model_dump() == { + "metric_type": None, + "M": 16, + "efConstruction": 200, + "ef": 10, + "index": IndexType.HNSW, + } + with pytest.raises(ValidationError): + ValkeyHNSWConfig(M=0, efConstruction=200) + with pytest.raises(ValidationError): + ValkeyHNSWConfig(M=16, efConstruction=200, ef=0) + with pytest.raises(ValidationError): + ValkeyConfig(host="localhost", collection_name="unsafe*") + with pytest.raises(ValidationError): + ValkeyConfig(host="localhost", port=0) + with pytest.raises(ValidationError): + ValkeyConfig(host="localhost", request_timeout_ms=0) + with pytest.raises(ValidationError): + ValkeyConfig(host="localhost", connection_timeout_ms=0) + with pytest.raises(ValidationError, match="requires ssl"): + ValkeyConfig(host="localhost", ssl=False, insecure_tls=True) + + +def test_valkey_index_insert_and_search(): + adapter, context, setup_client, runtime_client, create = make_adapter() + try: + assert setup_client.closed is True + assert setup_client.created is True + create_args = next(args for args in setup_client.custom_commands if command_name(args[0]) == "FT.CREATE") + assert create_args[1:7] == [ + "vdbbench_valkey", + "ON", + "HASH", + "PREFIX", + "1", + "vdbbench_valkey:", + ] + + glide_config = create.call_args_list[0].args[0] + assert glide_config.addresses[0].host == "localhost" + assert glide_config.addresses[0].port == 6379 + assert glide_config.use_tls is False + assert glide_config.request_timeout == 12_000 + assert glide_config.advanced_config.connection_timeout == 13_000 + assert glide_config.advanced_config.tls_config is None + assert glide_config.database_id == 0 + + embeddings = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] + inserted, error = adapter.insert_embeddings(embeddings, [1, 2]) + assert (inserted, error) == (2, None) + assert len(runtime_client.exec_batches) == 1 + assert len(runtime_client.exec_batches[0]) == 2 + assert runtime_client.exec_batches[0][0][1][0] == "vdbbench_valkey:1" + assert runtime_client.exec_batches[0][0][1][-1] == np.asarray(embeddings[0], dtype=np.float32).tobytes() + + adapter.prepare_filter(non_filter) + assert adapter.search_embedding([1.0, 2.0, 3.0], k=5) == [7] + search_args = next(args for args in runtime_client.custom_commands if command_name(args[0]) == "FT.SEARCH") + assert search_args[:3] == [ + "FT.SEARCH", + "vdbbench_valkey", + "*=>[KNN 5 @vector $vec EF_RUNTIME 40]", + ] + assert "NOCONTENT" in search_args + assert search_args[-2:] == ["DIALECT", "2"] + finally: + context.__exit__(None, None, None) + + assert runtime_client.closed is True + + +def test_valkey_validates_insert_data(): + adapter, context, _, _, _ = make_adapter() + try: + with pytest.raises(ValueError, match="same length"): + adapter.insert_embeddings([[1.0, 2.0, 3.0]], []) + finally: + context.__exit__(None, None, None) + + adapter, context, _, _, _ = make_adapter() + adapter.with_scalar_labels = True + try: + with pytest.raises(ValueError, match="Scalar labels"): + adapter.insert_embeddings([[1.0, 2.0, 3.0]], [1]) + finally: + context.__exit__(None, None, None) + + +def test_valkey_filters_and_metrics(): + adapter, context, _, runtime_client, _ = make_adapter() + try: + adapter.prepare_filter(IntFilter(filter_rate=0.5, int_value=42)) + adapter.search_embedding([1.0, 2.0, 3.0]) + assert runtime_client.custom_commands[-1][2].startswith("@metadata:[42 +inf]=>[") + + adapter.prepare_filter(LabelFilter(label_percentage=0.1)) + adapter.search_embedding([1.0, 2.0, 3.0]) + assert runtime_client.custom_commands[-1][2].startswith("@label:{label_10p}=>[") + finally: + context.__exit__(None, None, None) + + config_kwargs = {"M": 16, "efConstruction": 200} + assert ValkeyHNSWConfig(metric_type=MetricType.COSINE, **config_kwargs).parse_metric() == "COSINE" + assert ValkeyHNSWConfig(metric_type=MetricType.L2, **config_kwargs).parse_metric() == "L2" + assert ValkeyHNSWConfig(metric_type=MetricType.IP, **config_kwargs).parse_metric() == "IP" + with pytest.raises(ValueError, match="Unsupported metric type"): + ValkeyHNSWConfig(metric_type=MetricType.HAMMING, **config_kwargs).parse_metric() + + +def test_valkey_cluster_client_and_drop_old(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr("vectordb_bench.backend.clients.valkey.valkey._DROP_CHUNK_SIZE", 1) + client = FakeClient() + client.created = True + client.scan_keys = [b"vdbbench_valkey:1", b"vdbbench_valkey:2"] + db_config = { + "host": "localhost", + "port": 6379, + "password": None, + "ssl": True, + "insecure_tls": True, + "request_timeout_ms": 20_000, + "connection_timeout_ms": 21_000, + "cmd": False, + } + + with ( + patch( + "vectordb_bench.backend.clients.valkey.valkey.GlideClusterClient.create", + return_value=client, + ) as create, + patch( + "vectordb_bench.backend.clients.valkey.valkey.ClusterScanCursor", + FakeClusterCursor, + ), + ): + Valkey( + dim=3, + db_config=db_config, + db_case_config=ValkeyHNSWConfig(M=16, efConstruction=200, ef=10), + drop_old=True, + ) + + assert client.closed is True + commands = [command_name(args[0]) for args in client.custom_commands] + assert commands == ["FT._LIST", "FT.DROPINDEX", "FT._LIST", "FT.CREATE"] + assert client.custom_commands[1] == ["FT.DROPINDEX", "vdbbench_valkey"] + assert len(client.exec_batches) == 2 + assert all(len(batch) == 1 for batch in client.exec_batches) + unlink_command = Batch(is_atomic=False).unlink(["key"]).commands[0][0] + assert all(batch[0][0] == unlink_command for batch in client.exec_batches) + + glide_config = create.call_args.args[0] + assert glide_config.use_tls is True + assert glide_config.request_timeout == 20_000 + assert glide_config.advanced_config.connection_timeout == 21_000 + assert glide_config.advanced_config.tls_config.use_insecure_tls is True + assert glide_config.database_id is None + + +def test_valkey_uses_config_defaults(): + client = FakeClient() + with patch( + "vectordb_bench.backend.clients.valkey.valkey.GlideClient.create", + return_value=client, + ) as create: + Valkey( + dim=3, + db_config=ValkeyConfig(host="localhost", cmd=True).to_dict(), + db_case_config=ValkeyHNSWConfig(M=16, efConstruction=200, ef=10), + ) + + glide_config = create.call_args.args[0] + assert glide_config.use_tls is True + assert glide_config.request_timeout == 600_000 + assert glide_config.advanced_config.connection_timeout == 10_000 + assert glide_config.advanced_config.tls_config is None + + +def test_valkey_checks_index_list(): + adapter = object.__new__(Valkey) + adapter.collection_name = "vdbbench_valkey" + client = FakeClient() + + assert adapter._index_exists(client) is False + client.created = True + assert adapter._index_exists(client) is True diff --git a/vectordb_bench/backend/clients/__init__.py b/vectordb_bench/backend/clients/__init__.py index beacb37af..546908079 100644 --- a/vectordb_bench/backend/clients/__init__.py +++ b/vectordb_bench/backend/clients/__init__.py @@ -35,6 +35,7 @@ class DB(Enum): PgDiskANN = "PgDiskANN" AlloyDB = "AlloyDB" Redis = "Redis" + Valkey = "Valkey" MemoryDB = "MemoryDB" Chroma = "Chroma" AWSOpenSearch = "OpenSearch" @@ -129,6 +130,11 @@ def init_cls(self) -> type[VectorDB]: # noqa: PLR0911, PLR0912, C901, PLR0915 return Redis + if self == DB.Valkey: + from .valkey.valkey import Valkey + + return Valkey + if self == DB.MemoryDB: from .memorydb.memorydb import MemoryDB @@ -347,6 +353,11 @@ def config_cls(self) -> type[DBConfig]: # noqa: PLR0911, PLR0912, C901, PLR0915 return RedisConfig + if self == DB.Valkey: + from .valkey.config import ValkeyConfig + + return ValkeyConfig + if self == DB.MemoryDB: from .memorydb.config import MemoryDBConfig @@ -719,6 +730,11 @@ def case_config_cls( # noqa: C901, PLR0911, PLR0912, PLR0915 return AdbpgIndexConfig + if self == DB.Valkey: + from .valkey.config import ValkeyHNSWConfig + + return ValkeyHNSWConfig + # DB.Pinecone, DB.Redis return EmptyDBCaseConfig diff --git a/vectordb_bench/backend/clients/valkey/cli.py b/vectordb_bench/backend/clients/valkey/cli.py new file mode 100644 index 000000000..c21df0fe6 --- /dev/null +++ b/vectordb_bench/backend/clients/valkey/cli.py @@ -0,0 +1,111 @@ +from typing import Annotated, TypedDict, Unpack + +import click +from pydantic import SecretStr + +from ....cli.cli import ( + CommonTypedDict, + HNSWFlavor2, + cli, + click_parameter_decorators_from_typed_dict, + run, +) +from .. import DB +from .config import ValkeyHNSWConfig + + +class ValkeyTypedDict(TypedDict): + host: Annotated[str, click.option("--host", type=str, help="Db host", required=True)] + password: Annotated[str, click.option("--password", type=str, help="Db password")] + port: Annotated[int, click.option("--port", type=int, default=6379, help="Db Port")] + collection_name: Annotated[ + str, + click.option( + "--collection-name", + type=str, + default="vdbbench_valkey", + show_default=True, + help="Valkey search index and key prefix", + ), + ] + ssl: Annotated[ + bool, + click.option( + "--ssl/--no-ssl", + is_flag=True, + show_default=True, + default=True, + help="Enable or disable SSL for Valkey", + ), + ] + insecure_tls: Annotated[ + bool, + click.option( + "--insecure-tls", + is_flag=True, + show_default=True, + default=False, + help="Disable TLS certificate verification", + ), + ] + request_timeout_ms: Annotated[ + int, + click.option( + "--request-timeout-ms", + type=int, + default=600_000, + show_default=True, + help="GLIDE request timeout in milliseconds", + ), + ] + connection_timeout_ms: Annotated[ + int, + click.option( + "--connection-timeout-ms", + type=int, + default=10_000, + show_default=True, + help="GLIDE connection timeout in milliseconds", + ), + ] + cmd: Annotated[ + bool, + click.option( + "--cmd", + is_flag=True, + show_default=True, + default=False, + help="Cluster Mode Disabled (CMD) for Valkey doesn't use Cluster conn", + ), + ] + + +class ValkeyHNSWTypedDict(CommonTypedDict, ValkeyTypedDict, HNSWFlavor2): ... + + +@cli.command() +@click_parameter_decorators_from_typed_dict(ValkeyHNSWTypedDict) +def Valkey(**parameters: Unpack[ValkeyHNSWTypedDict]): + from .config import ValkeyConfig + + run( + db=DB.Valkey, + db_config=ValkeyConfig( + db_label=parameters["db_label"], + password=SecretStr(parameters["password"]) if parameters["password"] else None, + host=SecretStr(parameters["host"]), + port=parameters["port"], + collection_name=parameters["collection_name"], + ssl=parameters["ssl"], + insecure_tls=parameters["insecure_tls"], + request_timeout_ms=parameters["request_timeout_ms"], + connection_timeout_ms=parameters["connection_timeout_ms"], + cmd=parameters["cmd"], + ), + db_case_config=ValkeyHNSWConfig( + M=parameters["m"] if parameters["m"] is not None else 16, + efConstruction=parameters["ef_construction"] if parameters["ef_construction"] is not None else 200, + ef=parameters["ef_runtime"] if parameters["ef_runtime"] is not None else 10, + ), + **parameters, + ) diff --git a/vectordb_bench/backend/clients/valkey/config.py b/vectordb_bench/backend/clients/valkey/config.py new file mode 100644 index 000000000..2c73ac4a1 --- /dev/null +++ b/vectordb_bench/backend/clients/valkey/config.py @@ -0,0 +1,67 @@ +from pydantic import BaseModel, Field, SecretStr, model_validator + +from ..api import DBCaseConfig, DBConfig, IndexType, MetricType + + +class ValkeyConfig(DBConfig): + password: SecretStr | None = None + host: SecretStr + port: int = Field(default=6379, ge=1, le=65535) + ssl: bool = True + insecure_tls: bool = False + cmd: bool = False + request_timeout_ms: int = Field(default=600_000, gt=0) + connection_timeout_ms: int = Field(default=10_000, gt=0) + collection_name: str = Field(default="vdbbench_valkey", pattern=r"^[A-Za-z0-9_.-]+$") + + @model_validator(mode="after") + def validate_tls(self) -> "ValkeyConfig": + if self.insecure_tls and not self.ssl: + msg = "insecure_tls requires ssl=True" + raise ValueError(msg) + return self + + def to_dict(self) -> dict: + return { + "host": self.host.get_secret_value(), + "port": self.port, + "password": self.password.get_secret_value() if self.password is not None else None, + "ssl": self.ssl, + "insecure_tls": self.insecure_tls, + "cmd": self.cmd, + "request_timeout_ms": self.request_timeout_ms, + "connection_timeout_ms": self.connection_timeout_ms, + "collection_name": self.collection_name, + } + + +class ValkeyIndexConfig(BaseModel): + metric_type: MetricType | None = None + + def parse_metric(self) -> str: + if self.metric_type in {None, MetricType.COSINE}: + return "COSINE" + if self.metric_type in {MetricType.L2, MetricType.IP}: + return self.metric_type.value + msg = f"Unsupported metric type: {self.metric_type}" + raise ValueError(msg) + + +class ValkeyHNSWConfig(ValkeyIndexConfig, DBCaseConfig): + M: int = Field(default=16, gt=0) + efConstruction: int = Field(default=200, gt=0) + ef: int = Field(default=10, gt=0) + index: IndexType = IndexType.HNSW + + def index_param(self) -> dict: + return { + "metric_type": self.parse_metric(), + "index_type": self.index.value, + "params": {"M": self.M, "efConstruction": self.efConstruction}, + } + + def search_param(self) -> dict: + return { + "metric_type": self.parse_metric(), + "params": {"ef": self.ef}, + } diff --git a/vectordb_bench/backend/clients/valkey/valkey.py b/vectordb_bench/backend/clients/valkey/valkey.py new file mode 100644 index 000000000..2e6246a08 --- /dev/null +++ b/vectordb_bench/backend/clients/valkey/valkey.py @@ -0,0 +1,272 @@ +import logging +from collections.abc import Generator +from contextlib import contextmanager +from typing import Any + +import numpy as np +from glide_sync import ( + AdvancedGlideClientConfiguration, + AdvancedGlideClusterClientConfiguration, + Batch, + ClusterBatch, + ClusterScanCursor, + DataType, + DistanceMetricType, + FtCreateOptions, + FtSearchLimit, + FtSearchOptions, + GlideClient, + GlideClientConfiguration, + GlideClusterClient, + GlideClusterClientConfiguration, + NodeAddress, + NumericField, + ServerCredentials, + TagField, + TlsAdvancedConfiguration, + VectorAlgorithm, + VectorField, + VectorFieldAttributesHnsw, + VectorType, + ft, +) + +from vectordb_bench.backend.filter import Filter, FilterOp + +from ..api import VectorDB +from .config import ValkeyHNSWConfig + +log = logging.getLogger(__name__) +ValkeyConnection = GlideClient | GlideClusterClient +ValkeyBatch = Batch | ClusterBatch +_DROP_CHUNK_SIZE = 10_000 + + +class Valkey(VectorDB): + supported_filter_types: list[FilterOp] = [ + FilterOp.NonFilter, + FilterOp.NumGE, + FilterOp.StrEqual, + ] + + def __init__( + self, + dim: int, + db_config: dict, + db_case_config: ValkeyHNSWConfig, + collection_name: str = "vdbbench_valkey", + drop_old: bool = False, + with_scalar_labels: bool = False, + **kwargs, + ): + self.name = "Valkey" + self.db_config = db_config + self.case_config = db_case_config + self.collection_name = collection_name + self.key_prefix = f"{self.collection_name}:" + self.with_scalar_labels = with_scalar_labels + self.filter_query = "*" + self.cluster_mode = not self.db_config["cmd"] + self.conn: ValkeyConnection | None = None + + conn = self._create_client() + try: + if drop_old: + self._drop_index(conn) + self._make_index(dim, conn) + finally: + conn.close() + + def _create_client(self) -> ValkeyConnection: + password = self.db_config["password"] + config_kwargs = { + "addresses": [NodeAddress(self.db_config["host"], self.db_config["port"])], + "use_tls": self.db_config["ssl"], + "credentials": ServerCredentials(password=password) if password else None, + "request_timeout": self.db_config["request_timeout_ms"], + } + connection_timeout = self.db_config["connection_timeout_ms"] + tls_config = TlsAdvancedConfiguration(use_insecure_tls=True) if self.db_config["insecure_tls"] else None + if self.cluster_mode: + config_kwargs["advanced_config"] = AdvancedGlideClusterClientConfiguration( + connection_timeout=connection_timeout, + tls_config=tls_config, + ) + return GlideClusterClient.create(GlideClusterClientConfiguration(**config_kwargs)) + config_kwargs["advanced_config"] = AdvancedGlideClientConfiguration( + connection_timeout=connection_timeout, + tls_config=tls_config, + ) + return GlideClient.create(GlideClientConfiguration(database_id=0, **config_kwargs)) + + def _drop_index(self, conn: ValkeyConnection) -> None: + if not self._index_exists(conn): + log.info("Valkey index did not exist: %s", self.collection_name) + else: + ft.dropindex(conn, self.collection_name) + log.info("Valkey dropped old index: %s", self.collection_name) + self._delete_prefixed_documents(conn) + + def _index_exists(self, conn: ValkeyConnection) -> bool: + encoded_name = self.collection_name.encode() + return any(index_name in (self.collection_name, encoded_name) for index_name in ft.list(conn)) + + def _new_batch(self) -> ValkeyBatch: + return ClusterBatch(is_atomic=False) if self.cluster_mode else Batch(is_atomic=False) + + def _scan_keys(self, conn: ValkeyConnection) -> Generator[bytes, None, None]: + match = f"{self.key_prefix}*" + if self.cluster_mode: + cursor = ClusterScanCursor() + while not cursor.is_finished(): + cursor, keys = conn.scan(cursor, match=match, count=_DROP_CHUNK_SIZE) + yield from keys + return + + cursor: str | bytes | int = "0" + while True: + cursor, keys = conn.scan(cursor, match=match, count=_DROP_CHUNK_SIZE) + yield from keys + if cursor in {"0", b"0", 0}: + return + + def _delete_prefixed_documents(self, conn: ValkeyConnection) -> None: + deleted = 0 + pending = 0 + batch = self._new_batch() + for key in self._scan_keys(conn): + batch.unlink([key]) + deleted += 1 + pending += 1 + if pending == _DROP_CHUNK_SIZE: + conn.exec(batch, raise_on_error=True) + batch = self._new_batch() + pending = 0 + if pending: + conn.exec(batch, raise_on_error=True) + if deleted: + log.info("Valkey deleted %s documents with prefix %s", deleted, self.key_prefix) + + def _make_index(self, vector_dimensions: int, conn: ValkeyConnection) -> None: + if self._index_exists(conn): + return + + index_param = self.case_config.index_param() + schema = [ + NumericField("metadata"), + VectorField( + "vector", + VectorAlgorithm.HNSW, + VectorFieldAttributesHnsw( + dimensions=vector_dimensions, + distance_metric=DistanceMetricType(index_param["metric_type"]), + type=VectorType.FLOAT32, + number_of_edges=index_param["params"]["M"], + vectors_examined_on_construction=index_param["params"]["efConstruction"], + ), + ), + ] + if self.with_scalar_labels: + schema.append(TagField("label")) + + ft.create( + conn, + self.collection_name, + schema, + FtCreateOptions(data_type=DataType.HASH, prefixes=[self.key_prefix]), + ) + + @contextmanager + def init(self) -> Generator[None, None, None]: + conn = self._create_client() + self.conn = conn + try: + yield + finally: + self.conn = None + conn.close() + + def optimize(self, data_size: int | None = None) -> None: + """Valkey does not require a post-load optimization step.""" + return + + def insert_embeddings( + self, + embeddings: list[list[float]], + metadata: list[int], + labels_data: list[str] | None = None, + **kwargs: Any, + ) -> tuple[int, Exception | None]: + assert self.conn is not None, "Please call self.init() before insert_embeddings" + if len(embeddings) != len(metadata): + msg = "Embeddings and metadata must have the same length" + raise ValueError(msg) + if self.with_scalar_labels and (labels_data is None or len(labels_data) != len(embeddings)): + msg = "Scalar labels must be provided for every embedding" + raise ValueError(msg) + if not embeddings: + return 0, None + + try: + batch = self._new_batch() + for i, embedding in enumerate(embeddings): + doc_id = metadata[i] + mapping = { + "metadata": str(doc_id), + "vector": np.asarray(embedding, dtype=np.float32).tobytes(), + } + if self.with_scalar_labels and labels_data is not None: + mapping["label"] = labels_data[i] + batch.hset(f"{self.key_prefix}{doc_id}", mapping) + self.conn.exec(batch, raise_on_error=True) + except Exception as e: + return 0, e + + return len(embeddings), None + + def prepare_filter(self, filters: Filter) -> None: + if filters.type == FilterOp.NonFilter: + self.filter_query = "*" + elif filters.type == FilterOp.NumGE: + self.filter_query = f"@metadata:[{filters.int_value} +inf]" + elif filters.type == FilterOp.StrEqual: + label = filters.label_value.replace("\\", "\\\\").replace(" ", "\\ ") + self.filter_query = f"@label:{{{label}}}" + else: + msg = f"Unsupported filter for Valkey: {filters}" + raise ValueError(msg) + + def search_embedding( + self, + query: list[float], + k: int = 100, + filters: dict | None = None, + timeout: int | None = None, + **kwargs: Any, + ) -> list[int]: + assert self.conn is not None, "Please call self.init() before search_embedding" + + ef_runtime = self.case_config.search_param()["params"]["ef"] + knn = f"KNN {k} @vector $vec EF_RUNTIME {ef_runtime}" + query_vector = np.asarray(query, dtype=np.float32).tobytes() + result = ft.search( + self.conn, + self.collection_name, + f"{self.filter_query}=>[{knn}]", + FtSearchOptions( + nocontent=True, + timeout=timeout, + params={"vec": query_vector}, + limit=FtSearchLimit(0, k), + dialect=2, + ), + ) + documents = result[1] if len(result) > 1 else {} + return [self._parse_document_id(document_id) for document_id in documents] + + def _parse_document_id(self, document_id: str | bytes) -> int: + if isinstance(document_id, bytes): + document_id = document_id.decode() + if document_id.startswith(self.key_prefix): + document_id = document_id[len(self.key_prefix) :] + return int(document_id) diff --git a/vectordb_bench/cli/vectordbbench.py b/vectordb_bench/cli/vectordbbench.py index e0cb98652..2ec5d7d7b 100644 --- a/vectordb_bench/cli/vectordbbench.py +++ b/vectordb_bench/cli/vectordbbench.py @@ -47,6 +47,7 @@ from ..backend.clients.test.cli import Test from ..backend.clients.tidb.cli import TiDB from ..backend.clients.turbopuffer.cli import TurboPuffer, TurboPufferUnpin +from ..backend.clients.valkey.cli import Valkey from ..backend.clients.vectorchord.cli import VectorChordGraph, VectorChordRQ from ..backend.clients.vespa.cli import Vespa from ..backend.clients.volc_mysql.cli import VolcMySQLHNSW @@ -61,6 +62,7 @@ cli.add_command(PgVectoRSHNSW) cli.add_command(PgVectoRSIVFFlat) cli.add_command(Redis) +cli.add_command(Valkey) cli.add_command(MemoryDB) cli.add_command(Weaviate) cli.add_command(Test)