Main

.env.example

@@ -1,6 +1,5 @@
-# LOG_LEVEL=
-# LOG_PATH=
-# LOG_NAME=
+LOG_LEVEL=INFO
+LOG_FILE="logs/vectordb_bench.log"
 # TIMEZONE=
 
 # NUM_PER_BATCH=

.github/workflows/pull_request.yml

@@ -4,6 +4,7 @@ on:
   pull_request:
     branches:
       - main
+      - vdbbench_*
 
 jobs:
   build:

install/requirements_py3.11.txt

@@ -3,7 +3,7 @@ grpcio-tools==1.53.0
 qdrant-client
 pinecone-client
 weaviate-client
-elasticsearch
+elasticsearch==8.16.0
 pgvector
 pgvecto_rs[psycopg3]>=0.2.1
 sqlalchemy
@@ -24,3 +24,6 @@ scikit-learn
 pymilvus
 clickhouse_connect
 pyvespa
+mysql-connector-python
+packaging
+hdrhistogram>=0.10.1

pyproject.toml

@@ -4,6 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [tool.setuptools.package-data]
 "vectordb_bench.results" = ["*.json"]
+"vectordb_bench.fig.homepage" = ["*.png"]
 
 [tool.setuptools.packages.find]
 where = ["."]
@@ -12,7 +13,7 @@ include = ["vectordb_bench", "vectordb_bench.cli"]
 [project]
 name = "vectordb-bench"
 authors = [
-  {name="XuanYang-cn", email="xuan.yang@zilliz.com"},
+  {name="VDBBench teams", email="pymilvus@zilliz.com"},
 ]
 description = "VectorDBBench is not just an offering of benchmark results for mainstream vector databases and cloud services, it's your go-to tool for the ultimate performance and cost-effectiveness comparison. Designed with ease-of-use in mind, VectorDBBench is devised to help users, even non-professionals, reproduce results or test new systems, making the hunt for the optimal choice amongst a plethora of cloud services and open-source vector databases a breeze."
 
@@ -35,10 +36,12 @@ dependencies = [
     "psutil",
     "polars",
     "plotly",
-    "environs<14.1.0",
+    "environs",
     "pydantic<v2",
     "scikit-learn",
     "pymilvus", # with pandas, numpy, ujson
+    "ujson",
+    "hdrhistogram>=0.10.1",
 ]
 dynamic = ["version"]
 
@@ -48,6 +51,7 @@ test = [
     "ruff",
     "pytest",
 ]
+restful = [ "flask" ]
 
 all = [
     "grpcio==1.53.0", # for qdrant-client and pymilvus
@@ -67,11 +71,13 @@ all = [
     "opensearch-py",
     "memorydb",
     "alibabacloud_ha3engine_vector",
-    "alibabacloud_searchengine20211025",
     "mariadb",
     "PyMySQL",
     "clickhouse-connect",
     "pyvespa",
+    "lancedb",
+    "mysql-connector-python",
+    "turbopuffer[fast]",
 ]
 
 qdrant          = [ "qdrant-client" ]
@@ -88,18 +94,26 @@ redis           = [ "redis" ]
 memorydb        = [ "memorydb" ]
 chromadb        = [ "chromadb" ]
 opensearch      = [ "opensearch-py" ]
-aliyun_opensearch = [ "alibabacloud_ha3engine_vector", "alibabacloud_searchengine20211025"]
+aliyun_opensearch = [ "alibabacloud_ha3engine_vector" ]
 mongodb         = [ "pymongo" ]
 mariadb         = [ "mariadb" ]
 tidb            = [ "PyMySQL" ]
+cockroachdb     = [ "psycopg[binary,pool]", "pgvector" ]
 clickhouse      = [ "clickhouse-connect" ]
 vespa           = [ "pyvespa" ]
+lancedb         = [ "lancedb" ]
+oceanbase       = [ "mysql-connector-python" ]
+alisql          = [ "mysql-connector-python" ]
+doris           = [ "doris-vector-search" ]
+turbopuffer     = [ "turbopuffer" ]
 
 [project.urls]
-"repository" = "https://github.com/zilliztech/VectorDBBench"
+Repository  = "https://github.com/zilliztech/VectorDBBench"
+Leaderborad = "https://zilliz.com/vdbbench-leaderboard"
 
 [project.scripts]
 init_bench = "vectordb_bench.__main__:main"
+init_bench_rest = "vectordb_bench.restful.app:main"
 vectordbbench = "vectordb_bench.cli.vectordbbench:cli"
 
 [tool.setuptools_scm]
@@ -216,3 +230,9 @@ builtins-ignorelist = [
     # "dict", # TODO
     # "filter",
 ]
+
+[tool.ruff.lint.per-file-ignores]
+"vectordb_bench/backend/clients/*" = ["PLC0415"]
+"vectordb_bench/cli/batch_cli.py" = ["PLC0415"]
+"vectordb_bench/backend/data_source.py" = ["PLC0415"]
+

tests/test_cockroachdb.py

@@ -0,0 +1,128 @@
+"""
+Tests for CockroachDB vector database client.
+
+Assumes CockroachDB is running on localhost:26257.
+
+To start CockroachDB locally:
+    cockroach start-single-node --insecure --listen-addr=localhost:26257
+"""
+
+import logging
+
+import numpy as np
+
+from vectordb_bench.models import DB
+
+log = logging.getLogger(__name__)
+
+
+class TestCockroachDB:
+    """Test suite for CockroachDB vector operations."""
+
+    def test_insert_and_search(self):
+        """Test basic insert and search operations."""
+        assert DB.CockroachDB.value == "CockroachDB"
+
+        dbcls = DB.CockroachDB.init_cls
+        dbConfig = DB.CockroachDB.config_cls
+
+        # Connection config (matches your local CockroachDB instance)
+        config = {
+            "host": "localhost",
+            "port": 26257,
+            "user_name": "root",
+            "password": "",
+            "db_name": "defaultdb",
+            "table_name": "test_cockroachdb",
+        }
+
+        # Note: sslmode=disable is handled in the client's connect_config options
+
+        dim = 128
+        count = 1000
+
+        # Initialize CockroachDB client
+        cockroachdb = dbcls(
+            dim=dim,
+            db_config=config,
+            db_case_config=None,
+            collection_name="test_cockroachdb",
+            drop_old=True,
+        )
+
+        embeddings = [[np.random.random() for _ in range(dim)] for _ in range(count)]
+
+        # Test insert
+        with cockroachdb.init():
+            res = cockroachdb.insert_embeddings(embeddings=embeddings, metadata=list(range(count)))
+
+            assert res[0] == count, f"Insert count mismatch: {res[0]} != {count}"
+            assert res[1] is None, f"Insert failed with error: {res[1]}"
+
+        # Test search
+        with cockroachdb.init():
+            test_id = np.random.randint(count)
+            q = embeddings[test_id]
+
+            res = cockroachdb.search_embedding(query=q, k=10)
+
+            assert len(res) > 0, "Search returned no results"
+            assert res[0] == int(test_id), f"Top result {res[0]} != query id {test_id}"
+
+        log.info("CockroachDB insert and search test passed")
+
+    def test_search_with_filter(self):
+        """Test search with filters."""
+        assert DB.CockroachDB.value == "CockroachDB"
+
+        dbcls = DB.CockroachDB.init_cls
+
+        config = {
+            "host": "localhost",
+            "port": 26257,
+            "user_name": "root",
+            "password": "",
+            "db_name": "defaultdb",
+            "table_name": "test_cockroachdb_filter",
+        }
+
+        dim = 128
+        count = 1000
+        filter_value = 0.9
+
+        cockroachdb = dbcls(
+            dim=dim,
+            db_config=config,
+            db_case_config=None,
+            collection_name="test_cockroachdb_filter",
+            drop_old=True,
+        )
+
+        embeddings = [[np.random.random() for _ in range(dim)] for _ in range(count)]
+
+        # Insert data
+        with cockroachdb.init():
+            res = cockroachdb.insert_embeddings(embeddings=embeddings, metadata=list(range(count)))
+            assert res[0] == count, f"Insert count mismatch"
+
+        # Search with filter
+        with cockroachdb.init():
+            filter_id = int(count * filter_value)
+            test_id = np.random.randint(filter_id, count)
+            q = embeddings[test_id]
+
+            from vectordb_bench.backend.filter import IntFilter
+
+            filters = IntFilter(int_value=filter_id, filter_rate=0.9)
+            cockroachdb.prepare_filter(filters)
+
+            res = cockroachdb.search_embedding(query=q, k=10)
+
+            assert len(res) > 0, "Filtered search returned no results"
+            assert res[0] == int(test_id), f"Top result {res[0]} != query id {test_id}"
+
+            # Verify all results are >= filter_value
+            for result_id in res:
+                assert int(result_id) >= filter_id, f"Result {result_id} < filter threshold {filter_id}"
+
+        log.info("CockroachDB filter test passed")

vectordb_bench/__init__.py

@@ -6,52 +6,35 @@
 from . import log_util
 
 env = environs.Env()
-env.read_env(".env", False)
+env.read_env(path=".env", recurse=False)
 
 
 class config:
     ALIYUN_OSS_URL = "assets.zilliz.com.cn/benchmark/"
     AWS_S3_URL = "assets.zilliz.com/benchmark/"
 
     LOG_LEVEL = env.str("LOG_LEVEL", "INFO")
+    LOG_FILE = env.str("LOG_FILE", "logs/vectordb_bench.log")
 
     DEFAULT_DATASET_URL = env.str("DEFAULT_DATASET_URL", AWS_S3_URL)
+    DATASET_SOURCE = env.str("DATASET_SOURCE", "S3")  # Options "S3" or "AliyunOSS"
     DATASET_LOCAL_DIR = env.path("DATASET_LOCAL_DIR", "/tmp/vectordb_bench/dataset")
     NUM_PER_BATCH = env.int("NUM_PER_BATCH", 100)
+    TIME_PER_BATCH = 1  # 1s. for streaming insertion.
+    MAX_INSERT_RETRY = 5
+    MAX_SEARCH_RETRY = 5
+
+    LOAD_MAX_TRY_COUNT = 10
 
     DROP_OLD = env.bool("DROP_OLD", True)
     USE_SHUFFLED_DATA = env.bool("USE_SHUFFLED_DATA", True)
 
-    NUM_CONCURRENCY = env.list(
-        "NUM_CONCURRENCY",
-        [
-            1,
-            5,
-            10,
-            15,
-            20,
-            25,
-            30,
-            35,
-            40,
-            45,
-            50,
-            55,
-            60,
-            65,
-            70,
-            75,
-            80,
-            85,
-            90,
-            95,
-            100,
-        ],
-        subcast=int,
-    )
+    NUM_CONCURRENCY = env.list("NUM_CONCURRENCY", [1, 5, 10, 20, 30, 40, 60, 80], subcast=int)
 
     CONCURRENCY_DURATION = 30
 
+    CONCURRENCY_TIMEOUT = 3600
+
     RESULTS_LOCAL_DIR = env.path(
         "RESULTS_LOCAL_DIR",
         pathlib.Path(__file__).parent.joinpath("results"),
@@ -66,21 +49,29 @@ class config:
 
     CAPACITY_TIMEOUT_IN_SECONDS = 24 * 3600  # 24h
     LOAD_TIMEOUT_DEFAULT = 24 * 3600  # 24h
+    LOAD_TIMEOUT_768D_100K = 24 * 3600  # 24h
     LOAD_TIMEOUT_768D_1M = 24 * 3600  # 24h
     LOAD_TIMEOUT_768D_10M = 240 * 3600  # 10d
     LOAD_TIMEOUT_768D_100M = 2400 * 3600  # 100d
 
     LOAD_TIMEOUT_1536D_500K = 24 * 3600  # 24h
     LOAD_TIMEOUT_1536D_5M = 240 * 3600  # 10d
 
+    LOAD_TIMEOUT_1024D_1M = 24 * 3600  # 24h
+    LOAD_TIMEOUT_1024D_10M = 240 * 3600  # 10d
+
     OPTIMIZE_TIMEOUT_DEFAULT = 24 * 3600  # 24h
+    OPTIMIZE_TIMEOUT_768D_100K = 24 * 3600  # 24h
     OPTIMIZE_TIMEOUT_768D_1M = 24 * 3600  # 24h
     OPTIMIZE_TIMEOUT_768D_10M = 240 * 3600  # 10d
     OPTIMIZE_TIMEOUT_768D_100M = 2400 * 3600  # 100d
 
     OPTIMIZE_TIMEOUT_1536D_500K = 24 * 3600  # 24h
     OPTIMIZE_TIMEOUT_1536D_5M = 240 * 3600  # 10d
 
+    OPTIMIZE_TIMEOUT_1024D_1M = 24 * 3600  # 24h
+    OPTIMIZE_TIMEOUT_1024D_10M = 240 * 3600  # 10d
+
     def display(self) -> str:
         return [
             i
@@ -89,4 +80,4 @@ def display(self) -> str:
         ]
 
 
-log_util.init(config.LOG_LEVEL)
+log_util.init(config.LOG_LEVEL, pathlib.Path(config.LOG_FILE))

vectordb_bench/__main__.py

@@ -17,7 +17,7 @@ def run_streamlit():
     cmd = [
         "streamlit",
         "run",
-        f"{pathlib.Path(__file__).parent}/frontend/vdb_benchmark.py",
+        f"{pathlib.Path(__file__).parent}/frontend/vdbbench.py",
         "--logger.level",
         "info",
         "--theme.base",